0d96a94e5f
All checks were successful
构建并部署到测试环境（无 SSH） / build-and-deploy (push) Successful in 8m25s
294 lines
8.0 KiB
Go
294 lines
8.0 KiB
Go
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 执行。")
|
||
}
|
||
}
|