junhong_cmp_fiber/cmd/worker/polling_ingest.go

package main

import (
	"context"
	"time"

	"go.uber.org/zap"

	"github.com/break/junhong_cmp_fiber/internal/infrastructure/nsqbus"
	"github.com/break/junhong_cmp_fiber/internal/store/postgres"
	"github.com/break/junhong_cmp_fiber/internal/task"
	"github.com/break/junhong_cmp_fiber/pkg/config"
	"github.com/break/junhong_cmp_fiber/pkg/constants"
)

// pollingIngestRuntime 是轮询观测入库异步化在 Worker 内的运行组件（add-polling-async-ingest）：
// 入库消费者、采集记录批量消费者、活动优先项索引重建与积压监控。入库消费者常驻订阅，
// 使切回同步模式后队列中已有的入库消息仍会被处理完。
type pollingIngestRuntime struct {
	runtime     *workerRuntime
	nsqdAddress string
	executor    *task.IngestExecutor
	collector   *task.CollectionBatchConsumer
	ingest      *nsqbus.ConsumerManager
	collection  *nsqbus.ConsumerManager
	rebuilder   *task.PriorityIndexRebuilder
	stopCollect context.CancelFunc
	grace       time.Duration
	logger      *zap.Logger
}

// startPollingIngest 注入入库发布器并启动入库与采集记录消费者；未配置 nsqd 时返回 nil，轮询保持同步入库。
func startPollingIngest(ctx context.Context, cfg *config.Config, runtime *workerRuntime, producer *nsqbus.Producer,
	singleton bool, grace time.Duration, logger *zap.Logger) *pollingIngestRuntime {
	carrierCfg := cfg.Polling.Carrier
	if producer == nil || carrierCfg.NSQDAddress == "" {
		return nil
	}
	checkIngestConnBudget(cfg, logger)
	publisher := task.NewIngestPublisher(producer)
	runtime.pollingBase.SetIngest(publisher, func(ctx context.Context) bool {
		return runtime.carrierSettings.Current(ctx).IngestAsync()
	})
	observation := runtime.workerResult.Services.CardObservation
	if observation != nil {
		observation.SetCollectionSink(publisher)
	}
	r := &pollingIngestRuntime{
		runtime: runtime, nsqdAddress: carrierCfg.NSQDAddress, grace: grace, logger: logger,
		executor:  task.NewIngestExecutor(runtime.pollingBase, observation, cfg.Polling.Ingest.EffectiveTimeout()),
		collector: task.NewCollectionBatchConsumer(runtime.db, cfg.Polling.Collection.EffectiveBatchSize(), logger),
		rebuilder: task.NewPriorityIndexRebuilder(postgres.NewPollingPriorityItemStore(runtime.db), runtime.redisClient, logger),
	}
	collectCtx, stopCollect := context.WithCancel(context.Background())
	r.stopCollect = stopCollect
	r.collector.Start(collectCtx)
	consumerCfg := nsqbus.ConsumerConfig{
		NSQDAddress: carrierCfg.NSQDAddress, NSQLookupdAddresses: nsqbus.SplitAddresses(carrierCfg.NSQLookupdAddresses),
		MsgTimeout: carrierCfg.MsgTimeout, UnlimitedAttempts: true,
	}
	r.ingest = nsqbus.NewRetryConsumerManager(consumerCfg, constants.NSQChannelWorker, r.executor.Handle, logger)
	r.collection = nsqbus.NewRetryConsumerManager(consumerCfg, constants.NSQChannelWorker, r.collector.Handle, logger)
	r.ingest.Sync(map[string]int{constants.NSQCardIngestTopic: cfg.Polling.Ingest.EffectiveConcurrency()})
	r.collection.Sync(map[string]int{constants.NSQCardCollectionTopic: constants.PollingCollectionMaxInFlight})
	logger.Info("轮询观测入库消费者已启动", zap.Int("ingest_concurrency", cfg.Polling.Ingest.EffectiveConcurrency()),
		zap.Duration("ingest_timeout", cfg.Polling.Ingest.EffectiveTimeout()),
		zap.Int("collection_batch_size", cfg.Polling.Collection.EffectiveBatchSize()))
	if singleton {
		go r.loop(ctx)
	}
	return r
}

// checkIngestConnBudget 启动时检查入库并发与采集记录写并发不超过数据库连接数减去预留，超过只记 ERROR、不阻止启动。
func checkIngestConnBudget(cfg *config.Config, logger *zap.Logger) {
	maxOpen := cfg.Database.MaxOpenConns
	if maxOpen <= 0 {
		maxOpen = constants.DefaultMaxOpenConns
	}
	need := cfg.Polling.Ingest.EffectiveConcurrency() + constants.PollingIngestCollectionWriters
	if budget := maxOpen - constants.PollingIngestConnReserve; need > budget {
		logger.Error("入库并发超出数据库连接预算，请调小 polling.ingest.concurrency 或调大数据库连接数",
			zap.Int("ingest_concurrency", cfg.Polling.Ingest.EffectiveConcurrency()),
			zap.Int("collection_writers", constants.PollingIngestCollectionWriters),
			zap.Int("max_open_conns", maxOpen), zap.Int("reserved", constants.PollingIngestConnReserve))
	}
}

// loop 在 leader 上定期重建活动优先项索引并检查入库积压。
func (r *pollingIngestRuntime) loop(ctx context.Context) {
	r.rebuilder.Rebuild(ctx)
	rebuild := time.NewTicker(constants.PollingPriorityIndexRebuildEvery)
	monitor := time.NewTicker(constants.PollingIngestMonitorInterval)
	defer rebuild.Stop()
	defer monitor.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-rebuild.C:
			r.rebuilder.Rebuild(ctx)
		case <-monitor.C:
			r.checkBacklog(ctx)
		}
	}
}

// checkBacklog 记录入库与采集记录积压，超过告警阈值或等待过久记 ERROR；入库积压超过反压阈值时暂停派发新卡任务，降到一半以下恢复。
func (r *pollingIngestRuntime) checkBacklog(ctx context.Context) {
	defer func() {
		if recovered := recover(); recovered != nil {
			r.logger.Error("入库积压检查 panic，已恢复", zap.Any("panic", recovered))
		}
	}()
	settings := r.runtime.carrierSettings.Current(ctx)
	ingestDepth, ingestErr := nsqbus.TopicBacklog(ctx, r.nsqdAddress, constants.NSQCardIngestTopic)
	collectionDepth, collectionErr := nsqbus.TopicBacklog(ctx, r.nsqdAddress, constants.NSQCardCollectionTopic)
	if ingestErr != nil || collectionErr != nil {
		r.logger.Warn("读取入库积压失败", zap.NamedError("ingest_error", ingestErr), zap.NamedError("collection_error", collectionErr))
		return
	}
	lag := r.executor.LastLag()
	fields := []zap.Field{zap.Int64("ingest_backlog", ingestDepth), zap.Int64("collection_backlog", collectionDepth),
		zap.Duration("ingest_lag", lag), zap.String("ingest_mode", settings.IngestMode)}
	r.logger.Info("轮询入库积压", append(fields, zap.String("metric", "polling.ingest.backlog"))...)
	alertDepth := int64(settings.IngestAlertDepth)
	if ingestDepth > alertDepth || collectionDepth > alertDepth || lag > constants.PollingIngestLagAlert {
		r.logger.Error("轮询入库积压超过告警阈值", append(fields, zap.Int64("alert_depth", alertDepth))...)
	}
	r.applyBackpressure(ctx, ingestDepth, int64(settings.IngestPauseDepth))
}

// applyBackpressure 按积压置位或清除暂停标记；标记带 TTL，leader 失联时自动解除。
func (r *pollingIngestRuntime) applyBackpressure(ctx context.Context, depth, pauseDepth int64) {
	redisClient := r.runtime.redisClient
	paused, _ := redisClient.Exists(ctx, constants.RedisPollingIngestPausedKey).Result()
	ttl := 3 * constants.PollingIngestMonitorInterval
	switch {
	case depth > pauseDepth:
		redisClient.Set(ctx, constants.RedisPollingIngestPausedKey, depth, ttl)
		if paused == 0 {
			r.logger.Error("入库积压超过反压阈值，暂停派发新卡任务", zap.Int64("ingest_backlog", depth), zap.Int64("pause_depth", pauseDepth))
		}
	case paused > 0 && depth > pauseDepth/2:
		redisClient.Expire(ctx, constants.RedisPollingIngestPausedKey, ttl)
	case paused > 0:
		redisClient.Del(ctx, constants.RedisPollingIngestPausedKey)
		r.logger.Info("入库积压已降到反压阈值一半以下，恢复派发卡任务", zap.Int64("ingest_backlog", depth), zap.Int64("pause_depth", pauseDepth))
	}
}

// shutdown 先停入库消费者（处理中的消息在宽限期内完成），再停采集记录消费者，最后停止攒批循环（写完已缓冲的一批）。
func (r *pollingIngestRuntime) shutdown() {
	if r == nil {
		return
	}
	r.ingest.StopAll(r.grace)
	r.collection.StopAll(r.grace)
	r.stopCollect()
}