golang 实现并发的websocket

Posted 2020-12-11 Poltu

公司要求使用golang做识别系统的web后端，我采用gin框架开发，其中部分功能使用了websocket实现前后端的实时消息推送刷新。

记录一下golang中使用channel和锁将websocket封装成为可并发读写的websocket：

package cws

import (
	"errors"
	"github.com/gorilla/websocket"
	"sync"
)

//封装websocket并发读写操作

type Connection struct {
	WsConn    *websocket.Conn
	InChan    chan []byte
	OutChan   chan models.BayDataS
	CloseChan chan byte
	Mutex     sync.Mutex
	IsClosed  bool
}

func InitConnection(wsConn *websocket.Conn) (conn *Connection, err error) {
	conn = &Connection{
		WsConn:    wsConn,
		InChan:    make(chan []byte, 1000),
		OutChan:   make(chan []byte, 1000),
		CloseChan: make(chan byte, 1),
	}
	//读协程
	go conn.ReadLoop()
	//写协程
	go conn.WriteLoop()
	return
}

func (conn *Connection) ReadMess() (data []byte, err error) {
	select {
	case data = <-conn.InChan:
	case <-conn.CloseChan:
		err = errors.New("connection is closed")
	}
	return
}

func (conn *Connection) WriteMes(data []byte) (err error) {
	select {
	case conn.OutChan <- data:
	case <-conn.CloseChan:
		err = errors.New("connection is closed")
	}
	return
}

func (conn *Connection) Close() {
	conn.Close() //本身线程安全，可重入
	//加锁，只能执行一次
	conn.Mutex.Lock()
	if !conn.IsClosed {
		close(conn.CloseChan)
		conn.IsClosed = true
	}
}

//具体实现读消息
func (conn *Connection) ReadLoop() {
	var (
		data []byte
		err  error
	)
	for {
		if _, data, err = conn.WsConn.ReadMessage(); err != nil {
			goto ERR
		}
		select {
		case conn.InChan <- data:
		case <-conn.CloseChan:
			goto ERR
		}
	}
ERR:
	conn.Close()
}

//具体实现写消息
func (conn *Connection) WriteLoop() {
	var (
		data models.BayDataS
		err  error
	)
	for {
		select {
		case data = <-conn.OutChan:
		case <-conn.CloseChan:
			goto ERR
		}
		if err = conn.WsConn.WriteMessage(websocket.TextMessage, data); err != nil {
			goto ERR
		}
	}
ERR:
	conn.Close()
}

当然websocket支持发送和接收的消息多种，像我的话常用的是Json 如：

conn.WsConn.WriteJSON(data) //data可以是struct,map等