junhong_cmp_fiber/cmd/worker/carrier_scheduler.go

package main

import (
	"context"
	"encoding/json"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"go.uber.org/zap"

	cardObservationApp "github.com/break/junhong_cmp_fiber/internal/application/cardobservation"
	"github.com/break/junhong_cmp_fiber/internal/infrastructure/messaging/outbox"
	"github.com/break/junhong_cmp_fiber/internal/infrastructure/nsqbus"
	"github.com/break/junhong_cmp_fiber/internal/polling"
	"github.com/break/junhong_cmp_fiber/internal/polling/carrierqueue"
	"github.com/break/junhong_cmp_fiber/internal/task"
	"github.com/break/junhong_cmp_fiber/pkg/config"
	"github.com/break/junhong_cmp_fiber/pkg/constants"
	"github.com/break/junhong_cmp_fiber/pkg/queue"
)

// carrierInitLockKey 新结构初始化锁：多实例时只由一个实例初始化。
const carrierInitLockKey = "polling:carrier:init_lock"

// carrierSchedulerRuntime 是按运营商账户调度在 Worker 内的运行组件：模式监听、初始化、分发器与 NSQ 消费者。
type carrierSchedulerRuntime struct {
	cfg         config.CarrierSchedulerConfig
	runtime     *workerRuntime
	initializer *polling.CarrierInitializer
	dispatcher  *carrierqueue.Dispatcher
	consumers   *nsqbus.ConsumerManager
	producer    *nsqbus.Producer
	alerts      *carrierqueue.AlertWriter
	executor    *task.CarrierCardExecutor
	singleton   bool
	instanceID  string
	logger      *zap.Logger

	mu             sync.Mutex
	lastMode       string
	reinitQueued   bool
	legacyRestarts func()
}

// newCarrierSchedulerParts 在 Worker 运行时初始化阶段创建调度参数、账户队列存储与模式接缝，
// 并注入旧调度相关组件（生命周期、轮询基类），使其在启动前即可按模式写入正确结构。
func newCarrierSchedulerParts(runtime *workerRuntime, appLogger *zap.Logger) {
	settings := carrierqueue.NewDefaultSettingsProvider(runtime.db, runtime.redisClient, appLogger)
	store := carrierqueue.NewStore(runtime.redisClient)
	scheduling := polling.NewCarrierScheduling(store, settings, runtime.pollingConfigMgr, appLogger)
	runtime.carrierSettings = settings
	runtime.carrierStore = store
	runtime.carrierScheduling = scheduling
	runtime.lifecycleSvc.SetCarrierScheduling(scheduling)
	runtime.pollingBase.SetSchedulerModeReader(func(ctx context.Context) string { return settings.Current(ctx).Mode })
	if runtime.gatewayLimit != nil && runtime.gatewayLimit.Classifier != nil {
		classifier := runtime.gatewayLimit.Classifier
		runtime.pollingBase.SetErrorClassifier(func(ctx context.Context, op, message string) string {
			return classifier.Classify(ctx, op, message).Category
		})
	}
}

// startCarrierScheduler 启动模式监听；分发器只在 leader/all 运行（多实例安全），消费者在所有处理任务的实例运行。
func startCarrierScheduler(ctx context.Context, cfg *config.Config, runtime *workerRuntime, steps queue.PollingStepHandlers,
	singleton bool, legacyRestart func(), appLogger *zap.Logger) *carrierSchedulerRuntime {
	carrierCfg := cfg.Polling.Carrier
	instanceID := cfg.Worker.InstanceName
	if instanceID == "" {
		instanceID, _ = os.Hostname()
	}
	instanceID += "-" + strconv.Itoa(os.Getpid())
	alerts := carrierqueue.NewAlertWriter(runtime.db, runtime.redisClient, runtime.carrierSettings, appLogger)
	alerts.Start(ctx)
	executor := task.NewCarrierCardExecutor(runtime.pollingBase, runtime.carrierStore, runtime.carrierSettings, alerts,
		steps.Realname, steps.Carddata, steps.CardStatus, steps.Package, carrierCfg, instanceID)
	c := &carrierSchedulerRuntime{
		cfg: carrierCfg, runtime: runtime, singleton: singleton, instanceID: instanceID, logger: appLogger,
		alerts: alerts, executor: executor, legacyRestarts: legacyRestart,
		initializer: polling.NewCarrierInitializer(runtime.carrierScheduling, runtime.pollingIotCardStore, appLogger),
	}
	if carrierCfg.NSQDAddress != "" {
		producer, err := nsqbus.NewProducer(carrierCfg.NSQDAddress, appLogger)
		if err != nil {
			appLogger.Error("创建 NSQ 生产者失败，按运营商账户调度不可用", zap.Error(err))
		} else {
			c.producer = producer
		}
		c.consumers = nsqbus.NewConsumerManager(nsqbus.ConsumerConfig{
			NSQDAddress: carrierCfg.NSQDAddress, NSQLookupdAddresses: nsqbus.SplitAddresses(carrierCfg.NSQLookupdAddresses),
			MsgTimeout: carrierCfg.MsgTimeout, MaxAttempts: carrierCfg.MaxAttempts,
		}, constants.NSQChannelWorker, executor.Handle, appLogger)
	}
	var budgets carrierqueue.BudgetProvider
	if runtime.gatewayLimit != nil && runtime.gatewayLimit.Limiter != nil {
		budgets = runtime.gatewayLimit.Limiter
	}
	if c.producer != nil {
		c.dispatcher = carrierqueue.NewDispatcher(runtime.db, runtime.carrierStore, runtime.carrierSettings, c.producer, budgets,
			carrierCfg, instanceID, appLogger)
		if singleton {
			go c.dispatcher.Run(ctx)
		}
	}
	go c.watch(ctx)
	return c
}

// RequestReinit 轮询配置变更时请求按新配置重新初始化新结构（仅 carrier 模式生效）。
func (c *carrierSchedulerRuntime) RequestReinit() {
	if c == nil {
		return
	}
	c.mu.Lock()
	c.reinitQueued = true
	c.mu.Unlock()
}

func (c *carrierSchedulerRuntime) watch(ctx context.Context) {
	ticker := time.NewTicker(3 * time.Second)
	defer ticker.Stop()
	for {
		c.step(ctx)
		select {
		case <-ctx.Done():
			c.shutdown()
			return
		case <-ticker.C:
		}
	}
}

func (c *carrierSchedulerRuntime) step(ctx context.Context) {
	defer func() {
		if r := recover(); r != nil {
			c.logger.Error("按运营商账户调度模式监听 panic，已恢复", zap.Any("panic", r))
		}
	}()
	mode := c.runtime.carrierSettings.Current(ctx).Mode
	c.mu.Lock()
	previous := c.lastMode
	c.lastMode = mode
	reinit := c.reinitQueued
	c.reinitQueued = false
	c.mu.Unlock()
	if mode == constants.PollingSchedulerModeCarrier {
		if previous != mode {
			c.logger.Info("卡轮询调度模式：按运营商账户调度", zap.String("previous", previous))
		}
		if c.singleton {
			c.ensureInitialized(ctx, reinit)
		}
		c.syncConsumers(ctx)
		return
	}
	if previous == constants.PollingSchedulerModeCarrier {
		c.logger.Info("卡轮询调度模式回到旧分片调度：停止领取与消费，重建旧分片队列")
		if c.consumers != nil {
			c.consumers.StopAll(c.grace())
			c.consumers.Reset()
		}
		if c.singleton {
			c.runtime.redisClient.Del(ctx, constants.RedisCarrierSchedulerInitStateKey)
			if c.legacyRestarts != nil {
				c.legacyRestarts()
			}
		}
	}
}

// ensureInitialized 切到 carrier 后初始化新结构；初始化完成前分发器不领取。多实例时用初始化锁协调。
func (c *carrierSchedulerRuntime) ensureInitialized(ctx context.Context, reinit bool) {
	redisClient := c.runtime.redisClient
	state, _ := redisClient.Get(ctx, constants.RedisCarrierSchedulerInitStateKey).Result()
	if state == carrierqueue.InitStateDone && !reinit {
		return
	}
	ok, err := redisClient.SetNX(ctx, carrierInitLockKey, c.instanceID, 30*time.Minute).Result()
	if err != nil || !ok {
		return
	}
	defer redisClient.Del(context.WithoutCancel(ctx), carrierInitLockKey)
	if state != carrierqueue.InitStateDone {
		redisClient.Set(ctx, constants.RedisCarrierSchedulerInitStateKey, carrierqueue.InitStateRunning, 0)
	}
	started := time.Now()
	written, initErr := c.initializer.Run(ctx)
	if initErr != nil {
		c.logger.Error("按运营商账户调度初始化失败，下一轮重试", zap.Error(initErr))
		return
	}
	redisClient.Set(ctx, constants.RedisCarrierSchedulerInitStateKey, carrierqueue.InitStateDone, 0)
	c.logger.Info("按运营商账户调度初始化完成，开始领取", zap.Int64("cards", written), zap.Duration("duration", time.Since(started)),
		zap.String("metric", "polling.carrier.initialized"))
}

// syncConsumers 按当前运营商账户订阅卡任务 topic，max_in_flight 取账户调度参数。
func (c *carrierSchedulerRuntime) syncConsumers(ctx context.Context) {
	if c.consumers == nil || c.dispatcher == nil {
		return
	}
	desired := map[string]int{}
	for _, carrierID := range c.dispatcher.CarrierIDs(ctx) {
		desired[constants.NSQCardTopic(carrierID)] = c.dispatcher.ParamsFor(carrierID).CardMaxInFlight
	}
	c.consumers.Sync(desired)
}

func (c *carrierSchedulerRuntime) grace() time.Duration {
	if c.cfg.ShutdownGrace > 0 {
		return c.cfg.ShutdownGrace
	}
	return 30 * time.Second
}

// shutdown 优雅退出：停止订阅，宽限期内等待执行中的卡任务完成；未完成的由租约到期后重新领取。
func (c *carrierSchedulerRuntime) shutdown() {
	if c.consumers != nil {
		c.consumers.StopAll(c.grace())
	}
	if c.producer != nil {
		c.producer.Stop()
	}
}

// InitProgress 返回新结构初始化进度。
func (c *carrierSchedulerRuntime) InitProgress() polling.InitProgress {
	return c.initializer.Progress()
}

// carrierAwareSeriesConsumer 是业务事件观测序列 Outbox 消费者的模式接缝：按运营商账户调度时把业务事件改为按场景请求检查
// （事件在原业务事务中写入，消费即在事务提交之后），旧分片调度时沿用原观测序列。
type carrierAwareSeriesConsumer struct {
	legacy    outbox.EventConsumer
	requester *polling.CheckRequester
}

// Consume 消费业务观测序列请求事件。
func (c *carrierAwareSeriesConsumer) Consume(ctx context.Context, envelope outbox.DeliveryEnvelope) error {
	if c.requester == nil || !c.requester.Active(ctx) {
		return c.legacy.Consume(ctx, envelope)
	}
	var event cardObservationApp.SeriesRequestedEvent
	if err := json.Unmarshal(envelope.Payload, &event); err != nil {
		return nil
	}
	cardIDs := []uint{event.ResourceID}
	if event.ResourceType == constants.CardObservationResourceTypeDevice {
		cardIDs = event.ResourceIDs
	}
	c.requester.RequestCards(ctx, businessEventCheckScene(event), cardIDs, event.ExpectedValue)
	return nil
}

// businessEventCheckScene 把业务观测事件映射为请求检查场景：订单支付、套餐激活、自动续费购包按事件来源区分，停复机成功按观测场景。
func businessEventCheckScene(event cardObservationApp.SeriesRequestedEvent) string {
	switch {
	case strings.HasPrefix(event.EventID, "card-observation:order:"):
		return constants.PollingCheckSceneOrderPaid
	case strings.HasPrefix(event.EventID, "card-observation:auto-purchase:"), strings.HasPrefix(event.EventID, "asset-auto-renewal:"):
		return constants.PollingCheckSceneAutoRenewalPurchase
	case event.Scene == constants.CardObservationSceneBusinessStop || event.Scene == constants.CardObservationSceneBusinessResume:
		return constants.PollingCheckSceneStopResumeSucceeded
	default:
		return constants.PollingCheckScenePackageActivated
	}
}