浅浅升级一下

2026-05-06 14:41:14 +08:00
parent b0bd37ec12
commit 3a2e3f2571
16 changed files with 1490 additions and 106 deletions
--- a/cmd/worker/main.go
+++ b/cmd/worker/main.go
@@ -2,6 +2,7 @@ package main

 import (
 	"context"
+	"fmt"
 	"os"
 	"os/signal"
 	"strconv"
@@ -25,14 +26,50 @@ import (
 	"github.com/break/junhong_cmp_fiber/pkg/logger"
 	"github.com/break/junhong_cmp_fiber/pkg/queue"
 	"github.com/break/junhong_cmp_fiber/pkg/storage"
+	"gorm.io/gorm"
 )

+const (
+	workerModuleQueueServer        = "queue_server"
+	workerModulePollingInitializer = "polling_initializer"
+	workerModulePollingScheduler   = "polling_scheduler"
+	workerModuleAsynqScheduler     = "asynq_scheduler"
+)
+
+// workerModuleStatus 描述当前 Worker 角色下的模块启停计划。
+type workerModuleStatus struct {
+	enabled  []string
+	disabled []string
+}
+
+// workerRuntime 保存所有角色共享的 Worker 运行时依赖。
+type workerRuntime struct {
+	redisAddr           string
+	redisClient         *redis.Client
+	db                  *gorm.DB
+	storageSvc          *storage.Service
+	gatewayClient       *gateway.Client
+	asynqClient         *asynq.Client
+	workerResult        *bootstrap.WorkerBootstrapResult
+	workerServer        *queue.Server
+	pollingConfigMgr    *polling.PollingConfigManager
+	pollingQueueMgr     *polling.PollingQueueManager
+	pollingIotCardStore *postgres.IotCardStore
+	pollingBase         *task.PollingBase
+	lifecycleSvc        *polling.PollingLifecycleService
+}
+
 func main() {
 	cfg, err := config.Load()
 	if err != nil {
 		panic("加载配置失败: " + err.Error())
 	}

+	runWorker(cfg)
+}
+
+// runWorker 编排 Worker 进程启动、模块启停和优雅关闭流程。
+func runWorker(cfg *config.Config) {
 	if _, err := pkgBootstrap.EnsureDirectories(cfg, nil); err != nil {
 		panic("初始化目录失败: " + err.Error())
 	}
@@ -62,9 +99,90 @@ func main() {
 	}()

 	appLogger := logger.GetAppLogger()
-	appLogger.Info("Worker 服务启动中...")
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()

-	// 连接 Redis
+	moduleStatus := buildWorkerModuleStatus(cfg.Worker.Role)
+	appLogger.Info("Worker 服务启动中...")
+	logWorkerRole(appLogger, cfg.Worker.Role, cfg.Worker.InstanceName, moduleStatus)
+
+	runtime := initWorkerRuntime(ctx, cfg, appLogger)
+	defer runtime.close(appLogger)
+
+	var pollingInitializer *polling.PollingInitializer
+	var pollingScheduler *polling.Scheduler
+	var asynqScheduler *asynq.Scheduler
+
+	if runsSingletonModules(cfg.Worker.Role) {
+		pollingInitializer = startPollingInitializer(ctx, runtime, appLogger)
+		pollingScheduler = startPollingScheduler(ctx, runtime, pollingInitializer, appLogger)
+		asynqScheduler = startAsynqScheduler(cfg, runtime.redisAddr, appLogger)
+	}
+
+	taskHandler := createTaskHandler(runtime, appLogger)
+	taskHandler.RegisterHandlers()
+
+	appLogger.Info("Worker 服务器配置完成",
+		zap.Int("concurrency", cfg.Queue.Concurrency),
+		zap.Any("queues", cfg.Queue.Queues))
+
+	quit := make(chan os.Signal, 1)
+	signal.Notify(quit, os.Interrupt, syscall.SIGTERM)
+	defer signal.Stop(quit)
+
+	go func() {
+		if err := runtime.workerServer.Run(taskHandler.GetMux()); err != nil {
+			appLogger.Fatal("Worker 服务器运行失败", zap.Error(err))
+		}
+	}()
+
+	appLogger.Info("Worker 服务器已启动")
+	<-quit
+
+	shutdownWorker(cancel, appLogger, runtime.workerServer, pollingInitializer, pollingScheduler, asynqScheduler)
+}
+
+// buildWorkerModuleStatus 根据角色生成启用/禁用模块列表，供启动日志和人工验收使用。
+func buildWorkerModuleStatus(role string) workerModuleStatus {
+	status := workerModuleStatus{
+		enabled: []string{workerModuleQueueServer},
+	}
+
+	if runsSingletonModules(role) {
+		status.enabled = append(
+			status.enabled,
+			workerModulePollingInitializer,
+			workerModulePollingScheduler,
+			workerModuleAsynqScheduler,
+		)
+		return status
+	}
+
+	status.disabled = []string{
+		workerModulePollingInitializer,
+		workerModulePollingScheduler,
+		workerModuleAsynqScheduler,
+	}
+	return status
+}
+
+// logWorkerRole 输出当前实例角色、实例名以及模块启停计划。
+func logWorkerRole(appLogger *zap.Logger, role string, instanceName string, moduleStatus workerModuleStatus) {
+	appLogger.Info("Worker 角色已确认",
+		zap.String("role", role),
+		zap.String("instance_name", instanceName))
+	appLogger.Info("Worker 模块启停计划",
+		zap.Strings("enabled_modules", moduleStatus.enabled),
+		zap.Strings("disabled_modules", moduleStatus.disabled))
+}
+
+// runsSingletonModules 判断当前角色是否承担主动调度和初始化职责。
+func runsSingletonModules(role string) bool {
+	return role == constants.WorkerRoleAll || role == constants.WorkerRoleLeader
+}
+
+// initWorkerRuntime 初始化所有角色共享的外部依赖和轮询运行时对象。
+func initWorkerRuntime(ctx context.Context, cfg *config.Config, appLogger *zap.Logger) *workerRuntime {
 	redisAddr := cfg.Redis.Address + ":" + strconv.Itoa(cfg.Redis.Port)
 	redisClient := redis.NewClient(&redis.Options{
 		Addr:         redisAddr,
@@ -76,52 +194,26 @@ func main() {
 		ReadTimeout:  cfg.Redis.ReadTimeout,
 		WriteTimeout: cfg.Redis.WriteTimeout,
 	})
-	defer func() {
-		if err := redisClient.Close(); err != nil {
-			appLogger.Error("关闭 Redis 客户端失败", zap.Error(err))
-		}
-	}()

-	// 测试 Redis 连接
-	ctx := context.Background()
 	if err := redisClient.Ping(ctx).Err(); err != nil {
 		appLogger.Fatal("连接 Redis 失败", zap.Error(err))
 	}
 	appLogger.Info("Redis 已连接", zap.String("address", redisAddr))

-	// 初始化 PostgreSQL 连接
 	db, err := database.InitPostgreSQL(&cfg.Database, appLogger)
 	if err != nil {
 		appLogger.Fatal("初始化 PostgreSQL 失败", zap.Error(err))
 	}
-	defer func() {
-		sqlDB, _ := db.DB()
-		if sqlDB != nil {
-			if err := sqlDB.Close(); err != nil {
-				appLogger.Error("关闭 PostgreSQL 连接失败", zap.Error(err))
-			}
-		}
-	}()

-	// 初始化对象存储服务（可选）
 	storageSvc := initStorage(cfg, appLogger)
-
-	// 初始化 Gateway 客户端（可选，用于轮询任务）
 	gatewayClient := initGateway(cfg, appLogger)

-	// 创建 Asynq 客户端（用于调度器提交任务）
 	asynqClient := asynq.NewClient(asynq.RedisClientOpt{
 		Addr:     redisAddr,
 		Password: cfg.Redis.Password,
 		DB:       cfg.Redis.DB,
 	})
-	defer func() {
-		if err := asynqClient.Close(); err != nil {
-			appLogger.Error("关闭 Asynq 客户端失败", zap.Error(err))
-		}
-	}()

-	// 创建 Worker 依赖
 	workerDeps := &bootstrap.WorkerDependencies{
 		DB:             db,
 		Redis:          redisClient,
@@ -131,16 +223,13 @@ func main() {
 		GatewayClient:  gatewayClient,
 	}

-	// Bootstrap Worker 组件
 	workerResult, err := bootstrap.BootstrapWorker(workerDeps)
 	if err != nil {
 		appLogger.Fatal("Worker Bootstrap 失败", zap.Error(err))
 	}

-	// 创建 Asynq Worker 服务器
 	workerServer := queue.NewServer(redisClient, &cfg.Queue, appLogger)

-	// 初始化轮询系统新组件（Phase 5.7）
 	pollingConfigStore := postgres.NewPollingConfigStore(db)
 	pollingConfigMgr := polling.NewPollingConfigManager(pollingConfigStore, redisClient, appLogger)
 	if err := pollingConfigMgr.Load(ctx); err != nil {
@@ -149,55 +238,125 @@ func main() {
 	pollingConfigMgr.Start(ctx)

 	pollingQueueMgr := polling.NewPollingQueueManager(redisClient, constants.PollingShardCount, appLogger)
-
 	pollingIotCardStore := postgres.NewIotCardStore(db, redisClient)
-	pollingBase := task.NewPollingBase(redisClient, pollingQueueMgr, pollingConfigMgr, pollingIotCardStore, appLogger, cfg.Polling.VerboseLog)
+	pollingBase := task.NewPollingBase(
+		redisClient,
+		pollingQueueMgr,
+		pollingConfigMgr,
+		pollingIotCardStore,
+		appLogger,
+		cfg.Polling.VerboseLog,
+	)

-	// 后台渐进式初始化（将全量卡数据写入分片 Sorted Set）
-	pollingInitializer := polling.NewPollingInitializer(pollingIotCardStore, redisClient, pollingConfigMgr, pollingQueueMgr, appLogger)
+	pollingDeviceSimBindingStore := postgres.NewDeviceSimBindingStore(db, redisClient)
+	pollingDeviceStore := postgres.NewDeviceStore(db, redisClient)
+	lifecycleSvc := polling.NewPollingLifecycleService(
+		pollingQueueMgr,
+		pollingConfigMgr,
+		pollingIotCardStore,
+		pollingDeviceSimBindingStore,
+		pollingDeviceStore,
+		appLogger,
+	)
+
+	return &workerRuntime{
+		redisAddr:           redisAddr,
+		redisClient:         redisClient,
+		db:                  db,
+		storageSvc:          storageSvc,
+		gatewayClient:       gatewayClient,
+		asynqClient:         asynqClient,
+		workerResult:        workerResult,
+		workerServer:        workerServer,
+		pollingConfigMgr:    pollingConfigMgr,
+		pollingQueueMgr:     pollingQueueMgr,
+		pollingIotCardStore: pollingIotCardStore,
+		pollingBase:         pollingBase,
+		lifecycleSvc:        lifecycleSvc,
+	}
+}
+
+// close 在 Worker 退出时关闭共享客户端与数据库连接。
+func (r *workerRuntime) close(appLogger *zap.Logger) {
+	if r.asynqClient != nil {
+		if err := r.asynqClient.Close(); err != nil {
+			appLogger.Error("关闭 Asynq 客户端失败", zap.Error(err))
+		}
+	}
+	if r.db != nil {
+		sqlDB, _ := r.db.DB()
+		if sqlDB != nil {
+			if err := sqlDB.Close(); err != nil {
+				appLogger.Error("关闭 PostgreSQL 连接失败", zap.Error(err))
+			}
+		}
+	}
+	if r.redisClient != nil {
+		if err := r.redisClient.Close(); err != nil {
+			appLogger.Error("关闭 Redis 客户端失败", zap.Error(err))
+		}
+	}
+}
+
+// startPollingInitializer 为 leader/all 启动渐进式初始化器，并保留配置变更重启逻辑。
+func startPollingInitializer(ctx context.Context, runtime *workerRuntime, appLogger *zap.Logger) *polling.PollingInitializer {
+	pollingInitializer := polling.NewPollingInitializer(
+		runtime.pollingIotCardStore,
+		runtime.redisClient,
+		runtime.pollingConfigMgr,
+		runtime.pollingQueueMgr,
+		appLogger,
+	)
 	pollingInitializer.StartBackground(ctx)

-	// 订阅配置变更事件：配置从空变为非空时重启 Initializer，将卡重新入队
-	pollingConfigMgr.WatchChanges(ctx, func(hadConfigs, hasConfigs bool) {
+	runtime.pollingConfigMgr.WatchChanges(ctx, func(hadConfigs, hasConfigs bool) {
 		if !hadConfigs && hasConfigs {
 			appLogger.Info("轮询配置从空变为非空，触发队列重新初始化")
 			pollingInitializer.Restart(ctx)
 		}
 	})

-	// 创建生命周期服务（Worker 进程用，功能更完整）
-	pollingDeviceSimBindingStore := postgres.NewDeviceSimBindingStore(db, redisClient)
-	pollingDeviceStore := postgres.NewDeviceStore(db, redisClient)
-	lifecycleSvc := polling.NewPollingLifecycleService(pollingQueueMgr, pollingConfigMgr, pollingIotCardStore, pollingDeviceSimBindingStore, pollingDeviceStore, appLogger)
+	return pollingInitializer
+}

-	// 初始化调度器（激活/重置 Handler 通过构造函数注入，防止遗漏）
-	dataResetHandler := polling.NewDataResetHandler(workerResult.Services.ResetService, appLogger)
+// startPollingScheduler 为 leader/all 启动轮询调度器。
+func startPollingScheduler(
+	ctx context.Context,
+	runtime *workerRuntime,
+	pollingInitializer *polling.PollingInitializer,
+	appLogger *zap.Logger,
+) *polling.Scheduler {
+	dataResetHandler := polling.NewDataResetHandler(runtime.workerResult.Services.ResetService, appLogger)
 	activationHandler := polling.NewPackageActivationHandler(
-		db, redisClient, asynqClient,
-		workerResult.Services.ActivationService,
-		workerResult.Services.StopResumeService,
+		runtime.db,
+		runtime.redisClient,
+		runtime.asynqClient,
+		runtime.workerResult.Services.ActivationService,
+		runtime.workerResult.Services.StopResumeService,
 		appLogger,
 	)
-	scheduler := polling.NewScheduler(redisClient, asynqClient, pollingQueueMgr, pollingConfigMgr, appLogger, activationHandler, dataResetHandler)
-	scheduler.SetInitializer(pollingInitializer)

-	if err := scheduler.Start(ctx); err != nil {
+	pollingScheduler := polling.NewScheduler(
+		runtime.redisClient,
+		runtime.asynqClient,
+		runtime.pollingQueueMgr,
+		runtime.pollingConfigMgr,
+		appLogger,
+		activationHandler,
+		dataResetHandler,
+	)
+	pollingScheduler.SetInitializer(pollingInitializer)
+
+	if err := pollingScheduler.Start(ctx); err != nil {
 		appLogger.Error("启动轮询调度器失败", zap.Error(err))
-	} else {
-		appLogger.Info("轮询调度器已启动")
+		return nil
 	}

-	// 类型断言：StopResumeService 实际是 *iotCardSvc.StopResumeService，实现了 EvaluateAndAct 接口
-	stopResumeSvc, _ := workerResult.Services.StopResumeService.(iot_card_svc.StopResumeServiceInterface)
-	// 创建任务处理器并注册
-	taskHandler := queue.NewHandler(db, redisClient, storageSvc, gatewayClient, lifecycleSvc, workerResult, asynqClient, appLogger, pollingBase, stopResumeSvc)
-	taskHandler.RegisterHandlers()
+	return pollingScheduler
+}

-	appLogger.Info("Worker 服务器配置完成",
-		zap.Int("concurrency", cfg.Queue.Concurrency),
-		zap.Any("queues", cfg.Queue.Queues))
-
-	// 创建 Asynq Scheduler（定时任务调度器：订单超时、告警检查、数据清理）
+// startAsynqScheduler 为 leader/all 创建并启动 Asynq Scheduler。
+func startAsynqScheduler(cfg *config.Config, redisAddr string, appLogger *zap.Logger) *asynq.Scheduler {
 	asynqScheduler := asynq.NewScheduler(
 		asynq.RedisClientOpt{
 			Addr:     redisAddr,
@@ -207,57 +366,82 @@ func main() {
 		&asynq.SchedulerOpts{Location: time.Local},
 	)

-	// 注册定时任务：订单超时检查（每分钟）
-	if _, err := asynqScheduler.Register("@every 1m", asynq.NewTask(constants.TaskTypeOrderExpire, nil)); err != nil {
-		appLogger.Fatal("注册订单超时定时任务失败", zap.Error(err))
-	}
-	// 注册定时任务：告警检查（每分钟）
-	if _, err := asynqScheduler.Register("@every 1m", asynq.NewTask(constants.TaskTypeAlertCheck, nil)); err != nil {
-		appLogger.Fatal("注册告警检查定时任务失败", zap.Error(err))
-	}
-	// 注册定时任务：数据清理（每天凌晨 2 点）
-	if _, err := asynqScheduler.Register("0 2 * * *", asynq.NewTask(constants.TaskTypeDataCleanup, nil)); err != nil {
-		appLogger.Fatal("注册数据清理定时任务失败", zap.Error(err))
-	}
-	// 注册定时任务：每日流量落盘（每天凌晨 2 点，超时 5 分钟）
-	if _, err := asynqScheduler.Register("0 2 * * *", asynq.NewTask(constants.TaskTypeDailyTrafficFlush, nil, asynq.MaxRetry(3), asynq.Timeout(5*time.Minute))); err != nil {
-		appLogger.Fatal("注册每日流量落盘定时任务失败", zap.Error(err))
+	if err := registerAsynqScheduleTasks(asynqScheduler); err != nil {
+		appLogger.Fatal("注册 Asynq 定时任务失败", zap.Error(err))
 	}

-	// 启动 Asynq Scheduler
 	go func() {
 		if err := asynqScheduler.Run(); err != nil {
 			appLogger.Fatal("Asynq Scheduler 启动失败", zap.Error(err))
 		}
 	}()
-	appLogger.Info("Asynq Scheduler 已启动（订单超时: @every 1m, 告警检查: @every 1m, 数据清理: 0 2 * * *）")

-	// 优雅关闭
-	quit := make(chan os.Signal, 1)
-	signal.Notify(quit, os.Interrupt, syscall.SIGTERM)
+	appLogger.Info("Asynq Scheduler 已启动（订单超时: @every 1m, 告警检查: @every 1m, 数据清理: 0 2 * * *, 每日流量落盘: 0 2 * * *）")
+	return asynqScheduler
+}

-	// 启动 Worker 服务器（阻塞运行）
-	go func() {
-		if err := workerServer.Run(taskHandler.GetMux()); err != nil {
-			appLogger.Fatal("Worker 服务器运行失败", zap.Error(err))
-		}
-	}()
+// registerAsynqScheduleTasks 注册 Worker 入口需要的全部定时任务。
+func registerAsynqScheduleTasks(asynqScheduler *asynq.Scheduler) error {
+	if _, err := asynqScheduler.Register("@every 1m", asynq.NewTask(constants.TaskTypeOrderExpire, nil)); err != nil {
+		return fmt.Errorf("注册订单超时定时任务失败: %w", err)
+	}
+	if _, err := asynqScheduler.Register("@every 1m", asynq.NewTask(constants.TaskTypeAlertCheck, nil)); err != nil {
+		return fmt.Errorf("注册告警检查定时任务失败: %w", err)
+	}
+	if _, err := asynqScheduler.Register("0 2 * * *", asynq.NewTask(constants.TaskTypeDataCleanup, nil)); err != nil {
+		return fmt.Errorf("注册数据清理定时任务失败: %w", err)
+	}
+	if _, err := asynqScheduler.Register(
+		"0 2 * * *",
+		asynq.NewTask(constants.TaskTypeDailyTrafficFlush, nil, asynq.MaxRetry(3), asynq.Timeout(5*time.Minute)),
+	); err != nil {
+		return fmt.Errorf("注册每日流量落盘定时任务失败: %w", err)
+	}
+	return nil
+}

-	appLogger.Info("Worker 服务器已启动")
+// createTaskHandler 创建并返回包含全部任务处理器的 Asynq Handler。
+func createTaskHandler(runtime *workerRuntime, appLogger *zap.Logger) *queue.Handler {
+	stopResumeSvc, _ := runtime.workerResult.Services.StopResumeService.(iot_card_svc.StopResumeServiceInterface)
+	return queue.NewHandler(
+		runtime.db,
+		runtime.redisClient,
+		runtime.storageSvc,
+		runtime.gatewayClient,
+		runtime.lifecycleSvc,
+		runtime.workerResult,
+		runtime.asynqClient,
+		appLogger,
+		runtime.pollingBase,
+		stopResumeSvc,
+	)
+}

-	// 等待关闭信号
-	<-quit
+// shutdownWorker 按当前实例实际启动过的模块执行优雅关闭。
+func shutdownWorker(
+	cancel context.CancelFunc,
+	appLogger *zap.Logger,
+	workerServer *queue.Server,
+	pollingInitializer *polling.PollingInitializer,
+	pollingScheduler *polling.Scheduler,
+	asynqScheduler *asynq.Scheduler,
+) {
 	appLogger.Info("正在关闭 Worker 服务器...")

-	// 停止 Asynq Scheduler
-	asynqScheduler.Shutdown()
+	if asynqScheduler != nil {
+		asynqScheduler.Shutdown()
+	}
+	if pollingScheduler != nil {
+		pollingScheduler.Stop()
+	}

-	// 停止轮询调度器
-	scheduler.Stop()
+	cancel()
+
+	if pollingInitializer != nil {
+		pollingInitializer.Stop()
+	}

-	// 优雅关闭 Worker 服务器（等待正在执行的任务完成）
 	workerServer.Shutdown()
-
 	appLogger.Info("Worker 服务器已停止")
 }