dev 分支暂存

backend/gateway/internal/ws/hub.go

@@ -0,0 +1,132 @@
+package ws
+
+import (
+	"log"
+	"sync"
+)
+
+// Hub WebSocket连接池
+type Hub struct {
+	mu         sync.RWMutex
+	clients    map[*Client]bool
+	broadcast  chan []byte
+	register   chan *Client
+	unregister chan *Client
+
+	// 按用户ID索引的客户端映射
+	userClients map[string]map[*Client]bool
+}
+
+// NewHub 创建WebSocket Hub
+func NewHub() *Hub {
+	return &Hub{
+		clients:     make(map[*Client]bool),
+		broadcast:   make(chan []byte, 256),
+		register:    make(chan *Client),
+		unregister:  make(chan *Client),
+		userClients: make(map[string]map[*Client]bool),
+	}
+}
+
+// Run 启动Hub主循环
+func (h *Hub) Run() {
+	for {
+		select {
+		case client := <-h.register:
+			h.mu.Lock()
+			h.clients[client] = true
+
+			// 用户索引
+			if h.userClients[client.UserID] == nil {
+				h.userClients[client.UserID] = make(map[*Client]bool)
+			}
+			h.userClients[client.UserID][client] = true
+			h.mu.Unlock()
+
+			log.Printf("[WS] 客户端连接: user=%s session=%s (当前连接数: %d)",
+				client.UserID, client.SessionID, len(h.clients))
+
+		case client := <-h.unregister:
+			h.mu.Lock()
+			if _, ok := h.clients[client]; ok {
+				delete(h.clients, client)
+				close(client.Send)
+
+				// 清理用户索引
+				if h.userClients[client.UserID] != nil {
+					delete(h.userClients[client.UserID], client)
+					if len(h.userClients[client.UserID]) == 0 {
+						delete(h.userClients, client.UserID)
+					}
+				}
+			}
+			h.mu.Unlock()
+
+			log.Printf("[WS] 客户端断开: user=%s session=%s (当前连接数: %d)",
+				client.UserID, client.SessionID, len(h.clients))
+
+		case message := <-h.broadcast:
+			h.mu.RLock()
+			for client := range h.clients {
+				select {
+				case client.Send <- message:
+				default:
+					// 客户端发送通道已满，跳过
+					close(client.Send)
+					delete(h.clients, client)
+				}
+			}
+			h.mu.RUnlock()
+		}
+	}
+}
+
+// SendToUser 向指定用户的所有连接发送消息
+func (h *Hub) SendToUser(userID string, message []byte) {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	if clients, ok := h.userClients[userID]; ok {
+		for client := range clients {
+			select {
+			case client.Send <- message:
+			default:
+				// 跳过阻塞的客户端
+			}
+		}
+	}
+}
+
+// SendToSession 向指定会话的连接发送消息
+func (h *Hub) SendToSession(userID, sessionID string, message []byte) {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	if clients, ok := h.userClients[userID]; ok {
+		for client := range clients {
+			if client.SessionID == sessionID {
+				select {
+				case client.Send <- message:
+				default:
+				}
+			}
+		}
+	}
+}
+
+// ClientCount 获取当前连接数
+func (h *Hub) ClientCount() int {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	return len(h.clients)
+}
+
+// UserClientCount 获取指定用户的连接数
+func (h *Hub) UserClientCount(userID string) int {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	if clients, ok := h.userClients[userID]; ok {
+		return len(clients)
+	}
+	return 0
+}