goalng中websocket的简单应用

goalng中websocket的简单应用

直接上干货

var gClients map[string]*websocket.Conn
var totalMtx sync.Mutex

var upGrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	//HandshakeTimeout: 10 * time.Second,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

func JoinTask(c *gin.Context) {
	//升级get请求为webSocket协议
	ws, err := upGrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}
	//读取ws中的数据
	_, message, err := ws.ReadMessage()
	if err != nil {
	}
	if gClients == nil {
		gClients = make(map[string]*websocket.Conn)
	}

	if client, ok := gClients[string(message)]; ok {
		client.Close()
	}
	gClients[string(message)] = ws
}

func LeaveTask(c *gin.Context) {
	ws, err := upGrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}
	defer ws.Close()

	_, message, err := ws.ReadMessage()
	if err != nil {
		return
	}
	delClientConn(string(message))
}

func delClientConn(taskName string) error {
	if client, ok := gClients[taskName]; ok {
		client.Close()
		delete(gClients, taskName)
	}
	return nil
}

func Push(taskName string, message []byte) error {
	if client, ok := gClients[taskName]; !ok {
		return errors.New("No client")
	} else {
		totalMtx.Lock()
		defer totalMtx.Unlock()

		err := client.WriteMessage(websocket.TextMessage, message)
		if err != nil {
			return err
		}
	}
	return nil
}

注意

在使用push进行给指定的taskName推送消息时，建议先去判断全局的gClients中是否有该taskName，有的话再推送，没有的话推送消息会报错。

[2020-07-09 02:08:25.722][+08][Error][  105] Websocket push error:

如有争议，欢迎留言讨论。