package runner import ( "context" "encoding/json" "fmt" "sync" "with_you/proto/runner" "github.com/redis/go-redis/v9" "go.uber.org/zap" ) // taskBusMsg Redis 任务频道消息格式 type taskBusMsg struct { OriginInstanceID string `json:"origin_instance_id"` TargetRunnerID string `json:"target_runner_id"` TaskID string `json:"task_id"` TaskType string `json:"task_type"` TaskPayload json.RawMessage `json:"task_payload"` CreatedAt int64 `json:"created_at"` TimeoutSeconds int32 `json:"timeout_seconds"` } // resultBusMsg Redis 结果频道消息格式 type resultBusMsg struct { OriginInstanceID string `json:"origin_instance_id"` TaskID string `json:"task_id"` Status string `json:"status"` CompletedAt int64 `json:"completed_at"` Data json.RawMessage `json:"data"` ErrorCode string `json:"error_code"` ErrorMessage string `json:"error_message"` } // TaskBus 组合 RunnerHub + Redis Pub/Sub,实现跨实例任务分发和结果路由 type TaskBus struct { hub *RunnerHub registry *RunnerRegistry rdb *redis.Client instanceID string taskChan string resultChan string logger *zap.Logger remoteTasks sync.Map // taskID → originInstanceID(从其他实例接收的任务) sub *redis.PubSub cancel context.CancelFunc mu sync.Mutex closed bool } // NewTaskBus 创建 TaskBus 实例 func NewTaskBus(hub *RunnerHub, rdb *redis.Client, registry *RunnerRegistry, instanceID, taskChan, resultChan string, logger *zap.Logger) *TaskBus { return &TaskBus{ hub: hub, registry: registry, rdb: rdb, instanceID: instanceID, taskChan: taskChan, resultChan: resultChan, logger: logger, } } // Start 启动 Redis 订阅 func (b *TaskBus) Start() error { if b.rdb == nil { b.logger.Info("TaskBus: no Redis client, running in standalone mode") return nil } ctx, cancel := context.WithCancel(context.Background()) b.cancel = cancel sub := b.rdb.Subscribe(ctx, b.taskChan, b.resultChan) b.sub = sub if _, err := sub.Receive(ctx); err != nil { cancel() return err } go b.subscribeLoop(ctx, sub) b.logger.Info("TaskBus started", zap.String("instance_id", b.instanceID), zap.String("task_channel", b.taskChan), zap.String("result_channel", b.resultChan), ) return nil } // Stop 关闭 TaskBus func (b *TaskBus) Stop() error { b.mu.Lock() defer b.mu.Unlock() if b.closed { return nil } b.closed = true if b.cancel != nil { b.cancel() } if b.sub != nil { b.sub.Close() } if b.registry != nil { b.registry.Stop() } b.logger.Info("TaskBus stopped", zap.String("instance_id", b.instanceID)) return nil } // InstanceID 返回实例 ID func (b *TaskBus) InstanceID() string { return b.instanceID } // DispatchTask 实现 TaskDispatcher 接口 // 本地优先 → Redis 注册中心查找 → 远程发布 func (b *TaskBus) DispatchTask(task *runner.Task) error { taskTypeStr := task.Type.String() // 1. 尝试本地查找 conn, err := b.hub.GetAvailableRunner(task.Type) if err == nil && conn != nil { return b.hub.SendTask(conn.ID, task) } // 2. 查询 Redis 注册中心 if b.registry != nil { runnerID, instanceID, err := b.registry.GetAvailableRunner(taskTypeStr) if err != nil { return fmt.Errorf("no available runner: %w", err) } // 如果 runner 在本实例(可能本地 map 还没更新),直接发送 if instanceID == b.instanceID { return b.hub.SendTask(runnerID, task) } // 发布到远程实例 return b.publishTask(task, runnerID) } return fmt.Errorf("no available runner for task type: %s", taskTypeStr) } // HandleResult 处理任务结果 // 如果任务来自远程实例,转发结果;否则交给 TaskManager func (b *TaskBus) HandleResult(result *runner.TaskResult) { originInstance, ok := b.remoteTasks.Load(result.TaskId) if ok { b.remoteTasks.Delete(result.TaskId) // 任务来自远程实例,转发结果 if originStr, ok := originInstance.(string); ok && originStr != b.instanceID { b.publishResult(result, originStr) return } } // 本地任务,交给 TaskManager if b.hub.taskManager != nil { b.hub.taskManager.HandleResult(result) } } // publishTask 发布任务到 Redis 频道 func (b *TaskBus) publishTask(task *runner.Task, targetRunnerID string) error { msg := taskBusMsg{ OriginInstanceID: b.instanceID, TargetRunnerID: targetRunnerID, TaskID: task.TaskId, TaskType: task.Type.String(), TaskPayload: task.Payload, CreatedAt: task.CreatedAt, TimeoutSeconds: task.TimeoutSeconds, } data, err := json.Marshal(msg) if err != nil { return fmt.Errorf("failed to marshal task message: %w", err) } if err := b.rdb.Publish(context.Background(), b.taskChan, data).Err(); err != nil { return fmt.Errorf("failed to publish task to Redis: %w", err) } return nil } // publishResult 发布结果到 Redis 频道 func (b *TaskBus) publishResult(result *runner.TaskResult, originInstanceID string) error { var data json.RawMessage if len(result.Data) > 0 { data = result.Data } else { data = json.RawMessage("{}") } msg := resultBusMsg{ OriginInstanceID: originInstanceID, TaskID: result.TaskId, Status: result.Status.String(), CompletedAt: result.CompletedAt, Data: data, ErrorCode: result.ErrorCode, ErrorMessage: result.ErrorMessage, } raw, err := json.Marshal(msg) if err != nil { b.logger.Error("Failed to marshal result message", zap.Error(err)) return err } if err := b.rdb.Publish(context.Background(), b.resultChan, raw).Err(); err != nil { b.logger.Error("Failed to publish result to Redis", zap.Error(err)) return err } return nil } // subscribeLoop 处理 Redis 订阅消息 func (b *TaskBus) subscribeLoop(ctx context.Context, sub *redis.PubSub) { ch := sub.Channel() for { select { case <-ctx.Done(): return case msg, ok := <-ch: if !ok { return } switch msg.Channel { case b.taskChan: b.handleRemoteTask(msg) case b.resultChan: b.handleRemoteResult(msg) } } } } // handleRemoteTask 处理从其他实例发来的任务 func (b *TaskBus) handleRemoteTask(msg *redis.Message) { var taskMsg taskBusMsg if err := json.Unmarshal([]byte(msg.Payload), &taskMsg); err != nil { b.logger.Error("TaskBus: failed to unmarshal task message", zap.Error(err)) return } // 过滤自己发送的消息 if taskMsg.OriginInstanceID == b.instanceID { return } // 检查目标 Runner 是否在本实例 conn, err := b.hub.GetRunnerByID(taskMsg.TargetRunnerID) if err != nil { b.logger.Warn("TaskBus: target runner not found locally", zap.String("runner_id", taskMsg.TargetRunnerID), zap.String("task_id", taskMsg.TaskID)) return } // 记录任务的来源实例,用于结果回传 b.remoteTasks.Store(taskMsg.TaskID, taskMsg.OriginInstanceID) // 构造 Task 并发送 task := &runner.Task{ TaskId: taskMsg.TaskID, Type: parseTaskType(taskMsg.TaskType), CreatedAt: taskMsg.CreatedAt, TimeoutSeconds: taskMsg.TimeoutSeconds, Payload: taskMsg.TaskPayload, } if err := b.hub.SendTask(conn.ID, task); err != nil { b.logger.Error("TaskBus: failed to send remote task to runner", zap.String("runner_id", conn.ID), zap.String("task_id", taskMsg.TaskID), zap.Error(err)) b.remoteTasks.Delete(taskMsg.TaskID) } } // handleRemoteResult 处理从其他实例发来的任务结果 func (b *TaskBus) handleRemoteResult(msg *redis.Message) { var resultMsg resultBusMsg if err := json.Unmarshal([]byte(msg.Payload), &resultMsg); err != nil { b.logger.Error("TaskBus: failed to unmarshal result message", zap.Error(err)) return } // 过滤自己发送的消息 if resultMsg.OriginInstanceID == b.instanceID { return } // 只有 origin 匹配本实例才需要处理 // (originInstanceID 标识的是发起任务的实例) // 收到其他实例转发的结果,且 origin 是我们 → 交给 TaskManager // 但这里收到的是其他实例广播的结果,我们只是中继 // 实际上 result 是发到 runner:result channel 的,所有实例都会收到 // 我们只关心 origin 是本实例的结果 if resultMsg.OriginInstanceID != b.instanceID { return } // 交给 TaskManager 处理 result := &runner.TaskResult{ TaskId: resultMsg.TaskID, Status: parseTaskStatus(resultMsg.Status), CompletedAt: resultMsg.CompletedAt, Data: resultMsg.Data, ErrorCode: resultMsg.ErrorCode, ErrorMessage: resultMsg.ErrorMessage, } if b.hub.taskManager != nil { b.hub.taskManager.HandleResult(result) } } // parseTaskType 将字符串转换为 runner.TaskType 枚举 func parseTaskType(s string) runner.TaskType { switch s { case "TASK_TYPE_LOGIN": return runner.TaskType_TASK_TYPE_LOGIN case "TASK_TYPE_GET_SCHEDULE": return runner.TaskType_TASK_TYPE_GET_SCHEDULE case "TASK_TYPE_GET_GRADES": return runner.TaskType_TASK_TYPE_GET_GRADES case "TASK_TYPE_GET_EXAMS": return runner.TaskType_TASK_TYPE_GET_EXAMS case "TASK_TYPE_GET_USER_INFO": return runner.TaskType_TASK_TYPE_GET_USER_INFO default: return runner.TaskType_TASK_TYPE_UNKNOWN } } // parseTaskStatus 将字符串转换为 runner.TaskStatus 枚举 func parseTaskStatus(s string) runner.TaskStatus { switch s { case "TASK_STATUS_SUCCESS": return runner.TaskStatus_TASK_STATUS_SUCCESS case "TASK_STATUS_FAILED": return runner.TaskStatus_TASK_STATUS_FAILED case "TASK_STATUS_TIMEOUT": return runner.TaskStatus_TASK_STATUS_TIMEOUT case "TASK_STATUS_CANCELLED": return runner.TaskStatus_TASK_STATUS_CANCELLED default: return runner.TaskStatus_TASK_STATUS_UNKNOWN } }