208 lines
5.7 KiB
Go
208 lines
5.7 KiB
Go
/**
|
||
* 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)
|
||
}
|