perf(queue): 按策略动态扩缩异步 Worker

将 River 执行容量按平台模型与用户组的有效并发策略动态调整，并为长任务续租、并发租约原子抢占和限流退避补充保护。\n\n统一平台模型限流继承语义，兼容历史 platformLimits/modelLimits，并为三个迁移图像模型建立独立并发租约。补充管理端显式继承/覆盖配置、指标、单元测试及隔离 PostgreSQL 验收。\n\n验证：go test ./... -count=1；pnpm lint；pnpm test；pnpm build；隔离 PostgreSQL 并发原子性/续租测试；128 任务与三分钟长任务动态 Worker 验收。
2026-07-24 12:26:56 +08:00
parent 290b8c1854
commit 6c5daf29ca
37 changed files with 2365 additions and 205 deletions
@@ -11,6 +11,7 @@ import (
 	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
 	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
 	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
 	"github.com/riverqueue/river"
 	"github.com/riverqueue/river/riverdriver/riverpgxv5"
 	"github.com/riverqueue/river/rivermigrate"
@@ -23,6 +24,12 @@ type asyncTaskArgs struct {
 	TaskID string `json:"task_id" river:"unique"`
 }

+type asyncExecutionClient interface {
+	Start(context.Context) error
+	Stop(context.Context) error
+	StopAndCancel(context.Context) error
+}
+
 func (asyncTaskArgs) Kind() string { return "gateway_task_run" }

 type asyncTaskWorker struct {
@@ -87,21 +94,70 @@ func (s *Service) startRiverQueue(ctx context.Context) error {
 		return err
 	}

-	workers := river.NewWorkers()
-	if err := river.AddWorkerSafely(workers, &asyncTaskWorker{service: s}); err != nil {
+	controlClient, err := river.NewClient(driver, &river.Config{
+		ID:       asyncWorkerID() + "-control",
+		Logger:   s.logger,
+		TestOnly: s.cfg.AppEnv == "test",
+	})
+	if err != nil {
 		return err
 	}
-	riverClient, err := river.NewClient(driver, &river.Config{
-		ID:                          asyncWorkerID(),
+	snapshot, err := s.loadAsyncWorkerCapacity(ctx)
+	if err != nil {
+		return fmt.Errorf("calculate initial async worker capacity: %w", err)
+	}
+	executionClient, err := s.makeAsyncExecutionClient(snapshot.Capacity)
+	if err != nil {
+		return err
+	}
+	if err := executionClient.Start(ctx); err != nil {
+		cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		_ = executionClient.StopAndCancel(cleanupCtx)
+		cancel()
+		return err
+	}
+	s.riverMu.Lock()
+	s.riverControlClient = controlClient
+	s.riverExecutionClient = executionClient
+	s.riverWorkerCapacity = snapshot.Capacity
+	s.riverDrainingClients = make(map[asyncExecutionClient]struct{})
+	s.riverMu.Unlock()
+	s.observeAsyncWorkerCapacity(snapshot)
+	s.logger.Info("async worker capacity initialized",
+		"capacity", snapshot.Capacity,
+		"desiredCapacity", snapshot.Desired,
+		"hardLimit", snapshot.HardLimit,
+		"enabledModels", snapshot.EnabledModels,
+		"unlimitedModels", snapshot.UnlimitedModels,
+		"modelDesired", snapshot.ModelDesired,
+		"enabledGroups", snapshot.EnabledGroups,
+		"unlimitedGroups", snapshot.UnlimitedGroups,
+		"groupDesired", snapshot.GroupDesired,
+		"capped", snapshot.Capped,
+	)
+	if err := s.recoverAsyncRiverJobs(ctx); err != nil {
+		cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		_ = executionClient.StopAndCancel(cleanupCtx)
+		cancel()
+		return err
+	}
+	go s.refreshAsyncWorkerCapacity(ctx)
+	go s.stopAsyncWorkersOnShutdown(ctx)
+	return nil
+}
+
+func (s *Service) newRiverAsyncExecutionClient(capacity int) (*river.Client[pgx.Tx], error) {
+	workers := river.NewWorkers()
+	if err := river.AddWorkerSafely(workers, &asyncTaskWorker{service: s}); err != nil {
+		return nil, err
+	}
+	return river.NewClient(riverpgxv5.New(s.store.Pool()), &river.Config{
+		ID:                          fmt.Sprintf("%s-exec-%d-%d", asyncWorkerID(), capacity, time.Now().UnixNano()),
 		JobTimeout:                  -1,
 		Logger:                      s.logger,
 		CompletedJobRetentionPeriod: 24 * time.Hour,
 		Queues: map[string]river.QueueConfig{
-			// Image providers may hold a worker while polling for several
-			// minutes. Keep enough workers available for production bursts so
-			// unrelated models do not remain queued behind long-running media
-			// tasks.
-			asyncTaskQueueName: {MaxWorkers: 96},
+			asyncTaskQueueName: {MaxWorkers: capacity},
 		},
 		// Provider-backed media jobs commonly poll for 10-20 minutes. River may
 		// execute a still-running job again once this window elapses, so keep the
@@ -110,32 +166,160 @@ func (s *Service) startRiverQueue(ctx context.Context) error {
 		TestOnly:             s.cfg.AppEnv == "test",
 		Workers:              workers,
 	})
+}
+
+func (s *Service) makeAsyncExecutionClient(capacity int) (asyncExecutionClient, error) {
+	if s.asyncClientFactory != nil {
+		return s.asyncClientFactory(capacity)
+	}
+	return s.newRiverAsyncExecutionClient(capacity)
+}
+
+func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorkerCapacitySnapshot, error) {
+	if s.asyncCapacityLoader != nil {
+		return s.asyncCapacityLoader(ctx, s.cfg.AsyncWorkerHardLimit)
+	}
+	return s.store.AsyncWorkerCapacity(ctx, s.cfg.AsyncWorkerHardLimit)
+}
+
+func (s *Service) refreshAsyncWorkerCapacity(ctx context.Context) {
+	ticker := time.NewTicker(time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * time.Second)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			s.resizeAsyncWorkerCapacity(ctx)
+		}
+	}
+}
+
+func (s *Service) resizeAsyncWorkerCapacity(ctx context.Context) {
+	snapshot, err := s.loadAsyncWorkerCapacity(ctx)
 	if err != nil {
-		return err
+		s.observeAsyncWorkerResize("refresh_failed")
+		s.logger.Warn("refresh async worker capacity failed; keeping current client", "error", err)
+		return
 	}
-	s.riverClient = riverClient
-	if err := riverClient.Start(ctx); err != nil {
-		return err
+	s.riverMu.RLock()
+	currentCapacity := s.riverWorkerCapacity
+	s.riverMu.RUnlock()
+	s.observeAsyncWorkerCapacity(snapshot)
+	if snapshot.Capacity == currentCapacity {
+		return
 	}
-	if err := s.recoverAsyncRiverJobs(ctx); err != nil {
-		return err
+	newClient, err := s.makeAsyncExecutionClient(snapshot.Capacity)
+	if err != nil {
+		s.observeAsyncWorkerResize("create_failed")
+		s.logger.Warn("create replacement async worker client failed; keeping current client",
+			"error", err, "currentCapacity", currentCapacity, "desiredCapacity", snapshot.Capacity)
+		return
 	}
-	go func() {
-		<-ctx.Done()
+	if err := newClient.Start(ctx); err != nil {
+		cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		_ = newClient.StopAndCancel(cleanupCtx)
+		cancel()
+		s.observeAsyncWorkerResize("start_failed")
+		s.logger.Warn("start replacement async worker client failed; keeping current client",
+			"error", err, "currentCapacity", currentCapacity, "desiredCapacity", snapshot.Capacity)
+		return
+	}
+	if ctx.Err() != nil {
+		cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		_ = newClient.StopAndCancel(cleanupCtx)
+		cancel()
+		return
+	}
+	s.riverMu.Lock()
+	oldClient := s.riverExecutionClient
+	s.riverExecutionClient = newClient
+	s.riverWorkerCapacity = snapshot.Capacity
+	if oldClient != nil {
+		if s.riverDrainingClients == nil {
+			s.riverDrainingClients = make(map[asyncExecutionClient]struct{})
+		}
+		s.riverDrainingClients[oldClient] = struct{}{}
+	}
+	s.riverMu.Unlock()
+	s.observeAsyncWorkerResize("success")
+	s.logger.Info("async worker capacity resized",
+		"previousCapacity", currentCapacity,
+		"capacity", snapshot.Capacity,
+		"desiredCapacity", snapshot.Desired,
+		"hardLimit", snapshot.HardLimit,
+		"modelDesired", snapshot.ModelDesired,
+		"groupDesired", snapshot.GroupDesired,
+		"capped", snapshot.Capped,
+	)
+	if oldClient != nil {
+		go s.drainAsyncWorkerClient(oldClient)
+	}
+}
+
+func (s *Service) drainAsyncWorkerClient(client asyncExecutionClient) {
+	stopCtx, cancel := context.WithTimeout(context.Background(), time.Hour)
+	defer cancel()
+	if err := client.Stop(stopCtx); err != nil {
+		s.logger.Warn("gracefully drain previous async worker client failed", "error", err)
+	}
+	s.riverMu.Lock()
+	delete(s.riverDrainingClients, client)
+	s.riverMu.Unlock()
+}
+
+func (s *Service) stopAsyncWorkersOnShutdown(ctx context.Context) {
+	<-ctx.Done()
+	s.riverMu.Lock()
+	clients := make([]asyncExecutionClient, 0, 1+len(s.riverDrainingClients))
+	if s.riverExecutionClient != nil {
+		clients = append(clients, s.riverExecutionClient)
+	}
+	for client := range s.riverDrainingClients {
+		clients = append(clients, client)
+	}
+	s.riverExecutionClient = nil
+	s.riverDrainingClients = nil
+	s.riverMu.Unlock()
+	for _, client := range clients {
 		stopCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-		defer cancel()
-		if err := riverClient.StopAndCancel(stopCtx); err != nil {
+		if err := client.StopAndCancel(stopCtx); err != nil {
 			s.logger.Warn("stop river async queue failed", "error", err)
 		}
-	}()
-	return nil
+		cancel()
+	}
+}
+
+func (s *Service) asyncControlClient() *river.Client[pgx.Tx] {
+	s.riverMu.RLock()
+	defer s.riverMu.RUnlock()
+	return s.riverControlClient
+}
+
+func (s *Service) observeAsyncWorkerCapacity(snapshot store.AsyncWorkerCapacitySnapshot) {
+	observer, ok := s.billingMetrics.(interface {
+		SetAsyncWorkerCapacity(current, desired, hardLimit int, capped bool)
+	})
+	if ok {
+		observer.SetAsyncWorkerCapacity(snapshot.Capacity, snapshot.Desired, snapshot.HardLimit, snapshot.Capped)
+	}
+}
+
+func (s *Service) observeAsyncWorkerResize(outcome string) {
+	observer, ok := s.billingMetrics.(interface {
+		ObserveAsyncWorkerResize(string)
+	})
+	if ok {
+		observer.ObserveAsyncWorkerResize(outcome)
+	}
 }

 func (s *Service) EnqueueAsyncTask(ctx context.Context, task store.GatewayTask) error {
-	if s.riverClient == nil {
+	riverClient := s.asyncControlClient()
+	if riverClient == nil {
 		return errors.New("river async queue is not started")
 	}
-	result, err := s.riverClient.Insert(ctx, asyncTaskArgs{TaskID: task.ID}, asyncTaskInsertOpts(task))
+	result, err := riverClient.Insert(ctx, asyncTaskArgs{TaskID: task.ID}, asyncTaskInsertOpts(task))
 	if err != nil {
 		return err
 	}
@@ -155,12 +339,16 @@ func (s *Service) RunAsyncTask(ctx context.Context, task store.GatewayTask, user
 }

 func (s *Service) recoverAsyncRiverJobs(ctx context.Context) error {
+	riverClient := s.asyncControlClient()
+	if riverClient == nil {
+		return errors.New("river async queue is not started")
+	}
 	items, err := s.store.ListRecoverableAsyncTasks(ctx, 1000)
 	if err != nil {
 		return err
 	}
 	for _, item := range items {
-		result, err := s.riverClient.Insert(ctx, asyncTaskArgs{TaskID: item.ID}, asyncTaskRecoveryInsertOpts(item, time.Now()))
+		result, err := riverClient.Insert(ctx, asyncTaskArgs{TaskID: item.ID}, asyncTaskRecoveryInsertOpts(item, time.Now()))
 		if err != nil {
 			return err
 		}