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 执行。") } }