Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/batch_image_worker.go

@@ -0,0 +1,286 @@
+package service
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
+	"go.uber.org/zap"
+)
+
+const (
+	defaultBatchImageWorkerLockTTL             = 5 * time.Minute
+	defaultBatchImageWorkerLockConflictDelay   = 5 * time.Second
+	defaultBatchImageWorkerErrorRetryDelay     = time.Minute
+	defaultBatchImageWorkerRequeueDelay        = 30 * time.Second
+	defaultBatchImageWorkerDelayedPollInterval = 5 * time.Second
+	defaultBatchImageWorkerRecoveryInterval    = 5 * time.Minute
+	defaultBatchImageWorkerStaleActiveAfter    = 10 * time.Minute
+	defaultBatchImageWorkerDelayedMoveLimit    = 100
+	defaultBatchImageWorkerRecoverLimit        = 100
+	defaultBatchImageWorkerErrorBackoff        = time.Second
+	defaultBatchImageWorkerReserveBlockTimeout = 5 * time.Second
+)
+
+type BatchImageProcessor interface {
+	Process(ctx context.Context, batchID string) (BatchImageProcessResult, error)
+}
+
+type BatchImageProcessResult struct {
+	RequeueAfter time.Duration
+	Terminal     bool
+}
+
+type BatchImageWorkerOptions struct {
+	ReserveBlockTimeout time.Duration
+	JobLockTTL          time.Duration
+	LockConflictDelay   time.Duration
+	DefaultRequeueDelay time.Duration
+	ErrorRetryDelay     time.Duration
+	ErrorBackoff        time.Duration
+	DelayedPollInterval time.Duration
+	RecoveryInterval    time.Duration
+	StaleActiveAfter    time.Duration
+	DelayedMoveLimit    int
+	RecoverLimit        int
+}
+
+type BatchImageWorker struct {
+	queue     BatchImageQueue
+	processor BatchImageProcessor
+	opts      BatchImageWorkerOptions
+}
+
+func NewBatchImageWorker(queue BatchImageQueue, processor BatchImageProcessor, opts BatchImageWorkerOptions) *BatchImageWorker {
+	return &BatchImageWorker{
+		queue:     queue,
+		processor: processor,
+		opts:      normalizeBatchImageWorkerOptions(opts),
+	}
+}
+
+func NewBatchImageWorkerOptionsFromConfig(cfg *config.Config) BatchImageWorkerOptions {
+	if cfg == nil {
+		return normalizeBatchImageWorkerOptions(BatchImageWorkerOptions{})
+	}
+	return normalizeBatchImageWorkerOptions(BatchImageWorkerOptions{
+		JobLockTTL:          time.Duration(cfg.BatchImage.JobLockTTLSeconds) * time.Second,
+		LockConflictDelay:   time.Duration(cfg.BatchImage.LockConflictDelaySeconds) * time.Second,
+		DefaultRequeueDelay: time.Duration(cfg.BatchImage.DefaultRequeueDelaySeconds) * time.Second,
+		ErrorRetryDelay:     time.Duration(cfg.BatchImage.ErrorRetryDelaySeconds) * time.Second,
+		DelayedPollInterval: time.Duration(cfg.BatchImage.DelayedMoverIntervalSeconds) * time.Second,
+		RecoveryInterval:    time.Duration(cfg.BatchImage.RecoveryIntervalSeconds) * time.Second,
+		StaleActiveAfter:    time.Duration(cfg.BatchImage.StaleActiveAfterSeconds) * time.Second,
+		DelayedMoveLimit:    cfg.BatchImage.DelayedMoveLimit,
+		RecoverLimit:        cfg.BatchImage.RecoverLimit,
+	})
+}
+
+func normalizeBatchImageWorkerOptions(opts BatchImageWorkerOptions) BatchImageWorkerOptions {
+	if opts.ReserveBlockTimeout <= 0 {
+		opts.ReserveBlockTimeout = defaultBatchImageWorkerReserveBlockTimeout
+	}
+	if opts.JobLockTTL <= 0 {
+		opts.JobLockTTL = defaultBatchImageWorkerLockTTL
+	}
+	if opts.LockConflictDelay <= 0 {
+		opts.LockConflictDelay = defaultBatchImageWorkerLockConflictDelay
+	}
+	if opts.DefaultRequeueDelay <= 0 {
+		opts.DefaultRequeueDelay = defaultBatchImageWorkerRequeueDelay
+	}
+	if opts.ErrorRetryDelay <= 0 {
+		opts.ErrorRetryDelay = defaultBatchImageWorkerErrorRetryDelay
+	}
+	if opts.ErrorBackoff <= 0 {
+		opts.ErrorBackoff = defaultBatchImageWorkerErrorBackoff
+	}
+	if opts.DelayedPollInterval <= 0 {
+		opts.DelayedPollInterval = defaultBatchImageWorkerDelayedPollInterval
+	}
+	if opts.RecoveryInterval <= 0 {
+		opts.RecoveryInterval = defaultBatchImageWorkerRecoveryInterval
+	}
+	if opts.StaleActiveAfter <= 0 {
+		opts.StaleActiveAfter = defaultBatchImageWorkerStaleActiveAfter
+	}
+	if opts.DelayedMoveLimit <= 0 {
+		opts.DelayedMoveLimit = defaultBatchImageWorkerDelayedMoveLimit
+	}
+	if opts.RecoverLimit <= 0 {
+		opts.RecoverLimit = defaultBatchImageWorkerRecoverLimit
+	}
+	return opts
+}
+
+func (w *BatchImageWorker) Run(ctx context.Context) {
+	if w == nil {
+		return
+	}
+	for {
+		if err := ctx.Err(); err != nil {
+			return
+		}
+		if err := w.RunOnce(ctx); err != nil && ctx.Err() == nil {
+			sleepOrDone(ctx, w.opts.ErrorBackoff)
+		}
+	}
+}
+
+func (w *BatchImageWorker) RunOnce(ctx context.Context) error {
+	if w == nil || w.queue == nil || w.processor == nil {
+		return nil
+	}
+
+	reserved, err := w.queue.Reserve(ctx, w.opts.ReserveBlockTimeout)
+	if errors.Is(err, ErrBatchImageQueueEmpty) {
+		return nil
+	}
+	if err != nil {
+		return err
+	}
+
+	lock, ok, err := w.queue.TryAcquireJobLock(ctx, reserved.BatchID, w.opts.JobLockTTL)
+	if err != nil {
+		if requeueErr := w.queue.RequeueAfter(ctx, reserved.BatchID, w.opts.LockConflictDelay); requeueErr != nil {
+			return requeueErr
+		}
+		return err
+	}
+	if !ok {
+		// 锁被其他实例持有：按冲突延迟重新入队。直接丢弃会让 job 滞留在
+		// active zset，最早要等 StaleActiveAfter 才被恢复，造成分钟级停摆。
+		return w.queue.RequeueAfter(ctx, reserved.BatchID, w.opts.LockConflictDelay)
+	}
+	defer func() {
+		_ = lock.Release(ctx)
+	}()
+
+	// 处理期间持续心跳：刷新 active zset 时间戳防止 stale 恢复把在处理的
+	// job 重投给其他 worker，并对支持续期的锁实现延长锁 TTL。
+	hbStop := make(chan struct{})
+	hbDone := make(chan struct{})
+	go w.runJobHeartbeat(ctx, reserved.BatchID, lock, hbStop, hbDone)
+
+	result, err := w.processor.Process(ctx, reserved.BatchID)
+	close(hbStop)
+	<-hbDone
+	if err != nil {
+		logger.L().Warn("batch_image.worker_process_failed",
+			zap.String("batch_id", reserved.BatchID),
+			zap.Error(err),
+		)
+		return w.queue.RequeueAfter(ctx, reserved.BatchID, w.opts.ErrorRetryDelay)
+	}
+	if result.Terminal {
+		return w.queue.Ack(ctx, reserved.BatchID)
+	}
+	delay := result.RequeueAfter
+	if delay <= 0 {
+		delay = w.opts.DefaultRequeueDelay
+	}
+	return w.queue.RequeueAfter(ctx, reserved.BatchID, delay)
+}
+
+// BatchImageJobLockRefresher 是可选的锁续期能力；由具体锁实现按需提供。
+type BatchImageJobLockRefresher interface {
+	Refresh(ctx context.Context, ttl time.Duration) error
+}
+
+func (w *BatchImageWorker) heartbeatInterval() time.Duration {
+	interval := w.opts.JobLockTTL
+	if w.opts.StaleActiveAfter < interval {
+		interval = w.opts.StaleActiveAfter
+	}
+	interval /= 3
+	if interval < time.Second {
+		interval = time.Second
+	}
+	return interval
+}
+
+func (w *BatchImageWorker) runJobHeartbeat(ctx context.Context, batchID string, lock BatchImageJobLock, stop <-chan struct{}, done chan<- struct{}) {
+	defer close(done)
+	ticker := time.NewTicker(w.heartbeatInterval())
+	defer ticker.Stop()
+	for {
+		select {
+		case <-stop:
+			return
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			if err := w.queue.Heartbeat(ctx, batchID); err != nil && ctx.Err() == nil {
+				logger.L().Warn("batch_image.worker_heartbeat_failed",
+					zap.String("batch_id", batchID),
+					zap.Error(err),
+				)
+			}
+			if refresher, ok := lock.(BatchImageJobLockRefresher); ok {
+				if err := refresher.Refresh(ctx, w.opts.JobLockTTL); err != nil && ctx.Err() == nil {
+					logger.L().Warn("batch_image.worker_lock_refresh_failed",
+						zap.String("batch_id", batchID),
+						zap.Error(err),
+					)
+				}
+			}
+		}
+	}
+}
+
+func (w *BatchImageWorker) MoveDueDelayedOnce(ctx context.Context) (int, error) {
+	if w == nil || w.queue == nil {
+		return 0, nil
+	}
+	return w.queue.MoveDueDelayedToReady(ctx, w.opts.DelayedMoveLimit)
+}
+
+func (w *BatchImageWorker) RunDelayedMover(ctx context.Context) {
+	if w == nil {
+		return
+	}
+	for {
+		if err := ctx.Err(); err != nil {
+			return
+		}
+		moved, _ := w.MoveDueDelayedOnce(ctx)
+		if moved > 0 {
+			continue
+		}
+		sleepOrDone(ctx, w.opts.DelayedPollInterval)
+	}
+}
+
+func (w *BatchImageWorker) RecoverStaleActiveOnce(ctx context.Context) (int, error) {
+	if w == nil || w.queue == nil {
+		return 0, nil
+	}
+	return w.queue.RecoverStaleActive(ctx, w.opts.StaleActiveAfter, w.opts.RecoverLimit)
+}
+
+func (w *BatchImageWorker) RunStaleActiveRecovery(ctx context.Context) {
+	if w == nil {
+		return
+	}
+	for {
+		if err := ctx.Err(); err != nil {
+			return
+		}
+		_, _ = w.RecoverStaleActiveOnce(ctx)
+		sleepOrDone(ctx, w.opts.RecoveryInterval)
+	}
+}
+
+func sleepOrDone(ctx context.Context, d time.Duration) {
+	if d <= 0 {
+		return
+	}
+	timer := time.NewTimer(d)
+	defer timer.Stop()
+	select {
+	case <-ctx.Done():
+	case <-timer.C:
+	}
+}