From 8aa85ca4b26da9837fb6d6360b5b6e7005e1944b Mon Sep 17 00:00:00 2001 From: lan Date: Thu, 7 May 2026 01:08:39 +0800 Subject: [PATCH] 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. --- cmd/server/app.go | 45 ++-- cmd/server/wire_gen.go | 11 +- internal/config/config.go | 13 + internal/config/runner.go | 28 ++ internal/dto/dto.go | 6 + internal/grpc/runner/bus.go | 365 +++++++++++++++++++++++++++ internal/grpc/runner/hub.go | 99 +++++++- internal/grpc/runner/registry.go | 305 ++++++++++++++++++++++ internal/grpc/runner/server.go | 14 +- internal/grpc/runner/task_manager.go | 44 ++-- internal/handler/group_handler.go | 44 ++++ internal/router/router.go | 1 + internal/service/hot_rank_worker.go | 136 ++++++++-- internal/wire/grpc.go | 91 ++++++- internal/wire/service.go | 9 +- 15 files changed, 1125 insertions(+), 86 deletions(-) create mode 100644 internal/config/runner.go create mode 100644 internal/grpc/runner/bus.go create mode 100644 internal/grpc/runner/registry.go diff --git a/cmd/server/app.go b/cmd/server/app.go index a3ef7ec..37cd8e3 100644 --- a/cmd/server/app.go +++ b/cmd/server/app.go @@ -18,15 +18,16 @@ import ( // App 应用程序结构体 type App struct { - Config *config.Config - DB *gorm.DB - Router *router.Router - PushService service.PushService - HotRankWorker *service.HotRankWorker - GRPCServer *runner.Server - Server *http.Server - Publisher ws.MessagePublisher - Logger *zap.Logger + Config *config.Config + DB *gorm.DB + Router *router.Router + PushService service.PushService + HotRankWorker *service.HotRankWorker + GRPCServer *runner.Server + Server *http.Server + Publisher ws.MessagePublisher + TaskDispatcher runner.TaskDispatcher + Logger *zap.Logger } // NewApp 创建应用程序 @@ -38,17 +39,19 @@ func NewApp( hotRankWorker *service.HotRankWorker, grpcServer *runner.Server, publisher ws.MessagePublisher, + taskDispatcher runner.TaskDispatcher, logger *zap.Logger, ) *App { return &App{ - Config: cfg, - DB: db, - Router: r, - PushService: pushService, - HotRankWorker: hotRankWorker, - GRPCServer: grpcServer, - Publisher: publisher, - Logger: logger, + Config: cfg, + DB: db, + Router: r, + PushService: pushService, + HotRankWorker: hotRankWorker, + GRPCServer: grpcServer, + Publisher: publisher, + TaskDispatcher: taskDispatcher, + Logger: logger, } } @@ -128,6 +131,14 @@ func (a *App) Stop(ctx context.Context) error { a.GRPCServer.Stop() } + // 4.1 停止 TaskBus（集群模式） + if a.TaskDispatcher != nil { + if bus, ok := a.TaskDispatcher.(*runner.TaskBus); ok { + a.Logger.Info("Stopping TaskBus") + bus.Stop() + } + } + // 5. 停止推送 Worker a.Logger.Info("Stopping push worker") a.PushService.StopPushWorker() diff --git a/cmd/server/wire_gen.go b/cmd/server/wire_gen.go index 2a24e2d..ff5bc72 100644 --- a/cmd/server/wire_gen.go +++ b/cmd/server/wire_gen.go @@ -101,8 +101,8 @@ func InitializeApp() (*App, error) { channelHandler := handler.NewChannelHandler(channelService) scheduleRepository := repository.NewScheduleRepository(db) scheduleService := wire.ProvideScheduleService(scheduleRepository) - server := wire.ProvideGRPCServer(config, logger) - taskManager := wire.ProvideTaskManager(server) + runnerHub := wire.ProvideRunnerHub(logger) + taskManager := wire.ProvideTaskManager(runnerHub, logger) scheduleSyncService := wire.ProvideScheduleSyncService(taskManager, scheduleRepository, logger) scheduleHandler := wire.ProvideScheduleHandler(scheduleService, scheduleSyncService) roleRepository := wire.ProvideRoleRepository(db) @@ -148,8 +148,11 @@ func InitializeApp() (*App, error) { tradeHandler := handler.NewTradeHandler(tradeService) wsHandler := wire.ProvideWSHandler(messagePublisher, chatService, groupService, jwtService, callService, userRepository) router := ProvideRouter(userRepository, userHandler, postHandler, commentHandler, messageHandler, notificationHandler, uploadHandler, jwtService, pushHandler, systemMessageHandler, groupHandler, stickerHandler, voteHandler, channelHandler, scheduleHandler, roleHandler, adminUserHandler, adminPostHandler, adminCommentHandler, adminGroupHandler, adminDashboardHandler, adminLogHandler, qrCodeHandler, materialHandler, callHandler, reportHandler, adminReportHandler, verificationHandler, adminVerificationHandler, adminProfileAuditHandler, setupHandler, tradeHandler, logService, userActivityService, casbinService, wsHandler) - hotRankWorker := wire.ProvideHotRankWorker(config, postRepository, channelRepository, cache) - app := NewApp(config, db, router, pushService, hotRankWorker, server, messagePublisher, logger) + hotRankWorker := wire.ProvideHotRankWorker(config, postRepository, channelRepository, cache, client) + server := wire.ProvideGRPCServer(config, logger, runnerHub) + runnerRegistry := wire.ProvideRunnerRegistry(config, client) + taskDispatcher := wire.ProvideTaskDispatcher(runnerHub, runnerRegistry, config, client, logger) + app := NewApp(config, db, router, pushService, hotRankWorker, server, messagePublisher, taskDispatcher, logger) return app, nil } diff --git a/internal/config/config.go b/internal/config/config.go index 5abbbcf..462b836 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -33,6 +33,7 @@ type Config struct { JPush JPushConfig `mapstructure:"jpush"` SetupSecret string `mapstructure:"setup_secret"` WebSocket WSConfig `mapstructure:"websocket"` + Runner RunnerConfig `mapstructure:"runner"` } // Load 加载配置文件 @@ -180,6 +181,15 @@ func Load(configPath string) (*Config, error) { viper.SetDefault("websocket.cluster.msg_channel", "ws:msg") viper.SetDefault("websocket.cluster.online_ttl", 60) viper.SetDefault("websocket.cluster.heartbeat_interval", 20) + // Runner 集群默认值（mode 和 instance_id 空则跟随 websocket） + viper.SetDefault("runner.mode", "") + viper.SetDefault("runner.cluster.instance_id", "") + viper.SetDefault("runner.cluster.task_channel", "runner:task") + viper.SetDefault("runner.cluster.result_channel", "runner:result") + viper.SetDefault("runner.cluster.registry_prefix", "runner:reg:") + viper.SetDefault("runner.cluster.cap_prefix", "runner:cap:") + viper.SetDefault("runner.cluster.registry_ttl", 90) + viper.SetDefault("runner.cluster.heartbeat_interval", 30) if err := viper.ReadInConfig(); err != nil { return nil, fmt.Errorf("failed to read config: %w", err) @@ -190,6 +200,9 @@ func Load(configPath string) (*Config, error) { return nil, fmt.Errorf("failed to unmarshal config: %w", err) } + // Runner mode 和 instance_id 空则跟随 WebSocket 配置 + cfg.Runner.ApplyWSDefaults(cfg.WebSocket) + // Convert seconds to duration cfg.JWT.AccessTokenExpire = time.Duration(viper.GetInt("jwt.access_token_expire")) * time.Second cfg.JWT.RefreshTokenExpire = time.Duration(viper.GetInt("jwt.refresh_token_expire")) * time.Second diff --git a/internal/config/runner.go b/internal/config/runner.go new file mode 100644 index 0000000..bee7c83 --- /dev/null +++ b/internal/config/runner.go @@ -0,0 +1,28 @@ +package config + +// RunnerConfig Runner 集群配置 +type RunnerConfig struct { + Mode string `mapstructure:"mode"` // "standalone" | "cluster"，空则跟随 websocket.mode + Cluster RunnerClusterConfig `mapstructure:"cluster"` +} + +// RunnerClusterConfig Runner 集群模式配置 +type RunnerClusterConfig struct { + InstanceID string `mapstructure:"instance_id"` // 空则跟随 websocket.cluster.instance_id + TaskChannel string `mapstructure:"task_channel"` + ResultChannel string `mapstructure:"result_channel"` + RegistryPrefix string `mapstructure:"registry_prefix"` + CapPrefix string `mapstructure:"cap_prefix"` + RegistryTTL int `mapstructure:"registry_ttl"` + HeartbeatInterval int `mapstructure:"heartbeat_interval"` +} + +// ApplyWSDefaults 用 WebSocket 配置填充 Runner 中未设置的共享字段 +func (r *RunnerConfig) ApplyWSDefaults(ws WSConfig) { + if r.Mode == "" { + r.Mode = ws.Mode + } + if r.Cluster.InstanceID == "" { + r.Cluster.InstanceID = ws.Cluster.InstanceID + } +} \ No newline at end of file diff --git a/internal/dto/dto.go b/internal/dto/dto.go index 55af1f9..06adc47 100644 --- a/internal/dto/dto.go +++ b/internal/dto/dto.go @@ -778,6 +778,12 @@ type SetGroupAvatarParams struct { Avatar string `json:"avatar"` // 头像URL } +// SetGroupDescriptionParams 设置群描述参数 +type SetGroupDescriptionParams struct { + GroupID string `json:"group_id"` // 群组ID + Description string `json:"description"` // 群描述 +} + // SetGroupLeaveParams 退出群组参数 type SetGroupLeaveParams struct { GroupID string `json:"group_id"` // 群组ID diff --git a/internal/grpc/runner/bus.go b/internal/grpc/runner/bus.go new file mode 100644 index 0000000..c577e01 --- /dev/null +++ b/internal/grpc/runner/bus.go @@ -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 + } +} \ No newline at end of file diff --git a/internal/grpc/runner/hub.go b/internal/grpc/runner/hub.go index 8baec1c..c7d81a1 100644 --- a/internal/grpc/runner/hub.go +++ b/internal/grpc/runner/hub.go @@ -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) } } } diff --git a/internal/grpc/runner/registry.go b/internal/grpc/runner/registry.go new file mode 100644 index 0000000..bed940c --- /dev/null +++ b/internal/grpc/runner/registry.go @@ -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 +} \ No newline at end of file diff --git a/internal/grpc/runner/server.go b/internal/grpc/runner/server.go index e6dba55..259bbb9 100644 --- a/internal/grpc/runner/server.go +++ b/internal/grpc/runner/server.go @@ -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 { diff --git a/internal/grpc/runner/task_manager.go b/internal/grpc/runner/task_manager.go index b3762eb..47d2901 100644 --- a/internal/grpc/runner/task_manager.go +++ b/internal/grpc/runner/task_manager.go @@ -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 清理超时的任务 diff --git a/internal/handler/group_handler.go b/internal/handler/group_handler.go index b78ada5..4969170 100644 --- a/internal/handler/group_handler.go +++ b/internal/handler/group_handler.go @@ -1764,6 +1764,50 @@ func (h *GroupHandler) HandleSetGroupAvatar(c *gin.Context) { response.Success(c, dto.GroupToResponse(group)) } +// HandleSetGroupDescription 设置群描述 +// PUT /api/v1/groups/:id/description +func (h *GroupHandler) HandleSetGroupDescription(c *gin.Context) { + userID := parseUserID(c) + if userID == "" { + response.Unauthorized(c, "") + return + } + + groupID := parseGroupID(c) + if groupID == "" { + response.BadRequest(c, "group_id is required") + return + } + + var params dto.SetGroupDescriptionParams + if err := c.ShouldBindJSON(¶ms); err != nil { + response.BadRequest(c, err.Error()) + return + } + + updates := map[string]any{ + "description": params.Description, + } + + err := h.groupService.UpdateGroup(userID, groupID, updates) + if err != nil { + if err == service.ErrNotGroupAdmin { + response.Forbidden(c, "没有权限修改群组信息") + return + } + if errors.Is(err, service.ErrGroupNotFound) { + response.NotFound(c, "群组不存在") + return + } + response.InternalServerError(c, err.Error()) + return + } + + // 获取更新后的群组信息 + group, _ := h.groupService.GetGroupByID(groupID) + response.Success(c, dto.GroupToResponse(group)) +} + // HandleSetGroupLeave 退出群组 // POST /api/v1/groups/:id/leave func (h *GroupHandler) HandleSetGroupLeave(c *gin.Context) { diff --git a/internal/router/router.go b/internal/router/router.go index 5254bfe..dae8d47 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -518,6 +518,7 @@ func (r *Router) setupRoutes() { groups.PUT("/:id/join-type", r.groupHandler.HandleSetJoinType) // 加群方式 groups.PUT("/:id/name", r.groupHandler.HandleSetGroupName) // 群名称 groups.PUT("/:id/avatar", r.groupHandler.HandleSetGroupAvatar) // 群头像 + groups.PUT("/:id/description", r.groupHandler.HandleSetGroupDescription) // 群描述 // 群公告 groups.POST("/:id/announcements", r.groupHandler.HandleCreateAnnouncement) // 创建公告 diff --git a/internal/service/hot_rank_worker.go b/internal/service/hot_rank_worker.go index b8ef04e..0873dfa 100644 --- a/internal/service/hot_rank_worker.go +++ b/internal/service/hot_rank_worker.go @@ -13,7 +13,7 @@ import ( "with_you/internal/config" "with_you/internal/repository" - "github.com/redis/go-redis/v9" + redislib "github.com/redis/go-redis/v9" "go.uber.org/zap" ) @@ -23,29 +23,84 @@ type hotRankScored struct { } // HotRankWorker 定时在候选池上重算热门分，写入 Redis ZSET（仅 TopN）及 top_ids（Redis+本地缓存） +// 多实例部署时通过 Redis 分布式锁保证仅主实例执行定时刷新，避免重复计算。 type HotRankWorker struct { cfg *config.Config postRepo repository.PostRepository channelRepo repository.ChannelRepository c cache.Cache + rdb *redislib.Client // 用于分布式锁（可为 nil，降级为本地运行） mu sync.Mutex stopOnce sync.Once stopCh chan struct{} doneCh chan struct{} } -func NewHotRankWorker(cfg *config.Config, postRepo repository.PostRepository, channelRepo repository.ChannelRepository, c cache.Cache) *HotRankWorker { +const ( + // hotRankLockKey Redis 分布式锁键 + hotRankLockKey = "hot_rank:leader_lock" + // hotRankLockTTL 锁持有时间（略大于刷新间隔，确保单实例持有） + hotRankLockTTL = 5 * time.Minute + // hotRankLockRetryInterval 竞争锁的重试间隔 + hotRankLockRetryInterval = 30 * time.Second +) + +func NewHotRankWorker(cfg *config.Config, postRepo repository.PostRepository, channelRepo repository.ChannelRepository, c cache.Cache, rdb *redislib.Client) *HotRankWorker { return &HotRankWorker{ cfg: cfg, postRepo: postRepo, channelRepo: channelRepo, c: c, + rdb: rdb, stopCh: make(chan struct{}), doneCh: make(chan struct{}), } } -// Start 启动定时任务（启动时立即执行一次） +// tryAcquireLock 尝试获取分布式锁，返回是否成功 +func (w *HotRankWorker) tryAcquireLock(ctx context.Context) bool { + if w.rdb == nil { + return true // 无 Redis 时降级为本地运行 + } + ok, err := w.rdb.SetNX(ctx, hotRankLockKey, "1", hotRankLockTTL).Result() + if err != nil { + zap.L().Warn("hot rank: failed to acquire leader lock, running locally", zap.Error(err)) + return true // Redis 出错时降级为本地运行 + } + return ok +} + +// keepLockAlive 周期性续期分布式锁 +func (w *HotRankWorker) keepLockAlive(ctx context.Context) { + if w.rdb == nil { + return + } + ticker := time.NewTicker(hotRankLockTTL / 2) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-w.stopCh: + return + case <-ticker.C: + w.rdb.Expire(ctx, hotRankLockKey, hotRankLockTTL) + } + } +} + +// releaseLock 释放分布式锁 +func (w *HotRankWorker) releaseLock() { + if w.rdb == nil { + return + } + w.rdb.Del(context.Background(), hotRankLockKey) +} + +// Start 启动定时任务 +// 多实例部署时通过 Redis 分布式锁竞争主实例角色： +// - 获得锁的实例负责执行定时刷新，并周期续期 +// - 未获得锁的实例定期重试竞争（当主实例宕机时可接管） func (w *HotRankWorker) Start(ctx context.Context) { if w == nil || w.cfg == nil || !w.cfg.HotRank.Enabled || w.c == nil { close(w.doneCh) @@ -59,25 +114,63 @@ func (w *HotRankWorker) Start(ctx context.Context) { go func() { defer close(w.doneCh) - if err := w.Refresh(ctx); err != nil { - if !errors.Is(err, cache.ErrKeyNotFound) { - zap.L().Warn("hot rank initial refresh failed", zap.Error(err)) - } - } - - ticker := time.NewTicker(interval) - defer ticker.Stop() + isLeader := false for { select { case <-w.stopCh: + if isLeader { + w.releaseLock() + } return - case <-ticker.C: - if err := w.Refresh(ctx); err != nil { - if !errors.Is(err, cache.ErrKeyNotFound) { - zap.L().Warn("hot rank refresh failed", zap.Error(err)) + default: + } + + acquired := w.tryAcquireLock(ctx) + if acquired { + if !isLeader { + zap.L().Info("hot rank: acquired leader lock, becoming primary") + isLeader = true + // 启动续期协程 + lockCtx, lockCancel := context.WithCancel(ctx) + go w.keepLockAlive(lockCtx) + + // 主实例立即执行一次刷新 + if err := w.Refresh(ctx); err != nil { + if !errors.Is(err, cache.ErrKeyNotFound) { + zap.L().Warn("hot rank initial refresh failed", zap.Error(err)) + } + } + + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-w.stopCh: + lockCancel() + w.releaseLock() + return + case <-ticker.C: + if err := w.Refresh(ctx); err != nil { + if !errors.Is(err, cache.ErrKeyNotFound) { + zap.L().Warn("hot rank refresh failed", zap.Error(err)) + } + } + } } } + // 不可能走到这里（isLeader==true 的分支里有无限循环） + } else { + // 未获得锁，作为备用实例等待重试 + zap.L().Debug("hot rank: another instance holds leader lock, retrying later") + } + + select { + case <-w.stopCh: + return + case <-time.After(hotRankLockRetryInterval): + // 重试竞争锁 } } }() @@ -92,6 +185,7 @@ func (w *HotRankWorker) Stop() { close(w.stopCh) }) <-w.doneCh + w.releaseLock() } // Refresh 立即重算热门分（线程安全） @@ -240,15 +334,15 @@ func (w *HotRankWorker) Refresh(ctx context.Context) error { } const pinnedScoreBase = 1000.0 - zMembers := make([]redis.Z, 0, len(orderedIDs)) + zMembers := make([]redislib.Z, 0, len(orderedIDs)) for i, id := range pinned { - zMembers = append(zMembers, redis.Z{ + zMembers = append(zMembers, redislib.Z{ Score: pinnedScoreBase - float64(i)*1e-6, Member: id, }) } for _, h := range hotPicks { - zMembers = append(zMembers, redis.Z{Score: h.norm * 0.99, Member: h.id}) + zMembers = append(zMembers, redislib.Z{Score: h.norm * 0.99, Member: h.id}) } if err := w.c.ZReplaceSortedSet(ctx, zsetKey, zMembers); err != nil { @@ -448,15 +542,15 @@ func (w *HotRankWorker) refreshChannelRanks( } const pinnedScoreBase = 1000.0 - chZMembers := make([]redis.Z, 0, len(chOrderedIDs)) + chZMembers := make([]redislib.Z, 0, len(chOrderedIDs)) for i, id := range chPinned { - chZMembers = append(chZMembers, redis.Z{ + chZMembers = append(chZMembers, redislib.Z{ Score: pinnedScoreBase - float64(i)*1e-6, Member: id, }) } for _, h := range hotPicks { - chZMembers = append(chZMembers, redis.Z{Score: h.norm * 0.99, Member: h.id}) + chZMembers = append(chZMembers, redislib.Z{Score: h.norm * 0.99, Member: h.id}) } chZSetKey := cache.HotRankZSetChannelKey(ch.ID) diff --git a/internal/wire/grpc.go b/internal/wire/grpc.go index 1e4e9b1..0f0d14a 100644 --- a/internal/wire/grpc.go +++ b/internal/wire/grpc.go @@ -1,25 +1,108 @@ package wire import ( + "time" + "with_you/internal/config" "with_you/internal/grpc/runner" + "with_you/internal/pkg/redis" "github.com/google/wire" + "github.com/google/uuid" "go.uber.org/zap" ) // GRPCSet gRPC 相关的 Provider Set var GRPCSet = wire.NewSet( ProvideGRPCServer, + ProvideRunnerHub, ProvideTaskManager, + ProvideRunnerRegistry, + ProvideTaskDispatcher, ) // ProvideGRPCServer 提供 gRPC 服务器 -func ProvideGRPCServer(cfg *config.Config, logger *zap.Logger) *runner.Server { - return runner.NewServer(&cfg.GRPC, logger) +func ProvideGRPCServer(cfg *config.Config, logger *zap.Logger, hub *runner.RunnerHub) *runner.Server { + return runner.NewServerWithDeps(&cfg.GRPC, logger, hub) +} + +// ProvideRunnerHub 创建 RunnerHub +func ProvideRunnerHub(logger *zap.Logger) *runner.RunnerHub { + return runner.NewRunnerHub(nil, logger) } // ProvideTaskManager 提供 TaskManager -func ProvideTaskManager(server *runner.Server) *runner.TaskManager { - return server.GetTaskManager() +func ProvideTaskManager(hub *runner.RunnerHub, logger *zap.Logger) *runner.TaskManager { + tm := runner.NewTaskManager(nil, 60*time.Second, logger) + tm.SetDispatcher(hub) + return tm +} + +// ProvideRunnerRegistry 提供 Runner 注册中心 +// standalone 模式或 Redis 不可用时返回 nil +func ProvideRunnerRegistry(cfg *config.Config, redisClient *redis.Client) *runner.RunnerRegistry { + if cfg.Runner.Mode != "cluster" || redisClient == nil { + return nil + } + + instanceID := cfg.Runner.Cluster.InstanceID + if instanceID == "" { + instanceID = uuid.New().String() + } + + rdb := redisClient.GetClient() + + return runner.NewRunnerRegistry( + rdb, + instanceID, + cfg.Runner.Cluster.RegistryPrefix, + cfg.Runner.Cluster.CapPrefix, + time.Duration(cfg.Runner.Cluster.RegistryTTL)*time.Second, + time.Duration(cfg.Runner.Cluster.HeartbeatInterval)*time.Second, + zap.L(), + ) +} + +// ProvideTaskDispatcher 提供任务分发器 +// standalone 模式返回 *RunnerHub，cluster 模式返回 *TaskBus +func ProvideTaskDispatcher( + hub *runner.RunnerHub, + registry *runner.RunnerRegistry, + cfg *config.Config, + redisClient *redis.Client, + logger *zap.Logger, +) runner.TaskDispatcher { + if cfg.Runner.Mode != "cluster" || redisClient == nil || registry == nil { + logger.Info("Runner running in standalone mode") + return hub + } + + instanceID := cfg.Runner.Cluster.InstanceID + if instanceID == "" { + instanceID = uuid.New().String() + } + + rdb := redisClient.GetClient() + + bus := runner.NewTaskBus( + hub, + rdb, + registry, + instanceID, + cfg.Runner.Cluster.TaskChannel, + cfg.Runner.Cluster.ResultChannel, + logger, + ) + if err := bus.Start(); err != nil { + logger.Error("Failed to start Runner Bus, falling back to standalone", zap.Error(err)) + return hub + } + + // 注入 Bus 和 Registry 到 Hub + hub.SetBus(bus) + hub.SetRegistry(registry) + registry.StartHeartbeat() + + logger.Info("Runner running in cluster mode", zap.String("instance_id", bus.InstanceID())) + return bus } diff --git a/internal/wire/service.go b/internal/wire/service.go index 8779e7b..9e17b7b 100644 --- a/internal/wire/service.go +++ b/internal/wire/service.go @@ -18,6 +18,7 @@ import ( "with_you/internal/service" "github.com/google/wire" + redislib "github.com/redis/go-redis/v9" "go.uber.org/zap" "gorm.io/gorm" ) @@ -376,8 +377,12 @@ func ProvideLogService( } // ProvideQRCodeLoginService 提供二维码登录服务 -func ProvideHotRankWorker(cfg *config.Config, postRepo repository.PostRepository, channelRepo repository.ChannelRepository, cacheBackend cache.Cache) *service.HotRankWorker { - return service.NewHotRankWorker(cfg, postRepo, channelRepo, cacheBackend) +func ProvideHotRankWorker(cfg *config.Config, postRepo repository.PostRepository, channelRepo repository.ChannelRepository, cacheBackend cache.Cache, redisClient *redis.Client) *service.HotRankWorker { + var rdb *redislib.Client + if redisClient != nil { + rdb = redisClient.GetClient() + } + return service.NewHotRankWorker(cfg, postRepo, channelRepo, cacheBackend, rdb) } func ProvideQRCodeLoginService(