junhong_cmp_fiber/internal/polling/initializer.go

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

	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 {
			continue
		}

		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()
	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()
}