/** * package ws * 作用:处理 WebSocket 协议层逻辑,并管理基于 ants 的协程工作池。 * 核心:将 I/O 读取与业务逻辑分离,防止海量连接导致 Goroutine 爆炸。 */ package ws import ( "encoding/json" "fmt" "log" "runtime/debug" "xk-websocket-v2/internal/manager" "xk-websocket-v2/internal/model" "xk-websocket-v2/internal/service" "github.com/panjf2000/ants/v2" ) // ========================================== // 1. 配置常量与全局变量 // ========================================== const ( // MaxWorkerPoolSize 定义最大的业务处理协程数 // 即使有 100万 个连接,同时也只有 5000 个协程在处理具体的业务逻辑(发消息、存库等) // 其余连接处于空闲或等待 I/O 状态,极大节省内存。 MaxWorkerPoolSize = 5000 ) // GlobalWorkerPool 全局协程池实例 var GlobalWorkerPool *ants.PoolWithFunc // ========================================== // 2. 任务结构定义 // ========================================== /** * Job * 结构体:代表一个具体的待处理任务 * 从 WebSocket 读到的原始 bytes 和对应的 Client 对象打包在一起 */ type Job struct { Client *manager.Client // 消息来源客户端 Payload []byte // 原始消息数据 } // ========================================== // 3. 初始化与启动 // ========================================== /** * StartWorkerPool * 功能:初始化 ants 协程池 * 注意:必须在 main.go 中尽早调用 */ func StartWorkerPool() { var err error // NewPoolWithFunc 创建一个带任务处理函数的池子 // 参数1: 池容量 // 参数2: 具体的任务处理函数 (消费者逻辑) GlobalWorkerPool, err = ants.NewPoolWithFunc(MaxWorkerPoolSize, func(i interface{}) { // 类型断言:将接口转换为具体的 Job 结构体 job, ok := i.(Job) if !ok { log.Printf("❌ [Worker] 类型断言失败") return } // 执行业务逻辑 handleMessage(job.Client, job.Payload) }) if err != nil { log.Fatalf("❌ [Worker] 启动协程池失败: %v", err) } log.Printf("⚙️ [WorkerPool] 启动成功 | 容量: %d", MaxWorkerPoolSize) } /** * StopWorkerPool * 功能:优雅关闭协程池 */ func StopWorkerPool() { if GlobalWorkerPool != nil { GlobalWorkerPool.Release() log.Println("🛑 [WorkerPool] 已释放") } } // ========================================== // 4. 业务处理逻辑 (消费者) // ========================================== /** * handleMessage * 功能:核心业务路由,解析 JSON 并调用 Service 层 * 特性:包含 Panic Recover,防止单个请求崩溃导致整个进程退出 */ func handleMessage(client *manager.Client, message []byte) { // 🛡️ Panic 恢复机制:这是生产环境必须的! defer func() { if r := recover(); r != nil { log.Printf("❌ [Worker] Panic 捕获: %v\n堆栈信息: %s", r, debug.Stack()) // 可选:发送错误日志到监控系统 (Sentry/Prometheus) } }() // 1. 基础格式校验 if len(message) == 0 { return } // 2. 解析通用载荷 (RequestType) var req model.WsPayload if err := json.Unmarshal(message, &req); err != nil { log.Printf("⚠️ [Worker] JSON解析失败: %v | 来源: %s", err, client.ID) return } // 3. 业务路由分发 switch req.RequestType { case "bind": // 处理用户绑定 (bind) // 格式: {"request_type": "bind", "data": {"user_id": "1001"}} var bindData model.BindReq if err := json.Unmarshal(jsonBytes(req.Data), &bindData); err == nil { // 调用 Service 层逻辑 service.ChatSvc.BindUser(client, bindData.UserID) } else { log.Printf("⚠️ [Worker] 绑定参数错误: %v", err) } case "send_message": // 处理消息发送 (send_message) // 包含:文本聊天、图片、视频通话信令等 var msgData model.SendMessageReq if err := json.Unmarshal(jsonBytes(req.Data), &msgData); err == nil { // 核心:调用 ChatService 处理消息(持久化、转发、多端同步) // 注意:这里传入的是 client 指针,以便 Service 层获取发送者的 UserID 和 ClientID service.ChatSvc.HandleUserMessage(client, &msgData) } else { log.Printf("⚠️ [Worker] 消息参数错误: %v", err) } case "heartbeat": // 处理心跳 (heartbeat) // 快速响应,不需要走复杂业务逻辑 // 直接往客户端写回 Pong 消息 pong := []byte(`{"request_type":"heartbeat_ack"}`) manager.Manager.SendToClient(client.ID, pong) default: log.Printf("❓ [Worker] 未知请求类型: %s", req.RequestType) } } // ========================================== // 5. 辅助工具 // ========================================== /** * jsonBytes * 功能:辅助函数,将 map[string]interface{} 或其他类型转回 []byte * 场景:因为 WsPayload.Data 是 interface{},在 Unmarshal 时可能被解析为 map,需要二次转换 */ func jsonBytes(v interface{}) []byte { if v == nil { return []byte("{}") } // 如果已经是 string,直接转换 if str, ok := v.(string); ok { return []byte(str) } // 否则重新序列化 b, _ := json.Marshal(v) return b } /** * PushTask * 功能:对外暴露的投递函数,供 client_manager 或 main.go 调用 * 作用:非阻塞地将任务扔进池子 * @param client 消息来源客户端 * @param payload 原始消息 */ func PushTask(client *manager.Client, payload []byte) error { if GlobalWorkerPool == nil { return fmt.Errorf("worker pool not initialized") } // 包装任务 job := Job{ Client: client, Payload: payload, } // Invoke 会从池里取一个 worker 来执行 NewPoolWithFunc 注册的函数 // 如果池满了,Invoke 会阻塞,直到有 worker 空闲 (或者返回错误,取决于配置) return GlobalWorkerPool.Invoke(job) }