mirror of
https://gitee.com/dromara/mayfly-go
synced 2025-11-04 00:10:25 +08:00
75 lines
1.4 KiB
Go
75 lines
1.4 KiB
Go
package ws
|
||
|
||
import (
|
||
"encoding/json"
|
||
"mayfly-go/pkg/global"
|
||
"net/http"
|
||
"time"
|
||
|
||
"github.com/gorilla/websocket"
|
||
)
|
||
|
||
var Upgrader = websocket.Upgrader{
|
||
ReadBufferSize: 1024,
|
||
WriteBufferSize: 1024 * 1024 * 10,
|
||
CheckOrigin: func(r *http.Request) bool {
|
||
return true
|
||
},
|
||
}
|
||
|
||
var conns = make(map[uint64]*websocket.Conn, 100)
|
||
|
||
func init() {
|
||
checkConn()
|
||
}
|
||
|
||
// 放置ws连接
|
||
func Put(userId uint64, conn *websocket.Conn) {
|
||
existConn := conns[userId]
|
||
if existConn != nil {
|
||
Delete(userId)
|
||
}
|
||
|
||
conn.SetCloseHandler(func(code int, text string) error {
|
||
Delete(userId)
|
||
return nil
|
||
})
|
||
conns[userId] = conn
|
||
}
|
||
|
||
func checkConn() {
|
||
heartbeat := time.Duration(60) * time.Second
|
||
tick := time.NewTicker(heartbeat)
|
||
go func() {
|
||
for range tick.C {
|
||
// 遍历所有连接，ping失败的则删除掉
|
||
for uid, conn := range conns {
|
||
err := conn.WriteControl(websocket.PingMessage, []byte("ping"), time.Now().Add(heartbeat/2))
|
||
if err != nil {
|
||
Delete(uid)
|
||
return
|
||
}
|
||
}
|
||
}
|
||
}()
|
||
}
|
||
|
||
// 删除ws连接
|
||
func Delete(userid uint64) {
|
||
global.Log.Debug("移除websocket连接：uid = ", userid)
|
||
conn := conns[userid]
|
||
if conn != nil {
|
||
conn.Close()
|
||
delete(conns, userid)
|
||
}
|
||
}
|
||
|
||
// 对指定用户发送消息
|
||
func SendMsg(userId uint64, msg *Msg) {
|
||
conn := conns[userId]
|
||
if conn != nil {
|
||
bytes, _ := json.Marshal(msg)
|
||
conn.WriteMessage(websocket.TextMessage, bytes)
|
||
}
|
||
}
|