Golang Gin 引入Gorilla WebSocket 实现指定连接通信
Gorilla引入:
github.com/gorilla/websocket
代码完全时个人逻辑,存在一些问题,目前没有考虑线程锁,测试效果仅供娱乐。
package service
import (
"fmt"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
"net/http"
"sync"
)
//待优化:加入 缓存、线程锁、MQ
type message struct {
msg []byte
massageType int
msgSize int
}
type session struct {
Request *http.Request
conn *websocket.Conn
outputDone chan struct{}
open bool
rwmutex *sync.RWMutex
}
// SessionMap 保存WebSocket链接
var SessionMap = make(map[string]session)
var keyMap = make(map[*websocket.Conn]string)
func WsGorilla(ctx *gin.Context) {
//http 协议升级
upgraded := websocket.Upgrader{
CheckOrigin: func(r *http.Request) bool {
return true
},
ReadBufferSize: 1024,
WriteBufferSize: 1024,
}
Conn, err := upgraded.Upgrade(ctx.Writer, ctx.Request, nil)
addSession(Conn, ctx.Request, "小明")
if err != nil {
return
}
var msg []byte
var messageType int
//读取并发送消息
for {
messageType, msg, err = Conn.ReadMessage()
if err != nil {
_ = Conn.Close()
//清除会话信息
clearSession(Conn)
return
} else if msg != nil {
err := SendMessage(messageType, msg, "小明")
if err != nil {
return
}
}
}
}
// 添加会话
func addSession(coon *websocket.Conn, request *http.Request, key string) {
fmt.Println("addSession")
SessionMap[key] = session{
conn: coon,
Request: request,
}
keyMap[coon] = key
}
// 清除指定会话
func clearSession(coon *websocket.Conn) {
delete(SessionMap, keyMap[coon])
delete(keyMap, coon)
}
func SendMessage(messageType int, msg []byte, key string) error {
conn := SessionMap[key].conn
fmt.Println("SendMessage")
err := conn.WriteMessage(messageType, msg)
if msg = nil; err != nil {
return err
}
return nil
}