137 lines
3.1 KiB
Go
137 lines
3.1 KiB
Go
// Package websocket 提供通用的 WebSocket Hub，支持按 session_id 分房间广播
|
|
package websocket
|
|
|
|
import (
|
|
"encoding/json"
|
|
"log"
|
|
"sync"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
// Hub 管理所有 WebSocket 连接和房间
|
|
// 每个 session_id 对应一个房间，房间内的连接互相广播消息
|
|
type Hub struct {
|
|
mu sync.RWMutex
|
|
rooms map[string]map[*Conn]struct{} // session_id → 连接集合
|
|
register chan *Conn
|
|
unregister chan *Conn
|
|
}
|
|
|
|
// Conn 封装一条 WebSocket 连接
|
|
type Conn struct {
|
|
Hub *Hub
|
|
SessionID string
|
|
Role string // "A" 或 "B"，可选
|
|
RoleMu sync.Mutex
|
|
WS *websocket.Conn
|
|
Send chan []byte
|
|
closeOnce sync.Once
|
|
}
|
|
|
|
// NewHub 创建一个 Hub 并启动后台处理协程
|
|
func NewHub() *Hub {
|
|
h := &Hub{
|
|
rooms: make(map[string]map[*Conn]struct{}),
|
|
register: make(chan *Conn),
|
|
unregister: make(chan *Conn),
|
|
}
|
|
go h.run()
|
|
return h
|
|
}
|
|
|
|
func (h *Hub) run() {
|
|
for {
|
|
select {
|
|
case conn := <-h.register:
|
|
h.mu.Lock()
|
|
if h.rooms[conn.SessionID] == nil {
|
|
h.rooms[conn.SessionID] = make(map[*Conn]struct{})
|
|
}
|
|
h.rooms[conn.SessionID][conn] = struct{}{}
|
|
h.mu.Unlock()
|
|
|
|
case conn := <-h.unregister:
|
|
h.mu.Lock()
|
|
if room, ok := h.rooms[conn.SessionID]; ok {
|
|
delete(room, conn)
|
|
if len(room) == 0 {
|
|
delete(h.rooms, conn.SessionID)
|
|
}
|
|
}
|
|
// 用 sync.Once 保证 Send channel 只被 close 一次
|
|
conn.closeOnce.Do(func() { close(conn.Send) })
|
|
h.mu.Unlock()
|
|
}
|
|
}
|
|
}
|
|
|
|
// Register 将连接加入指定 session 的房间
|
|
func (h *Hub) Register(conn *Conn) {
|
|
h.register <- conn
|
|
}
|
|
|
|
// Unregister 将连接从房间移除（幂等，可重复调用）
|
|
func (h *Hub) Unregister(conn *Conn) {
|
|
h.unregister <- conn
|
|
}
|
|
|
|
// Broadcast 向指定 session 房间内的所有连接发送消息
|
|
// msg 会被序列化为 JSON
|
|
func (h *Hub) Broadcast(sessionID string, msg interface{}) {
|
|
data, err := json.Marshal(msg)
|
|
if err != nil {
|
|
log.Printf("[ws] marshal broadcast msg error: %v", err)
|
|
return
|
|
}
|
|
h.mu.RLock()
|
|
defer h.mu.RUnlock()
|
|
if room, ok := h.rooms[sessionID]; ok {
|
|
for conn := range room {
|
|
select {
|
|
case conn.Send <- data:
|
|
default:
|
|
// send channel 满了，丢弃，避免阻塞
|
|
log.Printf("[ws] send channel full for session %s, dropping", sessionID)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// RoomSize 返回指定 session 房间内的连接数（调试用）
|
|
func (h *Hub) RoomSize(sessionID string) int {
|
|
h.mu.RLock()
|
|
defer h.mu.RUnlock()
|
|
return len(h.rooms[sessionID])
|
|
}
|
|
|
|
// WritePump 从 Send channel 读消息并写入 WebSocket
|
|
// 在独立 goroutine 中运行，直到连接关闭
|
|
func WritePump(c *Conn) {
|
|
defer func() {
|
|
c.Hub.Unregister(c)
|
|
c.WS.Close()
|
|
}()
|
|
for msg := range c.Send {
|
|
if err := c.WS.WriteMessage(websocket.TextMessage, msg); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
// ReadPump 从 WebSocket 读消息并丢弃
|
|
// 生产环境如果需要客户端主动发消息可以在这里处理
|
|
// 当前只负责检测断连
|
|
func ReadPump(c *Conn) {
|
|
defer func() {
|
|
c.Hub.Unregister(c)
|
|
c.WS.Close()
|
|
}()
|
|
c.WS.SetPongHandler(func(string) error { return nil })
|
|
for {
|
|
if _, _, err := c.WS.ReadMessage(); err != nil {
|
|
return
|
|
}
|
|
}
|
|
}
|