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

backend/internal/service/batch_image_worker_runtime.go

@@ -0,0 +1,143 @@
+package service
+
+import (
+	"context"
+	"sync"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+)
+
+type BatchImageWorkerRuntime struct {
+	worker          *BatchImageWorker
+	billingRecovery *BatchImageBillingRecoveryService
+	cfg             *config.Config
+
+	mu     sync.Mutex
+	cancel context.CancelFunc
+	done   chan struct{}
+}
+
+func NewBatchImageWorkerRuntime(worker *BatchImageWorker, cfg *config.Config) *BatchImageWorkerRuntime {
+	return &BatchImageWorkerRuntime{worker: worker, cfg: cfg}
+}
+
+func ProvideBatchImageWorkerRuntime(
+	repo BatchImageRepository,
+	accountRepo AccountRepository,
+	queue BatchImageQueue,
+	billingRepo UsageBillingRepository,
+	usageLogRepo UsageLogRepository,
+	pricing *BatchImageModelPricingResolver,
+	authCache APIKeyAuthCacheInvalidator,
+	cfg *config.Config,
+) *BatchImageWorkerRuntime {
+	processor := &BatchImagePipelineProcessor{
+		ProviderProcessor: &BatchImageProviderProcessor{
+			Repo:             repo,
+			ProviderRegistry: NewBatchImageProviderRegistryFromConfig(cfg),
+			AccountResolver:  &BatchImageAccountRepositoryResolver{Repo: accountRepo},
+			BillingRepo:      billingRepo,
+			AuthCache:        authCache,
+		},
+		SettlementService: &BatchImageSettlementService{
+			Repo:         repo,
+			BillingRepo:  billingRepo,
+			UsageLogRepo: usageLogRepo,
+			Pricing:      pricing,
+			AuthCache:    authCache,
+			Config:       cfg,
+		},
+	}
+	runtime := NewBatchImageWorkerRuntime(NewBatchImageWorker(queue, processor, NewBatchImageWorkerOptionsFromConfig(cfg)), cfg)
+	runtime.billingRecovery = &BatchImageBillingRecoveryService{
+		Repo:       repo,
+		Billing:    billingRepo,
+		AuthCache:  authCache,
+		Queue:      queue,
+		StaleAfter: NewBatchImageWorkerOptionsFromConfig(cfg).StaleActiveAfter,
+		Limit:      NewBatchImageWorkerOptionsFromConfig(cfg).RecoverLimit,
+	}
+	runtime.Start()
+	return runtime
+}
+
+func (r *BatchImageWorkerRuntime) Start() {
+	if r == nil || r.worker == nil || r.cfg == nil || !r.cfg.BatchImage.QueueEnabled {
+		return
+	}
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if r.cancel != nil {
+		return
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	r.cancel = cancel
+	r.done = done
+
+	var wg sync.WaitGroup
+	wg.Add(4)
+	go func() {
+		defer wg.Done()
+		r.worker.Run(ctx)
+	}()
+	go func() {
+		defer wg.Done()
+		r.worker.RunDelayedMover(ctx)
+	}()
+	go func() {
+		defer wg.Done()
+		r.worker.RunStaleActiveRecovery(ctx)
+	}()
+	go func() {
+		defer wg.Done()
+		r.runBillingRecovery(ctx)
+	}()
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+}
+
+func (r *BatchImageWorkerRuntime) runBillingRecovery(ctx context.Context) {
+	if r == nil || r.worker == nil || r.billingRecovery == nil {
+		return
+	}
+	interval := r.worker.opts.RecoveryInterval
+	for {
+		if err := ctx.Err(); err != nil {
+			return
+		}
+		_, _ = r.billingRecovery.ReleaseStaleUnsubmittedOnce(ctx)
+		sleepOrDone(ctx, interval)
+	}
+}
+
+func (r *BatchImageWorkerRuntime) Stop() {
+	if r == nil {
+		return
+	}
+	r.mu.Lock()
+	cancel := r.cancel
+	done := r.done
+	r.cancel = nil
+	r.done = nil
+	r.mu.Unlock()
+
+	if cancel != nil {
+		cancel()
+	}
+	if done != nil {
+		<-done
+	}
+}
+
+func (r *BatchImageWorkerRuntime) Running() bool {
+	if r == nil {
+		return false
+	}
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return r.cancel != nil
+}