dual check v0.1

2026-08-29 23:34:01 +08:00
1 parent 78bad1fa3d
commit c137299016
@@ -0,0 +1,136 @@
+// 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
+		}
+	}
+}