0d96a94e5f
All checks were successful
构建并部署到测试环境（无 SSH） / build-and-deploy (push) Successful in 8m25s
349 lines
9.9 KiB
Go
349 lines
9.9 KiB
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
|
||
restartQueued 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()
|
||
}
|
||
|
||
// Restart 重新执行初始化，用于配置变更后补齐新增任务类型的分片队列。
|
||
// 若初始化正在进行中，则记录一次待重启请求，当前轮完成后再自动执行。
|
||
func (p *PollingInitializer) Restart(ctx context.Context) {
|
||
if !p.initCompleted.CompareAndSwap(true, false) {
|
||
p.restartQueued.Store(true)
|
||
p.logger.Info("轮询初始化仍在进行中，已标记完成后重新初始化")
|
||
return
|
||
}
|
||
p.setStatus("pending", "")
|
||
p.wg.Add(1)
|
||
go p.run(ctx)
|
||
p.logger.Info("轮询初始化已重新启动（配置变更触发）")
|
||
}
|
||
|
||
// 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,
|
||
}
|
||
}
|
||
|
||
// InitCards 将指定卡批量写入轮询分片队列和卡缓存。
|
||
// 用于迁移收尾等按批次重建场景；调用方负责决定卡集合和是否先清理旧队列成员。
|
||
func (p *PollingInitializer) InitCards(ctx context.Context, cards []*model.IotCard) error {
|
||
return p.initBatch(ctx, cards)
|
||
}
|
||
|
||
// 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.Error("批量初始化失败，该批次卡未入轮询队列",
|
||
zap.Uint("batch_start_id", lastID), 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)))
|
||
|
||
if p.restartQueued.Swap(false) {
|
||
select {
|
||
case <-ctx.Done():
|
||
return
|
||
default:
|
||
}
|
||
if p.initCompleted.CompareAndSwap(true, false) {
|
||
p.setStatus("pending", "")
|
||
p.wg.Add(1)
|
||
go p.run(ctx)
|
||
p.logger.Info("检测到初始化期间配置变更，完成后再次初始化")
|
||
}
|
||
}
|
||
}
|
||
|
||
// 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
|
||
}
|
||
cmds, execErr := pipe.Exec(ctx)
|
||
if execErr != nil {
|
||
p.logger.Error("Pipeline flush 失败，部分卡可能未入队", zap.Error(execErr))
|
||
}
|
||
// 逐条检查，记录具体失败命令（不中断批次）
|
||
for _, cmd := range cmds {
|
||
if cmdErr := cmd.Err(); cmdErr != nil && cmdErr != redis.Nil {
|
||
p.logger.Warn("Pipeline 单条命令失败",
|
||
zap.String("cmd", cmd.Name()), zap.Error(cmdErr))
|
||
}
|
||
}
|
||
pipe = p.redis.Pipeline()
|
||
cmdCount = 0
|
||
}
|
||
|
||
for _, card := range cards {
|
||
intervals := p.configMgr.MergedTaskIntervals(card)
|
||
if len(intervals) == 0 {
|
||
skippedCards++
|
||
continue
|
||
}
|
||
enqueuedCards++
|
||
|
||
shardID := int(card.ID) % p.queueMgr.shardCount
|
||
cardIDStr := fmt.Sprintf("%d", card.ID)
|
||
|
||
for taskType, info := range intervals {
|
||
lastCheckAt := lastCheckAtByTaskType(card, taskType)
|
||
nextCheck := calculateNextCheckTime(lastCheckAt, info.Interval, now)
|
||
pipe.ZAdd(ctx, constants.RedisPollingShardQueueKey(shardID, taskType), 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
|
||
}
|
||
|
||
// lastCheckAtByTaskType 根据 task type 返回对应的上次检查时间字段
|
||
func lastCheckAtByTaskType(card *model.IotCard, taskType string) *time.Time {
|
||
switch taskType {
|
||
case constants.TaskTypePollingRealname:
|
||
return card.LastRealNameCheckAt
|
||
case constants.TaskTypePollingCarddata, constants.TaskTypePollingPackage:
|
||
return card.LastDataCheckAt
|
||
case constants.TaskTypePollingProtect:
|
||
return card.LastProtectCheckAt
|
||
case constants.TaskTypePollingCardStatus:
|
||
return card.LastCardStatusCheckAt
|
||
default:
|
||
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()
|
||
}
|