golang 實現並發的websocket

公司要求使用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等