nl-im-service/internal/ws/worker.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)
}