package task import ( "context" "fmt" "strconv" "strings" "time" "github.com/hibiken/asynq" "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" ) // DailyTrafficFlushHandler 每日流量落盘任务处理器 // 凌晨 2 点分批 SCAN 昨日 Redis key → UPSERT 到 tb_card_daily_usage → 删 key type DailyTrafficFlushHandler struct { redis *redis.Client cardDailyUsageStore *postgres.CardDailyUsageStore logger *zap.Logger } func NewDailyTrafficFlushHandler( redisClient *redis.Client, cardDailyUsageStore *postgres.CardDailyUsageStore, logger *zap.Logger, ) *DailyTrafficFlushHandler { return &DailyTrafficFlushHandler{ redis: redisClient, cardDailyUsageStore: cardDailyUsageStore, logger: logger, } } func (h *DailyTrafficFlushHandler) HandleDailyTrafficFlush(ctx context.Context, _ *asynq.Task) error { startTime := time.Now() yesterday := time.Now().AddDate(0, 0, -1).Format("2006-01-02") pattern := fmt.Sprintf("traffic:daily:*:%s", yesterday) h.logger.Info("开始每日流量落盘", zap.String("date", yesterday), zap.String("pattern", pattern)) // 分批 SCAN 收集所有昨日 key var allKeys []string var cursor uint64 for { keys, nextCursor, err := h.redis.Scan(ctx, cursor, pattern, 500).Result() if err != nil { h.logger.Error("SCAN Redis key 失败", zap.Error(err)) return err } allKeys = append(allKeys, keys...) cursor = nextCursor if cursor == 0 { break } } if len(allKeys) == 0 { h.logger.Info("无需落盘,昨日无流量数据", zap.String("date", yesterday)) return nil } h.logger.Info("收集到待落盘 key", zap.String("date", yesterday), zap.Int("count", len(allKeys))) yesterdayDate, _ := time.Parse("2006-01-02", yesterday) // 分批处理:每批 200 条 flushedCount := 0 for i := 0; i < len(allKeys); i += 200 { end := i + 200 if end > len(allKeys) { end = len(allKeys) } batch := allKeys[i:end] var records []*model.CardDailyUsage var flushedKeys []string for _, key := range batch { cardID, err := parseCardIDFromKey(key) if err != nil { h.logger.Warn("解析 Redis key 失败", zap.String("key", key), zap.Error(err)) continue } val, err := h.redis.Get(ctx, key).Float64() if err != nil { h.logger.Warn("读取 Redis 值失败", zap.String("key", key), zap.Error(err)) continue } records = append(records, &model.CardDailyUsage{ IotCardID: cardID, Date: yesterdayDate, UsageMB: val, }) flushedKeys = append(flushedKeys, key) } // 批量 UPSERT 到 DB if len(records) > 0 { if err := h.cardDailyUsageStore.UpsertBatch(ctx, records); err != nil { h.logger.Error("批量 UPSERT 失败", zap.Int("batch_size", len(records)), zap.Error(err)) continue } } // 批量删除已落盘的 Redis key if len(flushedKeys) > 0 { pipe := h.redis.Pipeline() for _, key := range flushedKeys { pipe.Del(ctx, key) } if _, err := pipe.Exec(ctx); err != nil { h.logger.Warn("批量删除 Redis key 失败", zap.Error(err)) } } flushedCount += len(records) } h.logger.Info("每日流量落盘完成", zap.String("date", yesterday), zap.Int("flushed_count", flushedCount), zap.Duration("duration", time.Since(startTime))) return nil } // parseCardIDFromKey 从 Redis key "traffic:daily:{cardID}:{date}" 中解析 cardID func parseCardIDFromKey(key string) (uint, error) { parts := strings.Split(key, ":") if len(parts) != 4 { return 0, fmt.Errorf("invalid key format: %s", key) } id, err := strconv.ParseUint(parts[2], 10, 64) if err != nil { return 0, err } return uint(id), nil }