package polling import ( "context" "fmt" "sync" "sync/atomic" "time" "github.com/redis/go-redis/v9" "go.uber.org/zap" "github.com/break/junhong_cmp_fiber/internal/model" "github.com/break/junhong_cmp_fiber/internal/store/postgres" "github.com/break/junhong_cmp_fiber/pkg/constants" ) // initProgress 初始化进度内部状态(字段通过 GetProgress 快照对外,避免暴露锁) type initProgress struct { mu sync.RWMutex totalCards int64 loadedCards int64 startTime time.Time lastBatchTime time.Time status string errorMessage string } // InitProgress 初始化进度快照(GetProgress 返回的值拷贝,调用方无需加锁) type InitProgress struct { TotalCards int64 `json:"total_cards"` LoadedCards int64 `json:"loaded_cards"` StartTime time.Time `json:"start_time"` LastBatchTime time.Time `json:"last_batch_time"` Status string `json:"status"` ErrorMessage string `json:"error_message"` } // pipelineFlushSize 单次 Pipeline Exec 最大命令数(每张卡最多 6 条) // 10000 条约 2MB 内存峰值;千万级卡下约 6000 次 Exec,RTT 总开销可忽略 const pipelineFlushSize = 10000 // PollingInitializer 分片渐进式初始化器 // 启动时从 DB 分批加载全量卡数据到分片 Sorted Set // 使用 card_id % shardCount 分片,每 pipelineFlushSize 条命令 flush 一次 Pipeline type PollingInitializer struct { iotCardStore *postgres.IotCardStore redis *redis.Client configMgr *PollingConfigManager queueMgr *PollingQueueManager logger *zap.Logger progress initProgress initCompleted atomic.Bool stopChan chan struct{} wg sync.WaitGroup } // NewPollingInitializer 创建初始化器 func NewPollingInitializer( iotCardStore *postgres.IotCardStore, redisClient *redis.Client, configMgr *PollingConfigManager, queueMgr *PollingQueueManager, logger *zap.Logger, ) *PollingInitializer { p := &PollingInitializer{ iotCardStore: iotCardStore, redis: redisClient, configMgr: configMgr, queueMgr: queueMgr, logger: logger, stopChan: make(chan struct{}), } p.progress.status = "pending" return p } // StartBackground 启动后台渐进式初始化(非阻塞) func (p *PollingInitializer) StartBackground(ctx context.Context) { p.wg.Add(1) go p.run(ctx) } // Stop 停止初始化 func (p *PollingInitializer) Stop() { close(p.stopChan) p.wg.Wait() } // IsCompleted 检查初始化是否完成 func (p *PollingInitializer) IsCompleted() bool { return p.initCompleted.Load() } // GetProgress 返回当前初始化进度快照(加锁读取,返回值拷贝) func (p *PollingInitializer) GetProgress() InitProgress { p.progress.mu.RLock() defer p.progress.mu.RUnlock() return InitProgress{ TotalCards: p.progress.totalCards, LoadedCards: p.progress.loadedCards, StartTime: p.progress.startTime, LastBatchTime: p.progress.lastBatchTime, Status: p.progress.status, ErrorMessage: p.progress.errorMessage, } } // run 执行渐进式初始化 func (p *PollingInitializer) run(ctx context.Context) { defer p.wg.Done() p.setStatus("running", "") p.progress.mu.Lock() p.progress.startTime = time.Now() p.progress.mu.Unlock() p.logger.Info("开始分片渐进式初始化...") totalCards, err := p.iotCardStore.CountForPolling(ctx) if err != nil { p.logger.Error("获取卡总数失败", zap.Error(err)) p.setStatus("failed", err.Error()) return } p.progress.mu.Lock() p.progress.totalCards = totalCards p.progress.mu.Unlock() p.logger.Info("开始加载卡数据", zap.Int64("total_cards", totalCards)) const batchSize = 100000 const batchSleep = 500 * time.Millisecond var lastID uint batchCount := 0 for { select { case <-p.stopChan: p.logger.Info("渐进式初始化被中断") return default: } cards, fetchErr := p.iotCardStore.ListForPollingBatch(ctx, lastID, batchSize) if fetchErr != nil { p.logger.Error("加载卡数据失败", zap.Error(fetchErr)) p.setStatus("failed", fetchErr.Error()) return } if len(cards) == 0 { break } if initErr := p.initBatch(ctx, cards); initErr != nil { p.logger.Warn("批量初始化失败", zap.Error(initErr)) } lastID = cards[len(cards)-1].ID batchCount++ p.progress.mu.Lock() p.progress.loadedCards += int64(len(cards)) p.progress.lastBatchTime = time.Now() loaded := p.progress.loadedCards p.progress.mu.Unlock() if batchCount%10 == 0 || len(cards) < batchSize { p.logger.Info("初始化进度", zap.Int("batch", batchCount), zap.Int64("loaded", loaded), zap.Int64("total", totalCards)) } time.Sleep(batchSleep) } p.setStatus("completed", "") p.initCompleted.Store(true) snapshot := p.GetProgress() p.logger.Info("分片渐进式初始化完成", zap.Int64("total_loaded", snapshot.LoadedCards), zap.Duration("duration", time.Since(snapshot.StartTime))) } // initBatch 使用 Pipeline 将一批卡写入分片队列和缓存 // 每 pipelineFlushSize 条命令 Exec 一次,控制内存峰值并降低单次失败损失 func (p *PollingInitializer) initBatch(ctx context.Context, cards []*model.IotCard) error { if len(cards) == 0 { return nil } now := time.Now() cardCacheTTL := 7 * 24 * time.Hour pipe := p.redis.Pipeline() cmdCount := 0 enqueuedCards := 0 skippedCards := 0 flushPipe := func() { if cmdCount == 0 { return } if _, execErr := pipe.Exec(ctx); execErr != nil { p.logger.Warn("Pipeline flush 失败,继续下一批", zap.Error(execErr)) } pipe = p.redis.Pipeline() cmdCount = 0 } for _, card := range cards { cfg := p.configMgr.MatchConfig(card) if cfg == nil { skippedCards++ continue } enqueuedCards++ shardID := int(card.ID) % p.queueMgr.shardCount cardIDStr := fmt.Sprintf("%d", card.ID) if cfg.RealnameCheckInterval != nil && *cfg.RealnameCheckInterval > 0 { nextCheck := calculateNextCheckTime(card.LastRealNameCheckAt, *cfg.RealnameCheckInterval, now) pipe.ZAdd(ctx, constants.RedisPollingShardQueueKey(shardID, constants.TaskTypePollingRealname), redis.Z{ Score: float64(nextCheck.Unix()), Member: cardIDStr, }) cmdCount++ } if cfg.CarddataCheckInterval != nil && *cfg.CarddataCheckInterval > 0 { nextCheck := calculateNextCheckTime(card.LastDataCheckAt, *cfg.CarddataCheckInterval, now) pipe.ZAdd(ctx, constants.RedisPollingShardQueueKey(shardID, constants.TaskTypePollingCarddata), redis.Z{ Score: float64(nextCheck.Unix()), Member: cardIDStr, }) cmdCount++ } if cfg.PackageCheckInterval != nil && *cfg.PackageCheckInterval > 0 { nextCheck := calculateNextCheckTime(card.LastDataCheckAt, *cfg.PackageCheckInterval, now) pipe.ZAdd(ctx, constants.RedisPollingShardQueueKey(shardID, constants.TaskTypePollingPackage), redis.Z{ Score: float64(nextCheck.Unix()), Member: cardIDStr, }) cmdCount++ } if cfg.ProtectCheckInterval != nil && *cfg.ProtectCheckInterval > 0 { nextCheck := calculateNextCheckTime(card.LastProtectCheckAt, *cfg.ProtectCheckInterval, now) pipe.ZAdd(ctx, constants.RedisPollingShardQueueKey(shardID, constants.TaskTypePollingProtect), redis.Z{ Score: float64(nextCheck.Unix()), Member: cardIDStr, }) cmdCount++ } cacheKey := constants.RedisPollingCardInfoKey(card.ID) cacheData := map[string]interface{}{ "id": card.ID, "iccid": card.ICCID, "card_category": card.CardCategory, "real_name_status": card.RealNameStatus, "network_status": card.NetworkStatus, "carrier_id": card.CarrierID, "current_month_usage_mb": card.CurrentMonthUsageMB, "last_gateway_reading_mb": card.LastGatewayReadingMB, "data_usage_mb": card.DataUsageMB, "stop_reason": card.StopReason, "is_standalone": boolToStr(card.IsStandalone), "cached_at": now.Unix(), } if card.CurrentMonthStartDate != nil { cacheData["current_month_start_date"] = card.CurrentMonthStartDate.Unix() } pipe.HSet(ctx, cacheKey, cacheData) pipe.Expire(ctx, cacheKey, cardCacheTTL) cmdCount += 2 if cmdCount >= pipelineFlushSize { flushPipe() } } flushPipe() p.logger.Info("批量初始化完成", zap.Int("total", len(cards)), zap.Int("enqueued", enqueuedCards), zap.Int("skipped_no_config", skippedCards)) return nil } // calculateNextCheckTime 计算下次检查时间 func calculateNextCheckTime(lastCheckAt *time.Time, intervalSeconds int, now time.Time) time.Time { if lastCheckAt == nil { jitter := time.Duration(now.UnixNano()%int64(intervalSeconds)) * time.Second / 10 return now.Add(jitter) } nextCheck := lastCheckAt.Add(time.Duration(intervalSeconds) * time.Second) if nextCheck.Before(now) { return now } return nextCheck } func (p *PollingInitializer) setStatus(status, errMsg string) { p.progress.mu.Lock() p.progress.status = status p.progress.errorMessage = errMsg p.progress.mu.Unlock() }