dev 分支暂存

backend/gateway/internal/ws/client.go

@@ -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 // 通道满则丢弃
+	}
+}