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:
305
internal/grpc/runner/registry.go
Normal file
305
internal/grpc/runner/registry.go
Normal file
@@ -0,0 +1,305 @@
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user