feat(runner): implement cluster mode with Redis-based task dispatching and registry
Introduce a distributed task execution system for gRPC runners. This includes: - A `TaskBus` for cluster-mode task dispatching via Redis Pub/Sub. - A `RunnerRegistry` to track active runner instances in Redis. - Support for both standalone and cluster modes via configuration. - Refactored `RunnerHub` and `TaskManager` to use a `TaskDispatcher` interface, decoupling task submission from local execution. - Added distributed locking to `HotRankWorker` to ensure only one instance performs periodic hot rank recalculations in a multi-instance deployment. - Updated dependency injection (Wire) to support the new runner components and registry.
This commit is contained in:
365
internal/grpc/runner/bus.go
Normal file
365
internal/grpc/runner/bus.go
Normal file
@@ -0,0 +1,365 @@
|
||||
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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user