dev 分支暂存

2026-05-16 08:26:56 +08:00
parent 58c8caa570
commit eb4129176c
71 changed files with 8474 additions and 214 deletions
@@ -0,0 +1,138 @@
+package ws
+
+import (
+	"encoding/json"
+	"log"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+const (
+	// 写入超时
+	writeWait = 10 * time.Second
+
+	// 读取pong超时
+	pongWait = 60 * time.Second
+
+	// pong发送后等待下一次ping的间隔
+	pingPeriod = (pongWait * 9) / 10
+
+	// 最大消息大小
+	maxMessageSize = 65536
+)
+
+// Client WebSocket客户端
+type Client struct {
+	Hub      *Hub
+	Conn     *websocket.Conn
+	Send     chan []byte
+	UserID   string
+	SessionID string
+}
+
+// NewClient 创建WebSocket客户端
+func NewClient(hub *Hub, conn *websocket.Conn, userID, sessionID string) *Client {
+	return &Client{
+		Hub:       hub,
+		Conn:      conn,
+		Send:      make(chan []byte, 256),
+		UserID:    userID,
+		SessionID: sessionID,
+	}
+}
+
+// ReadPump 读取协程 —— 从WebSocket连接读取消息
+func (c *Client) ReadPump(onMessage func(client *Client, msg ClientMessage)) {
+	defer func() {
+		c.Hub.unregister <- c
+		c.Conn.Close()
+	}()
+
+	c.Conn.SetReadLimit(maxMessageSize)
+	c.Conn.SetReadDeadline(time.Now().Add(pongWait))
+	c.Conn.SetPongHandler(func(string) error {
+		c.Conn.SetReadDeadline(time.Now().Add(pongWait))
+		return nil
+	})
+
+	for {
+		_, rawMessage, err := c.Conn.ReadMessage()
+		if err != nil {
+			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure, websocket.CloseAbnormalClosure) {
+				log.Printf("[WS] 读取错误: user=%s err=%v", c.UserID, err)
+			}
+			break
+		}
+
+		// 解析消息
+		var msg ClientMessage
+		if err := json.Unmarshal(rawMessage, &msg); err != nil {
+			log.Printf("[WS] 消息解析失败: user=%s err=%v", c.UserID, err)
+			continue
+		}
+
+		// 处理ping
+		if msg.Type == "ping" {
+			pongMsg := ServerMessage{
+				Type:      "pong",
+				Timestamp: time.Now().UnixMilli(),
+			}
+			data, _ := json.Marshal(pongMsg)
+			c.Send <- data
+			continue
+		}
+
+		// 调用消息处理器
+		if onMessage != nil {
+			onMessage(c, msg)
+		}
+	}
+}
+
+// WritePump 写入协程 —— 向WebSocket连接写入消息
+func (c *Client) WritePump() {
+	ticker := time.NewTicker(pingPeriod)
+	defer func() {
+		ticker.Stop()
+		c.Conn.Close()
+	}()
+
+	for {
+		select {
+		case message, ok := <-c.Send:
+			c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if !ok {
+				// Hub关闭了通道
+				c.Conn.WriteMessage(websocket.CloseMessage, []byte{})
+				return
+			}
+
+			if err := c.Conn.WriteMessage(websocket.TextMessage, message); err != nil {
+				log.Printf("[WS] 写入错误: user=%s err=%v", c.UserID, err)
+				return
+			}
+
+		case <-ticker.C:
+			c.Conn.SetWriteDeadline(time.Now().Add(writeWait))
+			if err := c.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+				return
+			}
+		}
+	}
+}
+
+// SendMessage 向客户端发送消息
+func (c *Client) SendMessage(msg ServerMessage) error {
+	data, err := json.Marshal(msg)
+	if err != nil {
+		return err
+	}
+
+	select {
+	case c.Send <- data:
+		return nil
+	default:
+		return nil // 通道满则丢弃
+	}
+}
@@ -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
+}