fix(queue): 自动救援失活 Worker 遗留任务

新增基于任务状态、执行租约和 Worker 心跳的孤儿 River Job 精确识别，每 15 秒将确认失活的 Job 恢复为可重试，避免等待 River 一小时通用救援窗口。运行时恢复同步清除中断及终态任务的 admission、执行令牌和并发租约，防止状态残留污染队列指标。正常运行中的长任务仍由 running 状态和活跃心跳保护。验证通过完整 Go 测试、go vet、竞态检查、迁移安全检查和 PostgreSQL 18 集成测试。
2026-07-29 23:44:29 +08:00
parent 98820378b7
commit f3dd7cd262
6 changed files with 369 additions and 3 deletions
@@ -63,7 +63,7 @@ func TestAsyncQueueClientEnqueuesWithoutExecutionWorker(t *testing.T) {

 	var attemptedByCount int
 	if err := db.Pool().QueryRow(ctx, `
-SELECT cardinality(attempted_by)
+SELECT COALESCE(cardinality(attempted_by), 0)
 FROM river_job
 WHERE id = $1`, queued.RiverJobID).Scan(&attemptedByCount); err != nil {
 		t.Fatalf("load River job: %v", err)
@@ -72,3 +72,110 @@ WHERE id = $1`, queued.RiverJobID).Scan(&attemptedByCount); err != nil {
 		t.Fatalf("control-only queue client executed the job: attempted_by=%d", attemptedByCount)
 	}
 }
+
+func TestRecoverOrphanedAsyncRiverJobRequiresInactiveWorker(t *testing.T) {
+	databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
+	if databaseURL == "" {
+		t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run the orphaned River job integration test")
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	db, err := store.Connect(ctx, databaseURL)
+	if err != nil {
+		t.Fatalf("connect store: %v", err)
+	}
+	t.Cleanup(db.Close)
+	service := New(config.Config{AppEnv: "test"}, db, slog.New(slog.NewTextHandler(io.Discard, nil)))
+	service.StartAsyncQueueClient(ctx)
+
+	suffix := strconv.FormatInt(time.Now().UnixNano(), 10)
+	createQueuedTask := func(label string) store.GatewayTask {
+		t.Helper()
+		task, createErr := db.CreateTask(ctx, store.CreateTaskInput{
+			Kind:    "images.edits",
+			Model:   "orphan-recovery-" + label + "-" + suffix,
+			Request: map[string]any{"prompt": "recover orphaned River job"},
+			Async:   true,
+			RunMode: "simulation",
+		}, &auth.User{ID: "orphan-recovery-" + suffix, Source: "gateway"})
+		if createErr != nil {
+			t.Fatalf("create %s task: %v", label, createErr)
+		}
+		if enqueueErr := service.EnqueueAsyncTask(ctx, task); enqueueErr != nil {
+			t.Fatalf("enqueue %s task: %v", label, enqueueErr)
+		}
+		queued, getErr := db.GetTask(ctx, task.ID)
+		if getErr != nil {
+			t.Fatalf("load %s task: %v", label, getErr)
+		}
+		return queued
+	}
+
+	orphaned := createQueuedTask("orphaned")
+	protected := createQueuedTask("protected")
+	staleWorkerID := "orphan-recovery-stale-" + suffix
+	activeWorkerID := "orphan-recovery-active-" + suffix
+	t.Cleanup(func() {
+		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cleanupCancel()
+		_, _ = db.Pool().Exec(cleanupCtx, `DELETE FROM gateway_worker_instances WHERE instance_id = ANY($1::text[])`, []string{staleWorkerID, activeWorkerID})
+		_, _ = db.Pool().Exec(cleanupCtx, `DELETE FROM river_job WHERE id = ANY($1::bigint[])`, []int64{orphaned.RiverJobID, protected.RiverJobID})
+		_, _ = db.Pool().Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = ANY($1::uuid[])`, []string{orphaned.ID, protected.ID})
+	})
+
+	if _, err := db.Pool().Exec(ctx, `
+INSERT INTO gateway_worker_instances (
+  instance_id, status, desired_capacity, allocated_capacity, heartbeat_at, updated_at
+)
+VALUES
+  ($1, 'active', 1, 0, now() - interval '2 minutes', now() - interval '2 minutes'),
+  ($2, 'active', 1, 1, now(), now())`,
+		staleWorkerID,
+		activeWorkerID,
+	); err != nil {
+		t.Fatalf("seed worker heartbeats: %v", err)
+	}
+	if _, err := db.Pool().Exec(ctx, `
+UPDATE river_job
+SET state = 'running',
+    attempt = 1,
+    attempted_at = now() - interval '1 minute',
+    attempted_by = ARRAY[
+      CASE id
+        WHEN $1::bigint THEN $3
+        WHEN $2::bigint THEN $4
+      END
+    ]::text[]
+WHERE id = ANY($5::bigint[])`,
+		orphaned.RiverJobID,
+		protected.RiverJobID,
+		staleWorkerID+"-exec-1-test",
+		activeWorkerID+"-exec-1-test",
+		[]int64{orphaned.RiverJobID, protected.RiverJobID},
+	); err != nil {
+		t.Fatalf("mark River jobs running: %v", err)
+	}
+
+	recovered, err := db.RecoverOrphanedAsyncRiverJobs(ctx, 30*time.Second, 10)
+	if err != nil {
+		t.Fatalf("recover orphaned River jobs: %v", err)
+	}
+	if recovered != 1 {
+		t.Fatalf("recovered jobs=%d, want 1", recovered)
+	}
+	var orphanedState, protectedState string
+	if err := db.Pool().QueryRow(ctx, `SELECT state::text FROM river_job WHERE id = $1`, orphaned.RiverJobID).Scan(&orphanedState); err != nil {
+		t.Fatalf("read orphaned River job state: %v", err)
+	}
+	if err := db.Pool().QueryRow(ctx, `SELECT state::text FROM river_job WHERE id = $1`, protected.RiverJobID).Scan(&protectedState); err != nil {
+		t.Fatalf("read protected River job state: %v", err)
+	}
+	if orphanedState != "retryable" || protectedState != "running" {
+		t.Fatalf("River states orphaned=%s protected=%s, want retryable/running", orphanedState, protectedState)
+	}
+	recovered, err = db.RecoverOrphanedAsyncRiverJobs(ctx, 30*time.Second, 10)
+	if err != nil || recovered != 0 {
+		t.Fatalf("second orphan recovery count=%d err=%v, want idempotent zero", recovered, err)
+	}
+}
@@ -18,7 +18,12 @@ import (
 	"github.com/riverqueue/river/rivertype"
 )

-const asyncTaskQueueName = "gateway_tasks"
+const (
+	asyncTaskQueueName                = "gateway_tasks"
+	orphanedRiverJobScanInterval      = 15 * time.Second
+	orphanedRiverJobWorkerStaleAfter  = 45 * time.Second
+	orphanedRiverJobRecoveryBatchSize = 100
+)

 type asyncTaskArgs struct {
 	TaskID string `json:"task_id" river:"unique"`
@@ -173,6 +178,7 @@ func (s *Service) startRiverQueue(ctx context.Context, workerEnabled bool) error
 	go s.refreshAsyncWorkerCapacity(ctx)
 	go s.dispatchWaitingAsyncAdmissions(ctx)
 	go s.reapExpiredTaskAdmissions(ctx)
+	go s.recoverOrphanedAsyncRiverJobs(ctx)
 	go s.stopAsyncWorkersOnShutdown(ctx)
 	return nil
 }
@@ -457,6 +463,36 @@ func (s *Service) recoverAsyncRiverJobs(ctx context.Context) error {
 	return nil
 }

+func (s *Service) recoverOrphanedAsyncRiverJobs(ctx context.Context) {
+	recoverJobs := func() {
+		recovered, err := s.store.RecoverOrphanedAsyncRiverJobs(
+			ctx,
+			orphanedRiverJobWorkerStaleAfter,
+			orphanedRiverJobRecoveryBatchSize,
+		)
+		if err != nil {
+			if ctx.Err() == nil {
+				s.logger.Warn("recover orphaned river jobs failed", "error", err)
+			}
+			return
+		}
+		if recovered > 0 {
+			s.logger.Warn("orphaned river jobs recovered", "count", recovered)
+		}
+	}
+	recoverJobs()
+	ticker := time.NewTicker(orphanedRiverJobScanInterval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			recoverJobs()
+		}
+	}
+}
+
 func asyncTaskInsertOpts(task store.GatewayTask) *river.InsertOpts {
 	priority := 2
 	if task.ID == "" {