PandaX/kit/ws/ws.go

package ws

import (
	"encoding/json"
	"net/http"
	"sync"
	"time"

	"pandax/kit/logger"

	"github.com/gorilla/websocket"
)

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

type Connection struct {
	conn *websocket.Conn
	lock sync.Mutex
}

var connections = make(map[uint64]*Connection)
var connLock sync.Mutex

func init() {
	go checkConn()
}

// 放置ws连接
func Put(userId uint64, conn *websocket.Conn) {
	connLock.Lock()
	defer connLock.Unlock()

	Delete(userId)
	connections[userId] = &Connection{
		conn: conn,
	}
}

func checkConn() {
	heartbeat := time.Duration(20) * time.Second
	tick := time.NewTicker(heartbeat)
	for range tick.C {
		connLock.Lock()
		for uid, conn := range connections {
			conn.lock.Lock()
			err := conn.conn.WriteControl(websocket.PingMessage, []byte("ping"), time.Now().Add(heartbeat/2))
			conn.lock.Unlock()
			if err != nil {
				logger.Log.Info("删除ping失败的websocket连接：uid = ", uid)
				Delete(uid)
			}
		}
		connLock.Unlock()
	}
}

// 删除ws连接
func Delete(userid uint64) {
	connLock.Lock()
	defer connLock.Unlock()

	conn := connections[userid]
	if conn != nil {
		conn.lock.Lock()
		conn.conn.Close()
		conn.lock.Unlock()
		delete(connections, userid)
	}
}

// 对指定用户发送消息
func SendMsg(userId uint64, msg *Msg) {
	connLock.Lock()
	defer connLock.Unlock()

	conn := connections[userId]
	if conn != nil {
		conn.lock.Lock()
		defer conn.lock.Unlock()

		bytes, err := json.Marshal(msg)
		if err != nil {
			logger.Log.Error("发送消息失败：", err)
			return
		}
		err = conn.conn.WriteMessage(websocket.TextMessage, bytes)
		if err != nil {
			logger.Log.Error("发送消息失败：", err)
		}
	}
}