package ws import ( "encoding/json" "sync/atomic" "time" "github.com/echochat/backend/pkg/logs" "github.com/gorilla/websocket" "go.uber.org/zap" ) const ( writeWait = 10 * time.Second // 写超时 pongWait = 60 * time.Second // 等待 pong 的超时 pingPeriod = 30 * time.Second // 心跳发送间隔(必须小于 pongWait) maxMessageSize = 4096 // 单条消息最大字节数 sendBufSize = 256 // 发送缓冲区大小 ) // DisconnectHandler 断线回调函数类型 type DisconnectHandler func(userID int64) // Client 封装单个 WebSocket 客户端连接 // 每个连接持有两个 goroutine:readPump(读取客户端消息)和 writePump(写入消息到客户端) type Client struct { hub *Hub conn *websocket.Conn send chan []byte // 待发送消息缓冲队列 UserID int64 // 关联的用户 ID onDisconnect DisconnectHandler // 断线回调(清理在线状态等) closedByHub int32 // 原子标记:是否被 Hub 主动踢出(用于避免重复连接时的竞态清理) } // NewClient 创建客户端实例 func NewClient(hub *Hub, conn *websocket.Conn, userID int64) *Client { return &Client{ hub: hub, conn: conn, send: make(chan []byte, sendBufSize), UserID: userID, } } // SetOnDisconnect 设置断线回调 func (c *Client) SetOnDisconnect(handler DisconnectHandler) { c.onDisconnect = handler } // SetClosedByHub 标记此连接被 Hub 主动关闭(重复连接踢出场景) // 线程安全,使用 atomic 操作 func (c *Client) SetClosedByHub() { atomic.StoreInt32(&c.closedByHub, 1) } // IsClosedByHub 检查此连接是否被 Hub 主动关闭 // 断线回调中使用此方法判断是否需要执行下线清理 func (c *Client) IsClosedByHub() bool { return atomic.LoadInt32(&c.closedByHub) == 1 } // MessageHandler 消息处理回调函数类型 type MessageHandler func(client *Client, msg *Message) // ReadPump 读取客户端消息的循环 // 当连接断开或出错时退出,并触发注销流程 func (c *Client) ReadPump(onMessage MessageHandler) { defer func() { c.hub.Unregister(c) c.conn.Close() if c.onDisconnect != nil { c.onDisconnect(c.UserID) } }() 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 { _, rawMsg, err := c.conn.ReadMessage() if err != nil { if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) { logs.Warn(nil, "ws.client.ReadPump", "WebSocket 异常关闭", zap.Int64("user_id", c.UserID), zap.Error(err)) } return } var msg Message if err := json.Unmarshal(rawMsg, &msg); err != nil { logs.Warn(nil, "ws.client.ReadPump", "消息格式解析失败", zap.Int64("user_id", c.UserID), zap.Error(err)) continue } if onMessage != nil { onMessage(c, &msg) } } } // WritePump 向客户端写入消息的循环 // 从 send channel 读取消息并写入 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 { c.conn.WriteMessage(websocket.CloseMessage, []byte{}) return } if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil { logs.Warn(nil, "ws.client.WritePump", "写入消息失败", zap.Int64("user_id", c.UserID), zap.Error(err)) return } case <-ticker.C: c.conn.SetWriteDeadline(time.Now().Add(writeWait)) if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } } // Send 向客户端发送消息(非阻塞,缓冲区满时丢弃) func (c *Client) Send(data []byte) bool { select { case c.send <- data: return true default: logs.Warn(nil, "ws.client.Send", "发送缓冲区已满,丢弃消息", zap.Int64("user_id", c.UserID)) return false } }