feat(queue): 增加非文本模型分布式准入队列

使用 PostgreSQL 统一同步与异步非文本任务的并发准入、持久化等待和 Worker 容量分配，并将生产 API 与独立 Worker 角色拆分。补充策略管理、共享契约、OpenAPI、Kubernetes 双节点 Worker 清单及跨节点验收脚本；未默认启用任何生产 queue_size 策略。已在原基线完成 Go、前端、迁移、Shell、Kustomize 与长任务容量验收；合入最新主干后将重新执行发布门禁。
2026-07-29 16:15:43 +08:00
parent 2e5a90731b
commit 9e4fc7362d
55 changed files with 4565 additions and 237 deletions
@@ -102,7 +102,7 @@ func (s *Service) startRiverQueue(ctx context.Context, workerEnabled bool) error
 	}

 	controlClient, err := river.NewClient(driver, &river.Config{
-		ID:       asyncWorkerID() + "-control",
+		ID:       s.workerInstanceID + "-control",
 		Logger:   s.logger,
 		TestOnly: s.cfg.AppEnv == "test",
 	})
@@ -130,15 +130,18 @@ func (s *Service) startRiverQueue(ctx context.Context, workerEnabled bool) error
 	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
+	var executionClient asyncExecutionClient
+	if snapshot.Capacity > 0 {
+		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
@@ -161,11 +164,15 @@ func (s *Service) startRiverQueue(ctx context.Context, workerEnabled bool) error
 	)
 	if err := s.recoverAsyncRiverJobs(ctx); err != nil {
 		cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-		_ = executionClient.StopAndCancel(cleanupCtx)
+		if executionClient != nil {
+			_ = executionClient.StopAndCancel(cleanupCtx)
+		}
 		cancel()
 		return err
 	}
 	go s.refreshAsyncWorkerCapacity(ctx)
+	go s.dispatchWaitingAsyncAdmissions(ctx)
+	go s.reapExpiredTaskAdmissions(ctx)
 	go s.stopAsyncWorkersOnShutdown(ctx)
 	return nil
 }
@@ -176,7 +183,7 @@ func (s *Service) newRiverAsyncExecutionClient(capacity int) (*river.Client[pgx.
 		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()),
+		ID:                          fmt.Sprintf("%s-exec-%d-%d", s.workerInstanceID, capacity, time.Now().UnixNano()),
 		JobTimeout:                  -1,
 		Logger:                      s.logger,
 		CompletedJobRetentionPeriod: 24 * time.Hour,
@@ -203,7 +210,26 @@ func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorke
 	if s.asyncCapacityLoader != nil {
 		return s.asyncCapacityLoader(ctx, s.cfg.AsyncWorkerHardLimit)
 	}
-	return s.store.AsyncWorkerCapacity(ctx, s.cfg.AsyncWorkerHardLimit)
+	snapshot, err := s.store.AsyncWorkerCapacity(ctx, s.cfg.AsyncWorkerHardLimit)
+	if err != nil {
+		return store.AsyncWorkerCapacitySnapshot{}, err
+	}
+	allocation, err := s.store.RegisterWorkerInstance(ctx, store.WorkerRegistrationInput{
+		InstanceID:      s.workerInstanceID,
+		PodUID:          strings.TrimSpace(os.Getenv("POD_UID")),
+		PodName:         strings.TrimSpace(os.Getenv("POD_NAME")),
+		Site:            strings.TrimSpace(os.Getenv("EASYAI_SITE")),
+		Revision:        strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
+		DesiredCapacity: snapshot.Capacity,
+	})
+	if err != nil {
+		return store.AsyncWorkerCapacitySnapshot{}, err
+	}
+	snapshot.Capacity = allocation.Allocated
+	snapshot.GlobalCapacity = allocation.DesiredCapacity
+	snapshot.ActiveInstances = allocation.ActiveInstances
+	snapshot.InstanceID = allocation.InstanceID
+	return snapshot, nil
 }

 func (s *Service) refreshAsyncWorkerCapacity(ctx context.Context) {
@@ -233,25 +259,30 @@ func (s *Service) resizeAsyncWorkerCapacity(ctx context.Context) {
 	if snapshot.Capacity == currentCapacity {
 		return
 	}
-	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
-	}
-	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
+	var newClient asyncExecutionClient
+	if snapshot.Capacity > 0 {
+		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
+		}
+		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)
+		if newClient != nil {
+			_ = newClient.StopAndCancel(cleanupCtx)
+		}
 		cancel()
 		return
 	}
@@ -271,6 +302,8 @@ func (s *Service) resizeAsyncWorkerCapacity(ctx context.Context) {
 		"previousCapacity", currentCapacity,
 		"capacity", snapshot.Capacity,
 		"desiredCapacity", snapshot.Desired,
+		"activeInstances", snapshot.ActiveInstances,
+		"instanceID", snapshot.InstanceID,
 		"hardLimit", snapshot.HardLimit,
 		"modelDesired", snapshot.ModelDesired,
 		"groupDesired", snapshot.GroupDesired,
@@ -294,6 +327,13 @@ func (s *Service) drainAsyncWorkerClient(client asyncExecutionClient) {

 func (s *Service) stopAsyncWorkersOnShutdown(ctx context.Context) {
 	<-ctx.Done()
+	if s.store != nil && strings.TrimSpace(s.workerInstanceID) != "" {
+		drainCtx, drainCancel := context.WithTimeout(context.Background(), 5*time.Second)
+		if err := s.store.MarkWorkerDraining(drainCtx, s.workerInstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) {
+			s.logger.Warn("mark async worker draining failed", "error", err)
+		}
+		drainCancel()
+	}
 	s.riverMu.Lock()
 	s.riverControlClient = nil
 	clients := make([]asyncExecutionClient, 0, 1+len(s.riverDrainingClients))
@@ -328,6 +368,12 @@ func (s *Service) observeAsyncWorkerCapacity(snapshot store.AsyncWorkerCapacityS
 	if ok {
 		observer.SetAsyncWorkerCapacity(snapshot.Capacity, snapshot.Desired, snapshot.HardLimit, snapshot.Capped)
 	}
+	distributedObserver, ok := s.billingMetrics.(interface {
+		SetDistributedWorkerCapacity(activeInstances, globalCapacity, allocatedCapacity int)
+	})
+	if ok {
+		distributedObserver.SetDistributedWorkerCapacity(snapshot.ActiveInstances, snapshot.GlobalCapacity, snapshot.Capacity)
+	}
 }

 func (s *Service) observeAsyncWorkerResize(outcome string) {
@@ -340,16 +386,38 @@ func (s *Service) observeAsyncWorkerResize(outcome string) {
 }

 func (s *Service) EnqueueAsyncTask(ctx context.Context, task store.GatewayTask) error {
+	return s.enqueueAsyncTaskWithOptions(ctx, task.ID, asyncTaskInsertOpts(task))
+}
+
+func (s *Service) enqueueAsyncTaskWithOptions(ctx context.Context, taskID string, opts *river.InsertOpts) error {
+	tx, err := s.store.Pool().Begin(ctx)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback(ctx)
+	if err := s.enqueueAsyncTaskTx(ctx, tx, taskID, opts); err != nil {
+		return err
+	}
+	return tx.Commit(ctx)
+}
+
+func (s *Service) enqueueAsyncTaskTx(ctx context.Context, tx pgx.Tx, taskID string, opts *river.InsertOpts) error {
 	riverClient := s.asyncControlClient()
 	if riverClient == nil {
 		return errors.New("river async queue is not started")
 	}
-	result, err := riverClient.Insert(ctx, asyncTaskArgs{TaskID: task.ID}, asyncTaskInsertOpts(task))
+	result, err := riverClient.InsertTx(ctx, tx, asyncTaskArgs{TaskID: taskID}, opts)
 	if err != nil {
 		return err
 	}
 	if result.Job != nil {
-		return s.store.SetTaskRiverJobID(ctx, task.ID, result.Job.ID)
+		if _, err := tx.Exec(ctx, `
+UPDATE gateway_tasks
+SET river_job_id = $2,
+    updated_at = now()
+WHERE id = $1::uuid`, taskID, result.Job.ID); err != nil {
+			return err
+		}
 	}
 	return nil
 }
@@ -373,14 +441,13 @@ func (s *Service) recoverAsyncRiverJobs(ctx context.Context) error {
 		return err
 	}
 	for _, item := range items {
-		result, err := riverClient.Insert(ctx, asyncTaskArgs{TaskID: item.ID}, asyncTaskRecoveryInsertOpts(item, time.Now()))
-		if err != nil {
-			return err
+		if admission, admissionErr := s.store.GetTaskAdmission(ctx, item.ID); admissionErr == nil && admission.Status == "waiting" {
+			continue
+		} else if admissionErr != nil && !errors.Is(admissionErr, pgx.ErrNoRows) {
+			return admissionErr
 		}
-		if result.Job != nil {
-			if err := s.store.SetTaskRiverJobID(ctx, item.ID, result.Job.ID); err != nil {
-				return err
-			}
+		if err := s.enqueueAsyncTaskWithOptions(ctx, item.ID, asyncTaskRecoveryInsertOpts(item, time.Now())); err != nil {
+			return err
 		}
 	}
 	if len(items) > 0 {
@@ -418,10 +485,6 @@ func asyncTaskRecoveryInsertOpts(item store.AsyncTaskQueueItem, now time.Time) *
 	if item.NextRunAt.After(now) {
 		opts.ScheduledAt = item.NextRunAt
 	}
-	// A replacement process must not be blocked by a River row that the dead
-	// process left in running state. PostgreSQL execution leases still ensure
-	// that only one recovery job can call the upstream provider.
-	opts.UniqueOpts = river.UniqueOpts{}
 	return opts
 }