From 0d96a94e5f3b49eae0faa91060e11d889641a19d Mon Sep 17 00:00:00 2001 From: Break Date: Wed, 27 May 2026 10:59:19 +0800 Subject: [PATCH] =?UTF-8?q?=E5=85=B3=E4=BA=8E=E4=B8=8A=E6=B8=B8=E6=B5=81?= =?UTF-8?q?=E9=87=8F=E5=90=8C=E6=AD=A5=E8=A6=86=E7=9B=96=E4=BF=AE=E5=A4=8D?= =?UTF-8?q?,=E4=BB=A5=E5=8F=8A=E6=96=B0=E5=A2=9Eredis=E5=90=8C=E6=AD=A5?= =?UTF-8?q?=E8=BF=81=E7=A7=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cmd/migration-finalize/main.go | 293 ++++++++++++++++++++++ internal/polling/config_manager.go | 18 +- internal/polling/initializer.go | 6 + internal/polling/queue_manager.go | 20 ++ internal/service/iot_card/service.go | 7 +- internal/task/polling_carddata_handler.go | 11 +- 6 files changed, 349 insertions(+), 6 deletions(-) create mode 100644 cmd/migration-finalize/main.go diff --git a/cmd/migration-finalize/main.go b/cmd/migration-finalize/main.go new file mode 100644 index 0000000..d5e3271 --- /dev/null +++ b/cmd/migration-finalize/main.go @@ -0,0 +1,293 @@ +package main + +import ( + "context" + "flag" + "fmt" + "os" + "sort" + "strconv" + "strings" + "time" + + "github.com/redis/go-redis/v9" + "go.uber.org/zap" + "gorm.io/gorm" + + "github.com/break/junhong_cmp_fiber/internal/model" + "github.com/break/junhong_cmp_fiber/internal/polling" + "github.com/break/junhong_cmp_fiber/internal/store/postgres" + pkgbootstrap "github.com/break/junhong_cmp_fiber/pkg/bootstrap" + "github.com/break/junhong_cmp_fiber/pkg/config" + "github.com/break/junhong_cmp_fiber/pkg/constants" + "github.com/break/junhong_cmp_fiber/pkg/database" + "github.com/break/junhong_cmp_fiber/pkg/logger" +) + +type finalizeOptions struct { + batchNo string + apply bool + batchSize int + timeout time.Duration +} + +type finalizeStats struct { + importedTotal int64 + pollingCards int + enqueuedCards int + skippedNoConfig int + firstID uint + lastID uint + taskCounts map[string]int +} + +func main() { + if err := run(); err != nil { + fmt.Fprintf(os.Stderr, "迁移收尾失败: %v\n", err) + os.Exit(1) + } +} + +func run() error { + opts := parseOptions() + if strings.TrimSpace(opts.batchNo) == "" { + return fmt.Errorf("batch-no 不能为空") + } + if opts.batchSize <= 0 { + return fmt.Errorf("batch-size 必须大于 0") + } + + cfg, err := config.Load() + if err != nil { + return fmt.Errorf("加载配置失败: %w", err) + } + if _, err := pkgbootstrap.EnsureDirectories(cfg, nil); err != nil { + return fmt.Errorf("初始化目录失败: %w", err) + } + if err := initLogger(cfg); err != nil { + return err + } + defer func() { _ = logger.Sync() }() + + appLogger := logger.GetAppLogger() + ctx, cancel := context.WithTimeout(context.Background(), opts.timeout) + defer cancel() + + db, err := database.InitPostgreSQL(&cfg.Database, appLogger) + if err != nil { + return fmt.Errorf("初始化 PostgreSQL 失败: %w", err) + } + defer closeDB(db, appLogger) + + var redisClient *redis.Client + if opts.apply { + redisClient, err = initRedis(cfg, appLogger) + if err != nil { + return err + } + defer func() { + if closeErr := redisClient.Close(); closeErr != nil { + appLogger.Error("关闭 Redis 连接失败", zap.Error(closeErr)) + } + }() + } + + configStore := postgres.NewPollingConfigStore(db) + configMgr := polling.NewPollingConfigManager(configStore, redisClient, appLogger) + if opts.apply { + if err := configMgr.Load(ctx); err != nil { + return fmt.Errorf("加载轮询配置失败: %w", err) + } + } else { + if err := configMgr.LoadForInspection(ctx); err != nil { + return fmt.Errorf("加载轮询配置失败: %w", err) + } + } + + var queueMgr *polling.PollingQueueManager + var initializer *polling.PollingInitializer + if opts.apply { + queueMgr = polling.NewPollingQueueManager(redisClient, constants.PollingShardCount, appLogger) + cardStore := postgres.NewIotCardStore(db, redisClient) + initializer = polling.NewPollingInitializer(cardStore, redisClient, configMgr, queueMgr, appLogger) + } + + stats, err := finalizePolling(ctx, db, configMgr, queueMgr, initializer, opts) + if err != nil { + return err + } + printStats(stats, opts) + return nil +} + +func parseOptions() finalizeOptions { + opts := finalizeOptions{} + flag.StringVar(&opts.batchNo, "batch-no", "MIGRATION-QICHENG", "迁移批次号") + flag.BoolVar(&opts.apply, "apply", false, "实际写入 Redis 轮询队列和卡缓存；默认只 dry-run") + flag.IntVar(&opts.batchSize, "batch-size", 500, "每批处理卡数量") + flag.DurationVar(&opts.timeout, "timeout", 10*time.Minute, "迁移收尾超时时间") + flag.Parse() + return opts +} + +func initLogger(cfg *config.Config) error { + if err := logger.InitLoggers( + cfg.Logging.Level, + cfg.Logging.Development, + logger.LogRotationConfig{ + Filename: cfg.Logging.AppLog.Filename, + MaxSize: cfg.Logging.AppLog.MaxSize, + MaxBackups: cfg.Logging.AppLog.MaxBackups, + MaxAge: cfg.Logging.AppLog.MaxAge, + Compress: cfg.Logging.AppLog.Compress, + }, + logger.LogRotationConfig{ + Filename: cfg.Logging.AccessLog.Filename, + MaxSize: cfg.Logging.AccessLog.MaxSize, + MaxBackups: cfg.Logging.AccessLog.MaxBackups, + MaxAge: cfg.Logging.AccessLog.MaxAge, + Compress: cfg.Logging.AccessLog.Compress, + }, + ); err != nil { + return fmt.Errorf("初始化日志失败: %w", err) + } + return nil +} + +func initRedis(cfg *config.Config, appLogger *zap.Logger) (*redis.Client, error) { + redisAddr := cfg.Redis.Address + ":" + strconv.Itoa(cfg.Redis.Port) + client, err := database.NewRedisClient(database.RedisConfig{ + Address: redisAddr, + Password: cfg.Redis.Password, + DB: cfg.Redis.DB, + PoolSize: cfg.Redis.PoolSize, + MinIdleConns: cfg.Redis.MinIdleConns, + DialTimeout: cfg.Redis.DialTimeout, + ReadTimeout: cfg.Redis.ReadTimeout, + WriteTimeout: cfg.Redis.WriteTimeout, + }, appLogger) + if err != nil { + return nil, fmt.Errorf("初始化 Redis 失败: %w", err) + } + return client, nil +} + +func closeDB(db *gorm.DB, appLogger *zap.Logger) { + sqlDB, _ := db.DB() + if sqlDB == nil { + return + } + if err := sqlDB.Close(); err != nil { + appLogger.Error("关闭 PostgreSQL 连接失败", zap.Error(err)) + } +} + +func finalizePolling( + ctx context.Context, + db *gorm.DB, + configMgr *polling.PollingConfigManager, + queueMgr *polling.PollingQueueManager, + initializer *polling.PollingInitializer, + opts finalizeOptions, +) (finalizeStats, error) { + stats := finalizeStats{taskCounts: map[string]int{}} + if err := db.WithContext(ctx).Model(&model.IotCard{}). + Where("batch_no = ?", opts.batchNo). + Count(&stats.importedTotal).Error; err != nil { + return stats, fmt.Errorf("统计迁移卡失败: %w", err) + } + + var lastID uint + for { + cards, err := loadPollingCards(ctx, db, opts.batchNo, lastID, opts.batchSize) + if err != nil { + return stats, err + } + if len(cards) == 0 { + break + } + + updateStats(&stats, configMgr, cards) + if opts.apply { + cardIDs := collectCardIDs(cards) + if err := queueMgr.RemoveBatchFromCurrentShardQueues(ctx, cardIDs); err != nil { + return stats, fmt.Errorf("清理轮询旧队列失败: %w", err) + } + if err := initializer.InitCards(ctx, cards); err != nil { + return stats, fmt.Errorf("重建轮询队列失败: %w", err) + } + } + + lastID = cards[len(cards)-1].ID + } + + return stats, nil +} + +func loadPollingCards(ctx context.Context, db *gorm.DB, batchNo string, lastID uint, limit int) ([]*model.IotCard, error) { + var cards []*model.IotCard + err := db.WithContext(ctx). + Where("id > ? AND batch_no = ? AND enable_polling = true", lastID, batchNo). + Order("id ASC"). + Limit(limit). + Find(&cards).Error + if err != nil { + return nil, fmt.Errorf("查询迁移卡失败: %w", err) + } + return cards, nil +} + +func updateStats(stats *finalizeStats, configMgr *polling.PollingConfigManager, cards []*model.IotCard) { + for _, card := range cards { + stats.pollingCards++ + if stats.firstID == 0 { + stats.firstID = card.ID + } + stats.lastID = card.ID + + intervals := configMgr.MergedTaskIntervals(card) + if len(intervals) == 0 { + stats.skippedNoConfig++ + continue + } + stats.enqueuedCards++ + for taskType := range intervals { + stats.taskCounts[taskType]++ + } + } +} + +func collectCardIDs(cards []*model.IotCard) []uint { + ids := make([]uint, 0, len(cards)) + for _, card := range cards { + ids = append(ids, card.ID) + } + return ids +} + +func printStats(stats finalizeStats, opts finalizeOptions) { + mode := "dry-run" + if opts.apply { + mode = "apply" + } + fmt.Printf("迁移收尾完成: mode=%s batch_no=%s\n", mode, opts.batchNo) + fmt.Printf("导入卡总数: %d\n", stats.importedTotal) + fmt.Printf("启用轮询卡: %d\n", stats.pollingCards) + fmt.Printf("可入队卡: %d\n", stats.enqueuedCards) + fmt.Printf("无匹配轮询配置: %d\n", stats.skippedNoConfig) + if stats.firstID > 0 { + fmt.Printf("处理 ID 范围: %d-%d\n", stats.firstID, stats.lastID) + } + + taskTypes := make([]string, 0, len(stats.taskCounts)) + for taskType := range stats.taskCounts { + taskTypes = append(taskTypes, taskType) + } + sort.Strings(taskTypes) + for _, taskType := range taskTypes { + fmt.Printf("任务入队预估: %s=%d\n", taskType, stats.taskCounts[taskType]) + } + if !opts.apply { + fmt.Println("当前为 dry-run,没有写入 Redis；确认无误后增加 --apply 执行。") + } +} diff --git a/internal/polling/config_manager.go b/internal/polling/config_manager.go index 267a819..2d5f15b 100644 --- a/internal/polling/config_manager.go +++ b/internal/polling/config_manager.go @@ -45,6 +45,16 @@ func NewPollingConfigManager(configStore *postgres.PollingConfigStore, redisClie // Load 加载配置到内存缓存并同步到 Redis // 加载失败时保留原缓存不清空，确保可用性 func (m *PollingConfigManager) Load(ctx context.Context) error { + return m.load(ctx, true) +} + +// LoadForInspection 加载配置到内存缓存但不同步 Redis。 +// 用于迁移 dry-run 等只读检查场景，避免检查动作产生外部副作用。 +func (m *PollingConfigManager) LoadForInspection(ctx context.Context) error { + return m.load(ctx, false) +} + +func (m *PollingConfigManager) load(ctx context.Context, syncRedis bool) error { configs, err := m.configStore.ListEnabled(ctx) if err != nil { m.logger.Error("加载轮询配置失败", zap.Error(err)) @@ -55,8 +65,12 @@ func (m *PollingConfigManager) Load(ctx context.Context) error { m.configs = configs m.mu.Unlock() - if syncErr := m.syncToRedis(ctx, configs); syncErr != nil { - m.logger.Warn("同步配置到 Redis 失败", zap.Error(syncErr)) + if syncRedis && m.redis != nil { + if syncErr := m.syncToRedis(ctx, configs); syncErr != nil { + m.logger.Warn("同步配置到 Redis 失败", zap.Error(syncErr)) + } + } else if syncRedis { + m.logger.Warn("Redis 客户端为空，跳过轮询配置同步") } m.logger.Info("轮询配置已加载", zap.Int("count", len(configs))) diff --git a/internal/polling/initializer.go b/internal/polling/initializer.go index 54a3b12..4ebd892 100644 --- a/internal/polling/initializer.go +++ b/internal/polling/initializer.go @@ -123,6 +123,12 @@ func (p *PollingInitializer) GetProgress() InitProgress { } } +// InitCards 将指定卡批量写入轮询分片队列和卡缓存。 +// 用于迁移收尾等按批次重建场景；调用方负责决定卡集合和是否先清理旧队列成员。 +func (p *PollingInitializer) InitCards(ctx context.Context, cards []*model.IotCard) error { + return p.initBatch(ctx, cards) +} + // run 执行渐进式初始化 func (p *PollingInitializer) run(ctx context.Context) { defer p.wg.Done() diff --git a/internal/polling/queue_manager.go b/internal/polling/queue_manager.go index 7275893..a69aace 100644 --- a/internal/polling/queue_manager.go +++ b/internal/polling/queue_manager.go @@ -113,6 +113,26 @@ func (m *PollingQueueManager) RemoveFromAllQueues(ctx context.Context, cardID ui return err } +// RemoveBatchFromCurrentShardQueues 批量移除卡在当前分片队列中的旧成员，并清理卡缓存。 +// 迁移收尾重建队列前使用，避免重复执行时残留旧任务类型或旧缓存。 +func (m *PollingQueueManager) RemoveBatchFromCurrentShardQueues(ctx context.Context, cardIDs []uint) error { + if len(cardIDs) == 0 { + return nil + } + pipe := m.redis.Pipeline() + for _, cardID := range cardIDs { + member := fmt.Sprintf("%d", cardID) + shardID := int(cardID) % m.shardCount + for _, taskType := range allTaskTypes { + key := constants.RedisPollingShardQueueKey(shardID, taskType) + pipe.ZRem(ctx, key, member) + } + pipe.Del(ctx, constants.RedisPollingCardInfoKey(cardID)) + } + _, err := pipe.Exec(ctx) + return err +} + // EnqueueManual 手动触发入队（List RPUSH，调度器优先消费） func (m *PollingQueueManager) EnqueueManual(ctx context.Context, cardID uint, taskType string) error { key := constants.RedisPollingManualQueueKey(taskType) diff --git a/internal/service/iot_card/service.go b/internal/service/iot_card/service.go index 3d29f46..4bcbbe9 100644 --- a/internal/service/iot_card/service.go +++ b/internal/service/iot_card/service.go @@ -1499,6 +1499,7 @@ func (s *Service) RefreshCardDataFromGateway(ctx context.Context, iccid string) // 返回本次增量值（MB），用于触发套餐扣减 func (s *Service) calculateRefreshFlowUpdates(card *model.IotCard, gatewayFlowMB float64, now time.Time, updates map[string]any) float64 { increment := gatewayFlowMB - card.LastGatewayReadingMB + shouldUpdateGatewayReading := true if increment < 0 { // 当前值比上次小，检查是否为上游运营商正常重置 @@ -1518,11 +1519,13 @@ func (s *Service) calculateRefreshFlowUpdates(card *model.IotCard, gatewayFlowMB zap.Float64("gateway_flow", gatewayFlowMB), zap.Float64("last_reading", card.LastGatewayReadingMB)) increment = 0 + shouldUpdateGatewayReading = false } } - // 始终更新网关读数基准 - updates["last_gateway_reading_mb"] = gatewayFlowMB + if shouldUpdateGatewayReading { + updates["last_gateway_reading_mb"] = gatewayFlowMB + } // 检测跨自然月 currentMonthStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, now.Location()) diff --git a/internal/task/polling_carddata_handler.go b/internal/task/polling_carddata_handler.go index c3b8736..c50eceb 100644 --- a/internal/task/polling_carddata_handler.go +++ b/internal/task/polling_carddata_handler.go @@ -109,7 +109,10 @@ func (h *PollingCarddataHandler) Handle(ctx context.Context, t *asynq.Task) erro go h.base.insertDataUsageRecord(context.Background(), cardID, flowIncrementMB, now) - cacheUpdates := map[string]any{"last_gateway_reading_mb": gatewayFlowMB} + cacheUpdates := map[string]any{} + if reading, ok := updates["last_gateway_reading_mb"]; ok { + cacheUpdates["last_gateway_reading_mb"] = reading + } if isCrossMonth { // 跨月时同步更新缓存中的本月开始日期和本月用量，避免下次从缓存读取时再次误判跨月 currentMonthStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, now.Location()) @@ -150,6 +153,7 @@ func (h *PollingCarddataHandler) calculateFlowUpdates(card *model.IotCard, gatew updates := make(map[string]any) increment := gatewayFlowMB - card.LastGatewayReadingMB + shouldUpdateGatewayReading := true if increment < 0 { if isResetWindow(now, resetDay) { increment = gatewayFlowMB @@ -164,10 +168,13 @@ func (h *PollingCarddataHandler) calculateFlowUpdates(card *model.IotCard, gatew zap.Float64("gateway_flow", gatewayFlowMB), zap.Float64("last_reading", card.LastGatewayReadingMB)) increment = 0 + shouldUpdateGatewayReading = false } } - updates["last_gateway_reading_mb"] = gatewayFlowMB + if shouldUpdateGatewayReading { + updates["last_gateway_reading_mb"] = gatewayFlowMB + } currentMonthStart := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, now.Location()) isCrossMonth := card.CurrentMonthStartDate == nil ||