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
|
||||
}
|
||||
}
|
||||
@@ -26,22 +26,20 @@ type RunnerConnection struct {
|
||||
type RunnerHub struct {
|
||||
runner.UnimplementedRunnerHubServer // 嵌入以实现 gRPC 接口
|
||||
|
||||
mu sync.RWMutex
|
||||
runners map[string]*RunnerConnection
|
||||
taskManager *TaskManager
|
||||
registerChan chan *RunnerConnection
|
||||
unregisterChan chan string
|
||||
logger *zap.Logger
|
||||
mu sync.RWMutex
|
||||
runners map[string]*RunnerConnection
|
||||
taskManager *TaskManager
|
||||
bus *TaskBus
|
||||
registry *RunnerRegistry
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewRunnerHub 创建新的 RunnerHub
|
||||
func NewRunnerHub(taskManager *TaskManager, logger *zap.Logger) *RunnerHub {
|
||||
return &RunnerHub{
|
||||
runners: make(map[string]*RunnerConnection),
|
||||
taskManager: taskManager,
|
||||
registerChan: make(chan *RunnerConnection, 100),
|
||||
unregisterChan: make(chan string, 100),
|
||||
logger: logger,
|
||||
runners: make(map[string]*RunnerConnection),
|
||||
taskManager: taskManager,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -137,6 +135,15 @@ func (h *RunnerHub) handleRegister(req *runner.RegisterRequest, stream runner.Ru
|
||||
h.runners[req.RunnerId] = conn
|
||||
h.mu.Unlock()
|
||||
|
||||
// 同步到 Redis 注册中心
|
||||
if h.registry != nil {
|
||||
if err := h.registry.Register(req.RunnerId, req.Version, req.Capabilities); err != nil {
|
||||
h.logger.Error("Failed to register runner in registry",
|
||||
zap.String("runner_id", req.RunnerId),
|
||||
zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
// 发送注册响应
|
||||
if err := conn.send(&runner.StreamMessage{
|
||||
Message: &runner.StreamMessage_RegisterResponse{
|
||||
@@ -161,6 +168,15 @@ func (h *RunnerHub) handleHeartbeat(conn *RunnerConnection, hb *runner.Heartbeat
|
||||
}
|
||||
conn.LastSeen = time.Now()
|
||||
|
||||
// 刷新 Redis 注册中心的 TTL
|
||||
if h.registry != nil {
|
||||
if err := h.registry.RefreshHeartbeat(conn.ID); err != nil {
|
||||
h.logger.Debug("Failed to refresh runner heartbeat in registry",
|
||||
zap.String("runner_id", conn.ID),
|
||||
zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
if err := conn.send(&runner.StreamMessage{
|
||||
Message: &runner.StreamMessage_HeartbeatAck{
|
||||
HeartbeatAck: &runner.HeartbeatAck{
|
||||
@@ -179,6 +195,12 @@ func (h *RunnerHub) handleTaskResult(result *runner.TaskResult) {
|
||||
zap.String("task_id", result.TaskId),
|
||||
zap.String("status", result.Status.String()))
|
||||
|
||||
// 如果有 Bus，委托给 Bus 处理（可能需要转发到远程实例）
|
||||
if h.bus != nil {
|
||||
h.bus.HandleResult(result)
|
||||
return
|
||||
}
|
||||
|
||||
if h.taskManager != nil {
|
||||
h.taskManager.HandleResult(result)
|
||||
}
|
||||
@@ -187,12 +209,19 @@ func (h *RunnerHub) handleTaskResult(result *runner.TaskResult) {
|
||||
// unregisterRunner 注销 Runner
|
||||
func (h *RunnerHub) unregisterRunner(runnerID string) {
|
||||
h.mu.Lock()
|
||||
var capabilities map[string]string
|
||||
if conn, ok := h.runners[runnerID]; ok {
|
||||
capabilities = conn.Capabilities
|
||||
close(conn.sendChan)
|
||||
delete(h.runners, runnerID)
|
||||
}
|
||||
h.mu.Unlock()
|
||||
|
||||
// 从 Redis 注册中心移除
|
||||
if h.registry != nil && capabilities != nil {
|
||||
h.registry.Unregister(runnerID, capabilities)
|
||||
}
|
||||
|
||||
h.logger.Info("Runner unregistered", zap.String("runner_id", runnerID))
|
||||
}
|
||||
|
||||
@@ -257,6 +286,37 @@ func (h *RunnerHub) SetTaskManager(tm *TaskManager) {
|
||||
h.taskManager = tm
|
||||
}
|
||||
|
||||
// SetBus 设置 TaskBus（集群模式）
|
||||
func (h *RunnerHub) SetBus(bus *TaskBus) {
|
||||
h.bus = bus
|
||||
}
|
||||
|
||||
// SetRegistry 设置 RunnerRegistry（集群模式）
|
||||
func (h *RunnerHub) SetRegistry(registry *RunnerRegistry) {
|
||||
h.registry = registry
|
||||
}
|
||||
|
||||
// DispatchTask 实现 TaskDispatcher 接口（standalone 模式）
|
||||
func (h *RunnerHub) DispatchTask(task *runner.Task) error {
|
||||
conn, err := h.GetAvailableRunner(task.Type)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return h.SendTask(conn.ID, task)
|
||||
}
|
||||
|
||||
// GetRunnerByID 根据 ID 获取 Runner 连接
|
||||
func (h *RunnerHub) GetRunnerByID(runnerID string) (*RunnerConnection, error) {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
|
||||
conn, ok := h.runners[runnerID]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("runner not found: %s", runnerID)
|
||||
}
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
// send 向 RunnerConnection 发送消息
|
||||
func (c *RunnerConnection) send(msg *runner.StreamMessage) error {
|
||||
select {
|
||||
@@ -270,8 +330,11 @@ func (c *RunnerConnection) send(msg *runner.StreamMessage) error {
|
||||
// CheckStaleRunners 检查并清理超时的 Runner 连接
|
||||
func (h *RunnerHub) CheckStaleRunners(timeout time.Duration) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
|
||||
var toRemove []struct {
|
||||
id string
|
||||
capabilities map[string]string
|
||||
}
|
||||
now := time.Now()
|
||||
for id, conn := range h.runners {
|
||||
if now.Sub(conn.LastSeen) > timeout {
|
||||
@@ -280,6 +343,18 @@ func (h *RunnerHub) CheckStaleRunners(timeout time.Duration) {
|
||||
zap.Time("last_seen", conn.LastSeen))
|
||||
close(conn.sendChan)
|
||||
delete(h.runners, id)
|
||||
toRemove = append(toRemove, struct {
|
||||
id string
|
||||
capabilities map[string]string
|
||||
}{id: id, capabilities: conn.Capabilities})
|
||||
}
|
||||
}
|
||||
h.mu.Unlock()
|
||||
|
||||
// 从 Redis 注册中心批量移除
|
||||
if h.registry != nil {
|
||||
for _, r := range toRemove {
|
||||
h.registry.Unregister(r.id, r.capabilities)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
@@ -23,12 +23,11 @@ type Server struct {
|
||||
grpcSrv *grpc.Server
|
||||
}
|
||||
|
||||
// NewServer 创建新的 gRPC 服务器
|
||||
// NewServer 创建新的 gRPC 服务器（内部构造 Hub + TaskManager）
|
||||
func NewServer(cfg *config.GRPCConfig, logger *zap.Logger) *Server {
|
||||
// 创建 TaskManager 和 RunnerHub，解决循环依赖
|
||||
taskManager := NewTaskManager(nil, 60*time.Second, logger)
|
||||
hub := NewRunnerHub(taskManager, logger)
|
||||
taskManager.SetHub(hub)
|
||||
taskManager.SetDispatcher(hub)
|
||||
|
||||
return &Server{
|
||||
cfg: cfg,
|
||||
@@ -37,6 +36,15 @@ func NewServer(cfg *config.GRPCConfig, logger *zap.Logger) *Server {
|
||||
}
|
||||
}
|
||||
|
||||
// NewServerWithDeps 使用外部依赖创建 gRPC 服务器（Wire DI 使用）
|
||||
func NewServerWithDeps(cfg *config.GRPCConfig, logger *zap.Logger, hub *RunnerHub) *Server {
|
||||
return &Server{
|
||||
cfg: cfg,
|
||||
logger: logger,
|
||||
hub: hub,
|
||||
}
|
||||
}
|
||||
|
||||
// Start 启动 gRPC 服务器
|
||||
func (s *Server) Start() error {
|
||||
if !s.cfg.Enabled {
|
||||
|
||||
@@ -16,6 +16,11 @@ import (
|
||||
// TaskCallback 任务完成回调函数类型
|
||||
type TaskCallback func(result *runner.TaskResult)
|
||||
|
||||
// TaskDispatcher 任务分发接口
|
||||
type TaskDispatcher interface {
|
||||
DispatchTask(task *runner.Task) error
|
||||
}
|
||||
|
||||
// PendingTask 表示一个等待处理的任务
|
||||
type PendingTask struct {
|
||||
Task *runner.Task
|
||||
@@ -25,20 +30,20 @@ type PendingTask struct {
|
||||
|
||||
// TaskManager 管理任务的生命周期
|
||||
type TaskManager struct {
|
||||
mu sync.RWMutex
|
||||
pending map[string]*PendingTask
|
||||
hub *RunnerHub
|
||||
timeout time.Duration
|
||||
logger *zap.Logger
|
||||
mu sync.RWMutex
|
||||
pending map[string]*PendingTask
|
||||
dispatcher TaskDispatcher
|
||||
timeout time.Duration
|
||||
logger *zap.Logger
|
||||
}
|
||||
|
||||
// NewTaskManager 创建新的任务管理器
|
||||
func NewTaskManager(hub *RunnerHub, timeout time.Duration, logger *zap.Logger) *TaskManager {
|
||||
func NewTaskManager(dispatcher TaskDispatcher, timeout time.Duration, logger *zap.Logger) *TaskManager {
|
||||
return &TaskManager{
|
||||
pending: make(map[string]*PendingTask),
|
||||
hub: hub,
|
||||
timeout: timeout,
|
||||
logger: logger,
|
||||
pending: make(map[string]*PendingTask),
|
||||
dispatcher: dispatcher,
|
||||
timeout: timeout,
|
||||
logger: logger,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -69,17 +74,10 @@ func (tm *TaskManager) SubmitTask(ctx context.Context, taskType runner.TaskType,
|
||||
zap.String("task_id", task.TaskId),
|
||||
zap.String("type", taskType.String()))
|
||||
|
||||
// 获取可用 Runner
|
||||
conn, err := tm.hub.GetAvailableRunner(taskType)
|
||||
if err != nil {
|
||||
// 通过 dispatcher 分发任务
|
||||
if err := tm.dispatcher.DispatchTask(task); err != nil {
|
||||
tm.cleanupTask(task.TaskId)
|
||||
return nil, fmt.Errorf("no available runner: %w", err)
|
||||
}
|
||||
|
||||
// 发送任务
|
||||
if err := tm.hub.SendTask(conn.ID, task); err != nil {
|
||||
tm.cleanupTask(task.TaskId)
|
||||
return nil, fmt.Errorf("failed to send task: %w", err)
|
||||
return nil, fmt.Errorf("failed to dispatch task: %w", err)
|
||||
}
|
||||
|
||||
// 启动超时检查
|
||||
@@ -203,9 +201,9 @@ func (tm *TaskManager) GetPendingTaskCount() int {
|
||||
return len(tm.pending)
|
||||
}
|
||||
|
||||
// SetHub 设置 RunnerHub（用于解决循环依赖）
|
||||
func (tm *TaskManager) SetHub(hub *RunnerHub) {
|
||||
tm.hub = hub
|
||||
// SetDispatcher 设置任务分发器（用于解决循环依赖）
|
||||
func (tm *TaskManager) SetDispatcher(dispatcher TaskDispatcher) {
|
||||
tm.dispatcher = dispatcher
|
||||
}
|
||||
|
||||
// CleanupStaleTasks 清理超时的任务
|
||||
|
||||
Reference in New Issue
Block a user