junhong_cmp_fiber/internal/polling/carrierqueue/dispatcher.go

package carrierqueue

import (
	"context"
	"encoding/json"
	stderrors "errors"
	"strconv"
	"sync"
	"time"

	"github.com/redis/go-redis/v9"
	"go.uber.org/zap"
	"gorm.io/gorm"

	"github.com/break/junhong_cmp_fiber/internal/model"
	"github.com/break/junhong_cmp_fiber/pkg/config"
	"github.com/break/junhong_cmp_fiber/pkg/constants"
)

// CardMessage 是卡任务 NSQ 消息体。
type CardMessage struct {
	CardID       uint   `json:"card_id"`
	CarrierID    uint   `json:"carrier_id"`
	Token        string `json:"token"`
	DispatchedAt int64  `json:"dispatched_at"`
}

// Publisher 发布 NSQ 消息。
type Publisher interface {
	Publish(topic string, body []byte) error
}

// BudgetProvider 返回账户各查询操作所属限流组本秒剩余名额；limited=false 表示不按 QPS 限制领取。
type BudgetProvider interface {
	Budget(ctx context.Context, carrierID uint, ops []string) (groups map[string]int, opGroup map[string]string, limited bool)
}

// CarrierParams 是一个运营商账户生效的调度参数（已合并全局默认）。
type CarrierParams struct {
	InflightLimit   int
	CardMaxInFlight int
	PriorityRatio   float64
	Configured      bool // 是否有该账户的配置行

	// 停复机指令（add-stop-resume-command-queue）。
	SRInflightLimit int           // 账户停复机执行上限
	SRMaxInFlight   int           // 每进程停复机消费者 max_in_flight
	SRMinInterval   time.Duration // 单卡最小重试间隔（停机与复机合计）
	SRConfirmDelay  time.Duration // 结果未知后的确认延迟
}

// carrierEntry 是分发器缓存的账户信息。
type carrierEntry struct {
	id              uint
	realnameControl bool
}

// Dispatcher 每个分发间隔对每个运营商账户抢分发锁、领取到期卡并发布到账户 topic。多实例安全。
type Dispatcher struct {
	db         *gorm.DB
	store      *Store
	settings   *SettingsProvider
	publisher  Publisher
	budgets    BudgetProvider
	cfg        config.CarrierSchedulerConfig
	instanceID string
	logger     *zap.Logger

	mu        sync.Mutex
	carriers  []carrierEntry
	params    map[uint]model.CarrierPollingParam
	loadedAt  time.Time
	tightened map[uint]time.Time
}

// NewDispatcher 创建分发器。
func NewDispatcher(db *gorm.DB, store *Store, settings *SettingsProvider, publisher Publisher, budgets BudgetProvider,
	cfg config.CarrierSchedulerConfig, instanceID string, logger *zap.Logger) *Dispatcher {
	if logger == nil {
		logger = zap.NewNop()
	}
	return &Dispatcher{db: db, store: store, settings: settings, publisher: publisher, budgets: budgets, cfg: cfg,
		instanceID: instanceID, logger: logger, params: map[uint]model.CarrierPollingParam{}, tightened: map[uint]time.Time{}}
}

// Run 按分发间隔循环，直到 ctx 结束。
func (d *Dispatcher) Run(ctx context.Context) {
	interval := d.cfg.DispatchInterval
	if interval <= 0 {
		interval = time.Second
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			d.tick(ctx, interval)
		}
	}
}

// Ready 判断新调度是否可领取：模式为 carrier 且新结构初始化已完成。
func (d *Dispatcher) Ready(ctx context.Context) bool {
	if !d.settings.Current(ctx).CarrierMode() {
		return false
	}
	state, err := d.store.redis.Get(ctx, constants.RedisCarrierSchedulerInitStateKey).Result()
	return err == nil && state == InitStateDone
}

func (d *Dispatcher) tick(ctx context.Context, interval time.Duration) {
	defer func() {
		if r := recover(); r != nil {
			d.logger.Error("分发器 panic，已恢复", zap.Any("panic", r))
		}
	}()
	// 停复机指令与设备判断的分发只看停复机出口开关，不受 scheduler_mode 限制；卡任务只在新调度就绪时领取。
	stopResumeOn := d.settings.Current(ctx).StopResumeQueueEnabled
	cardsReady := d.Ready(ctx)
	if !stopResumeOn && !cardsReady {
		return
	}
	d.store.redis.HSet(ctx, constants.RedisCarrierSchedulerHeartbeatKey, d.instanceID, time.Now().UnixMilli())
	if cardsReady {
		d.drainLegacyHints(ctx)
	}
	d.refreshCarriers(ctx)
	d.mu.Lock()
	carriers := append([]carrierEntry(nil), d.carriers...)
	d.mu.Unlock()
	lockTTL := interval * 9 / 10
	if lockTTL <= 0 {
		lockTTL = 900 * time.Millisecond
	}
	if stopResumeOn {
		d.dispatchDeviceEvals(ctx, lockTTL)
	}
	for _, carrier := range carriers {
		if ctx.Err() != nil {
			return
		}
		ok, err := d.store.redis.SetNX(ctx, constants.RedisCarrierDispatchLockKey(carrier.id), d.instanceID, lockTTL).Result()
		if err != nil || !ok {
			continue
		}
		// 停复机优先：先领取停复机指令占用本秒名额，再领取卡任务。
		if stopResumeOn {
			d.dispatchStopResume(ctx, carrier.id)
		}
		if cardsReady {
			d.dispatchCarrier(ctx, carrier)
		}
	}
}

// StopResumeQueueEnabled 返回停复机出口开关当前值。
func (d *Dispatcher) StopResumeQueueEnabled(ctx context.Context) bool {
	return d.settings.Current(ctx).StopResumeQueueEnabled
}

func (d *Dispatcher) dispatchCarrier(ctx context.Context, carrier carrierEntry) {
	params := d.ParamsFor(carrier.id)
	limit := params.InflightLimit
	// 执行上限不得大于订阅进程 max_in_flight 之和，否则领出的消息在 NSQ 中排队到租约过期；不满足时收紧本轮上限并上报。
	processes := d.cfg.SubscriberProcesses
	if processes <= 0 {
		processes = 1
	}
	if capacity := params.CardMaxInFlight * processes; capacity > 0 && limit > capacity {
		limit = capacity
		d.reportTightened(carrier.id, params.InflightLimit, capacity)
	}
	input := ClaimInput{
		CarrierID: carrier.id, Lease: d.lease(), InflightLimit: limit, PriorityRatio: params.PriorityRatio,
		TokenPrefix:   d.instanceID + ":" + strconv.FormatInt(time.Now().UnixNano(), 36),
		CountRealname: carrier.realnameControl,
	}
	if d.budgets != nil {
		ops := []string{constants.GatewayOpQueryStatus, constants.GatewayOpQueryFlow, constants.GatewayOpQueryRealname}
		if groups, opGroup, limited := d.budgets.Budget(ctx, carrier.id, ops); limited {
			input.GroupBudget = groups
			input.TypeGroup = map[string]string{
				constants.CheckTypeCardStatus: opGroup[constants.GatewayOpQueryStatus],
				constants.CheckTypeCarddata:   opGroup[constants.GatewayOpQueryFlow],
				constants.CheckTypeRealname:   opGroup[constants.GatewayOpQueryRealname],
			}
		}
	}
	claimed, expired, err := d.store.Claim(ctx, input)
	if err != nil {
		d.logger.Warn("运营商账户领取失败", zap.Uint("carrier_id", carrier.id), zap.Error(err))
		return
	}
	if expired > 0 {
		minute := strconv.FormatInt(time.Now().Unix()/60, 10)
		key := constants.RedisCarrierExpiredClaimsKey(carrier.id)
		d.store.redis.HIncrBy(ctx, key, minute, expired)
		d.store.redis.Expire(ctx, key, 2*time.Hour)
		d.logger.Warn("运营商账户处理中标记到期（执行崩溃或超时）", zap.Uint("carrier_id", carrier.id), zap.Int64("count", expired),
			zap.String("metric", "polling.carrier.claim_expired"))
	}
	topic := constants.NSQCardTopic(carrier.id)
	for _, item := range claimed {
		body, _ := json.Marshal(CardMessage{CardID: item.CardID, CarrierID: carrier.id, Token: item.Token, DispatchedAt: time.Now().UnixMilli()})
		if err := d.publisher.Publish(topic, body); err != nil {
			// 发布失败：用独立短超时上下文按 token 删除处理中标记并恢复 score；恢复也失败则等租约到期后自然重新领取。
			releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 2*time.Second)
			if releaseErr := d.store.Release(releaseCtx, item.CardID, carrier.id, item.Token); releaseErr != nil {
				d.logger.Warn("发布失败后释放处理中标记失败", zap.Uint("card_id", item.CardID), zap.Error(releaseErr))
			}
			cancel()
			d.logger.Warn("发布卡任务失败，已释放", zap.Uint("card_id", item.CardID), zap.String("topic", topic), zap.Error(err))
		}
	}
}

func (d *Dispatcher) lease() time.Duration {
	if d.cfg.Lease > 0 {
		return d.cfg.Lease
	}
	return 5 * time.Minute
}

func (d *Dispatcher) reportTightened(carrierID uint, configured, capacity int) {
	d.mu.Lock()
	last := d.tightened[carrierID]
	if time.Since(last) < 10*time.Minute {
		d.mu.Unlock()
		return
	}
	d.tightened[carrierID] = time.Now()
	d.mu.Unlock()
	d.logger.Warn("账户执行上限大于订阅进程 max_in_flight 之和，已按消费能力收紧本轮领取上限",
		zap.Uint("carrier_id", carrierID), zap.Int("inflight_limit", configured), zap.Int("capacity", capacity),
		zap.String("metric", "polling.carrier.inflight_tightened"))
}

// refreshCarriers 每 10 秒重新读取运营商账户与调度参数。
func (d *Dispatcher) refreshCarriers(ctx context.Context) {
	d.mu.Lock()
	fresh := time.Since(d.loadedAt) < 10*time.Second
	d.mu.Unlock()
	if fresh {
		return
	}
	var carriers []model.Carrier
	if err := d.db.WithContext(ctx).Select("id", "realname_control_enabled").Order("id").Find(&carriers).Error; err != nil {
		d.logger.Warn("读取运营商账户失败，沿用上一次列表", zap.Error(err))
		return
	}
	var rows []model.CarrierPollingParam
	if err := d.db.WithContext(ctx).Find(&rows).Error; err != nil {
		d.logger.Warn("读取运营商账户调度参数失败，沿用上一次参数", zap.Error(err))
		return
	}
	entries := make([]carrierEntry, 0, len(carriers))
	for _, carrier := range carriers {
		entries = append(entries, carrierEntry{id: carrier.ID, realnameControl: carrier.RealnameControlEnabled == 1})
	}
	params := make(map[uint]model.CarrierPollingParam, len(rows))
	for _, row := range rows {
		params[row.CarrierID] = row
	}
	d.mu.Lock()
	d.carriers, d.params, d.loadedAt = entries, params, time.Now()
	d.mu.Unlock()
}

// CarrierIDs 返回当前缓存的运营商账户 ID（消费者据此订阅）。
func (d *Dispatcher) CarrierIDs(ctx context.Context) []uint {
	d.refreshCarriers(ctx)
	d.mu.Lock()
	defer d.mu.Unlock()
	ids := make([]uint, 0, len(d.carriers))
	for _, carrier := range d.carriers {
		ids = append(ids, carrier.id)
	}
	return ids
}

// ParamsFor 返回账户生效调度参数（未配置项用全局默认）。
func (d *Dispatcher) ParamsFor(carrierID uint) CarrierParams {
	d.mu.Lock()
	row, ok := d.params[carrierID]
	d.mu.Unlock()
	return MergeParams(row, ok, d.cfg)
}

// MergeParams 合并账户配置行与全局默认。
func MergeParams(row model.CarrierPollingParam, configured bool, cfg config.CarrierSchedulerConfig) CarrierParams {
	params := CarrierParams{InflightLimit: cfg.DefaultInflightLimit, CardMaxInFlight: cfg.DefaultCardMaxInFlight, Configured: configured,
		SRInflightLimit: constants.DefaultStopResumeInflightLimit, SRMaxInFlight: constants.DefaultStopResumeMaxInFlight,
		SRMinInterval:  constants.DefaultStopResumeMinIntervalSeconds * time.Second,
		SRConfirmDelay: constants.DefaultStopResumeConfirmDelaySeconds * time.Second}
	if params.InflightLimit <= 0 {
		params.InflightLimit = 20
	}
	if params.CardMaxInFlight <= 0 {
		params.CardMaxInFlight = 20
	}
	if !configured {
		return params
	}
	if row.InflightLimit != nil && *row.InflightLimit > 0 {
		params.InflightLimit = *row.InflightLimit
	}
	if row.CardMaxInFlight != nil && *row.CardMaxInFlight > 0 {
		params.CardMaxInFlight = *row.CardMaxInFlight
	}
	if row.PriorityMaxRatio != nil && *row.PriorityMaxRatio > 0 {
		params.PriorityRatio = *row.PriorityMaxRatio
	}
	if row.SRInflightLimit != nil && *row.SRInflightLimit > 0 {
		params.SRInflightLimit = *row.SRInflightLimit
	}
	if row.StopResumeMaxInFlight != nil && *row.StopResumeMaxInFlight > 0 {
		params.SRMaxInFlight = *row.StopResumeMaxInFlight
	}
	if row.StopResumeMinInterval != nil && *row.StopResumeMinInterval >= 0 {
		params.SRMinInterval = time.Duration(*row.StopResumeMinInterval) * time.Second
	}
	if row.StopResumeConfirmDelay != nil && *row.StopResumeConfirmDelay > 0 {
		params.SRConfirmDelay = time.Duration(*row.StopResumeConfirmDelay) * time.Second
	}
	return params
}

// drainLegacyHints 消费旧优先提示通道与手动触发队列：优先提示放入账户优先队列，手动触发在普通队列中立即到期。
func (d *Dispatcher) drainLegacyHints(ctx context.Context) {
	for _, checkType := range constants.AllCheckTypes() {
		taskType := constants.TaskTypeFromCheckType(checkType)
		d.drainList(ctx, constants.RedisPollingPriorityQueueKey(taskType), checkType, false)
		d.drainList(ctx, constants.RedisPollingManualQueueKey(taskType), checkType, true)
	}
}

func (d *Dispatcher) drainList(ctx context.Context, key, checkType string, normal bool) {
	for i := 0; i < 500; i++ {
		raw, err := d.store.redis.LPop(ctx, key).Result()
		if err != nil {
			if !stderrors.Is(err, redis.Nil) {
				d.logger.Warn("读取旧提示队列失败", zap.String("key", key), zap.Error(err))
			}
			return
		}
		cardID, convErr := strconv.ParseUint(raw, 10, 64)
		if convErr != nil {
			continue
		}
		due, dueErr := d.store.ReadDue(ctx, uint(cardID))
		if dueErr != nil || due.CarrierID == 0 {
			continue
		}
		if _, err := d.store.RequestCheck(ctx, CheckRequest{CardID: uint(cardID), CarrierID: due.CarrierID, CheckTypes: []string{checkType}, Normal: normal}); err != nil {
			d.logger.Warn("写入账户队列失败", zap.Uint64("card_id", cardID), zap.Error(err))
		}
	}
}

// 新结构初始化状态（constants.RedisCarrierSchedulerInitStateKey）。
const (
	InitStateRunning = "running"
	InitStateDone    = "done"
)