package runner import ( "context" "encoding/json" "fmt" "sync" "time" "github.com/redis/go-redis/v9" "go.uber.org/zap" ) // RunnerRegistry 分布式 Runner 注册中心 // 使用 Redis Hash 存储 Runner 信息,Redis Set 按能力索引 type RunnerRegistry struct { rdb *redis.Client instanceID string registryPrefix string capPrefix string onlineTTL time.Duration heartbeatInterval time.Duration localRunners sync.Map // runnerID → map[string]string (capabilities) logger *zap.Logger cancel context.CancelFunc } // runnerInfo 存储在 Redis Hash 中的 Runner 信息 type runnerInfo struct { InstanceID string `json:"instance_id"` Version string `json:"version"` Capabilities string `json:"capabilities"` // JSON-encoded map } // NewRunnerRegistry 创建 Runner 注册中心 func NewRunnerRegistry(rdb *redis.Client, instanceID, registryPrefix, capPrefix string, onlineTTL, heartbeatInterval time.Duration, logger *zap.Logger) *RunnerRegistry { return &RunnerRegistry{ rdb: rdb, instanceID: instanceID, registryPrefix: registryPrefix, capPrefix: capPrefix, onlineTTL: onlineTTL, heartbeatInterval: heartbeatInterval, logger: logger, } } // Register 注册 Runner 到分布式注册中心 func (r *RunnerRegistry) Register(runnerID, version string, capabilities map[string]string) error { if r.rdb == nil { return nil } ctx := context.Background() capJSON, _ := json.Marshal(capabilities) key := r.registryPrefix + runnerID info := runnerInfo{ InstanceID: r.instanceID, Version: version, Capabilities: string(capJSON), } infoJSON, _ := json.Marshal(info) pipe := r.rdb.Pipeline() pipe.Set(ctx, key, infoJSON, r.onlineTTL) // 为每个能力创建索引 taskTypeStr := parseTaskTypes(capabilities) for _, tt := range taskTypeStr { capKey := r.capPrefix + tt pipe.SAdd(ctx, capKey, runnerID) pipe.Expire(ctx, capKey, r.onlineTTL*2) } if _, err := pipe.Exec(ctx); err != nil { return fmt.Errorf("failed to register runner in redis: %w", err) } r.localRunners.Store(runnerID, capabilities) return nil } // Unregister 从分布式注册中心移除 Runner func (r *RunnerRegistry) Unregister(runnerID string, capabilities map[string]string) { if r.rdb == nil { return } ctx := context.Background() key := r.registryPrefix + runnerID pipe := r.rdb.Pipeline() pipe.Del(ctx, key) taskTypeStr := parseTaskTypes(capabilities) for _, tt := range taskTypeStr { capKey := r.capPrefix + tt pipe.SRem(ctx, capKey, runnerID) } if _, err := pipe.Exec(ctx); err != nil { r.logger.Error("Failed to unregister runner from redis", zap.String("runner_id", runnerID), zap.Error(err)) } r.localRunners.Delete(runnerID) } // RefreshHeartbeat 刷新 Runner 的 TTL func (r *RunnerRegistry) RefreshHeartbeat(runnerID string) error { if r.rdb == nil { return nil } ctx := context.Background() key := r.registryPrefix + runnerID return r.rdb.Expire(ctx, key, r.onlineTTL).Err() } // GetAvailableRunner 查找支持指定任务类型的 Runner // 返回 runnerID 和所在 instanceID func (r *RunnerRegistry) GetAvailableRunner(taskType string) (runnerID, instanceID string, err error) { if r.rdb == nil { return "", "", fmt.Errorf("registry not available") } ctx := context.Background() capKey := r.capPrefix + taskType members, err := r.rdb.SMembers(ctx, capKey).Result() if err != nil { return "", "", fmt.Errorf("failed to query runner capabilities: %w", err) } if len(members) == 0 { return "", "", fmt.Errorf("no runner available for task type: %s", taskType) } // 随机选一个(伪随机,map遍历本身随机性) for _, rid := range members { key := r.registryPrefix + rid val, err := r.rdb.Get(ctx, key).Bytes() if err != nil { continue // 已过期,跳过 } var info runnerInfo if err := json.Unmarshal(val, &info); err != nil { continue } return rid, info.InstanceID, nil } return "", "", fmt.Errorf("no runner available for task type: %s", taskType) } // GetRunnerInstance 获取 Runner 所在的实例 ID func (r *RunnerRegistry) GetRunnerInstance(runnerID string) (string, error) { if r.rdb == nil { return "", fmt.Errorf("registry not available") } ctx := context.Background() key := r.registryPrefix + runnerID val, err := r.rdb.Get(ctx, key).Bytes() if err != nil { return "", fmt.Errorf("runner not found: %s", runnerID) } var info runnerInfo if err := json.Unmarshal(val, &info); err != nil { return "", fmt.Errorf("failed to parse runner info: %w", err) } return info.InstanceID, nil } // GetAllRunners 获取所有注册的 Runner 信息 func (r *RunnerRegistry) GetAllRunners() ([]map[string]any, error) { if r.rdb == nil { return nil, nil } ctx := context.Background() pattern := r.registryPrefix + "*" keys, err := r.rdb.Keys(ctx, pattern).Result() if err != nil { return nil, err } result := make([]map[string]any, 0, len(keys)) for _, key := range keys { val, err := r.rdb.Get(ctx, key).Bytes() if err != nil { continue } var info runnerInfo if err := json.Unmarshal(val, &info); err != nil { continue } var caps map[string]string json.Unmarshal([]byte(info.Capabilities), &caps) runnerID := key[len(r.registryPrefix):] result = append(result, map[string]any{ "id": runnerID, "instance_id": info.InstanceID, "version": info.Version, "capabilities": caps, }) } return result, nil } // StartHeartbeat 启动心跳续期 func (r *RunnerRegistry) StartHeartbeat() { if r.rdb == nil { return } ctx, cancel := context.WithCancel(context.Background()) r.cancel = cancel go func() { ticker := time.NewTicker(r.heartbeatInterval) defer ticker.Stop() for { select { case <-ctx.Done(): return case <-ticker.C: r.heartbeat() } } }() } func (r *RunnerRegistry) heartbeat() { ctx := context.Background() pipe := r.rdb.Pipeline() count := 0 r.localRunners.Range(func(key, _ any) bool { runnerID := key.(string) regKey := r.registryPrefix + runnerID pipe.Expire(ctx, regKey, r.onlineTTL) count++ return true }) if count > 0 { if _, err := pipe.Exec(ctx); err != nil { r.logger.Error("Failed to heartbeat runner registry", zap.Error(err)) } } } // Stop 停止心跳并清理本地 Runner 的 Redis 注册 func (r *RunnerRegistry) Stop() { if r.cancel != nil { r.cancel() } if r.rdb == nil { return } r.localRunners.Range(func(key, value any) bool { runnerID := key.(string) capabilities, _ := value.(map[string]string) r.Unregister(runnerID, capabilities) return true }) } // parseTaskTypes 从 capabilities map 中提取支持的任务类型列表 func parseTaskTypes(capabilities map[string]string) []string { var types []string for capKey, capValue := range capabilities { // key = "TASK_TYPE_GET_SCHEDULE", value = "true" if capValue == "true" && isTaskType(capKey) { types = append(types, capKey) continue } // key = "schedule", value = "TASK_TYPE_GET_SCHEDULE" if isTaskType(capValue) { types = append(types, capValue) } } return types } // isTaskType 检查字符串是否为有效的 TaskType 枚举值 func isTaskType(s string) bool { switch s { case "TASK_TYPE_LOGIN", "TASK_TYPE_GET_SCHEDULE", "TASK_TYPE_GET_GRADES", "TASK_TYPE_GET_EXAMS", "TASK_TYPE_GET_USER_INFO": return true } return false }