From f3dd7cd2621f381bb6d0fadccd5e3d099e505262 Mon Sep 17 00:00:00 2001 From: wangbo Date: Wed, 29 Jul 2026 23:44:29 +0800 Subject: [PATCH] =?UTF-8?q?fix(queue):=20=E8=87=AA=E5=8A=A8=E6=95=91?= =?UTF-8?q?=E6=8F=B4=E5=A4=B1=E6=B4=BB=20Worker=20=E9=81=97=E7=95=99?= =?UTF-8?q?=E4=BB=BB=E5=8A=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增基于任务状态、执行租约和 Worker 心跳的孤儿 River Job 精确识别，每 15 秒将确认失活的 Job 恢复为可重试，避免等待 River 一小时通用救援窗口。运行时恢复同步清除中断及终态任务的 admission、执行令牌和并发租约，防止状态残留污染队列指标。正常运行中的长任务仍由 running 状态和活跃心跳保护。验证通过完整 Go 测试、go vet、竞态检查、迁移安全检查和 PostgreSQL 18 集成测试。 --- apps/api/cmd/gateway/main.go | 3 +- apps/api/internal/runner/queue_client_test.go | 109 +++++++++++++++++- apps/api/internal/runner/queue_worker.go | 38 +++++- .../store/admission_queue_integration_test.go | 77 +++++++++++++ apps/api/internal/store/rate_limits.go | 43 +++++++ apps/api/internal/store/worker_registry.go | 102 ++++++++++++++++ 6 files changed, 369 insertions(+), 3 deletions(-) diff --git a/apps/api/cmd/gateway/main.go b/apps/api/cmd/gateway/main.go index 16ce331..f117ab1 100644 --- a/apps/api/cmd/gateway/main.go +++ b/apps/api/cmd/gateway/main.go @@ -47,13 +47,14 @@ func main() { if recovery, err := db.RecoverInterruptedRuntimeState(ctx); err != nil { logger.Error("recover interrupted runtime state failed", "error", err) os.Exit(1) - } else if recovery.ReleasedConcurrencyLeases > 0 || recovery.ReleasedRateReservations > 0 || recovery.FailedAttempts > 0 || recovery.FailedTasks > 0 || recovery.RequeuedAsyncTasks > 0 { + } else if recovery.ReleasedConcurrencyLeases > 0 || recovery.ReleasedRateReservations > 0 || recovery.FailedAttempts > 0 || recovery.FailedTasks > 0 || recovery.RequeuedAsyncTasks > 0 || recovery.CleanedTaskAdmissions > 0 { logger.Warn("interrupted runtime state recovered", "releasedConcurrencyLeases", recovery.ReleasedConcurrencyLeases, "releasedRateReservations", recovery.ReleasedRateReservations, "failedAttempts", recovery.FailedAttempts, "failedTasks", recovery.FailedTasks, "requeuedAsyncTasks", recovery.RequeuedAsyncTasks, + "cleanedTaskAdmissions", recovery.CleanedTaskAdmissions, ) } diff --git a/apps/api/internal/runner/queue_client_test.go b/apps/api/internal/runner/queue_client_test.go index 9820f79..a254c92 100644 --- a/apps/api/internal/runner/queue_client_test.go +++ b/apps/api/internal/runner/queue_client_test.go @@ -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) + } +} diff --git a/apps/api/internal/runner/queue_worker.go b/apps/api/internal/runner/queue_worker.go index 13aaaa9..f21ea55 100644 --- a/apps/api/internal/runner/queue_worker.go +++ b/apps/api/internal/runner/queue_worker.go @@ -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 == "" { diff --git a/apps/api/internal/store/admission_queue_integration_test.go b/apps/api/internal/store/admission_queue_integration_test.go index d3b5b1e..2d728ce 100644 --- a/apps/api/internal/store/admission_queue_integration_test.go +++ b/apps/api/internal/store/admission_queue_integration_test.go @@ -446,6 +446,83 @@ WHERE id = $1::uuid`, atomicTask.ID).Scan(&riverJobID); err != nil { if err := first.DeleteTaskAdmission(ctx, atomicTask.ID); err != nil { t.Fatalf("release atomic admitted hook task: %v", err) } + + recoveryTask := createTask(true) + result, err = first.TryTaskAdmission(ctx, inputFor(recoveryTask, 100, "")) + if err != nil || !result.Admitted { + t.Fatalf("admit interrupted recovery task: result=%+v err=%v", result, err) + } + if _, err := first.pool.Exec(ctx, ` +UPDATE gateway_tasks +SET status = 'running', + execution_token = $2::uuid, + execution_lease_expires_at = now() - interval '1 second' +WHERE id = $1::uuid`, recoveryTask.ID, uuid.NewString()); err != nil { + t.Fatalf("expire interrupted recovery task lease: %v", err) + } + recovery, err := first.RecoverInterruptedRuntimeState(ctx) + if err != nil { + t.Fatalf("recover interrupted admitted task: %v", err) + } + if recovery.RequeuedAsyncTasks < 1 || recovery.CleanedTaskAdmissions < 1 { + t.Fatalf("interrupted task recovery = %+v, want requeue and admission cleanup", recovery) + } + var recoveredStatus string + var recoveredExecutionToken *string + if err := first.pool.QueryRow(ctx, ` +SELECT status, execution_token::text +FROM gateway_tasks +WHERE id = $1::uuid`, recoveryTask.ID).Scan(&recoveredStatus, &recoveredExecutionToken); err != nil { + t.Fatalf("read recovered task state: %v", err) + } + if recoveredStatus != "queued" || recoveredExecutionToken != nil { + t.Fatalf("recovered task status=%s execution_token=%v, want queued without token", recoveredStatus, recoveredExecutionToken) + } + var recoveredAdmissions, recoveredLeases int + if err := first.pool.QueryRow(ctx, ` +SELECT + (SELECT count(*) FROM gateway_task_admissions WHERE task_id = $1::uuid), + (SELECT count(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`, + recoveryTask.ID, + ).Scan(&recoveredAdmissions, &recoveredLeases); err != nil { + t.Fatalf("read recovered task cleanup: %v", err) + } + if recoveredAdmissions != 0 || recoveredLeases != 0 { + t.Fatalf("recovered task left admissions=%d leases=%d", recoveredAdmissions, recoveredLeases) + } + + terminalResidue := createTask(true) + result, err = first.TryTaskAdmission(ctx, inputFor(terminalResidue, 100, "")) + if err != nil || !result.Admitted { + t.Fatalf("admit terminal residue task: result=%+v err=%v", result, err) + } + if _, err := first.pool.Exec(ctx, ` +UPDATE gateway_tasks +SET status = 'failed', + error_code = 'integration_test', + finished_at = now(), + updated_at = now() +WHERE id = $1::uuid`, terminalResidue.ID); err != nil { + t.Fatalf("mark terminal residue task failed: %v", err) + } + recovery, err = first.RecoverInterruptedRuntimeState(ctx) + if err != nil { + t.Fatalf("clean terminal admission residue: %v", err) + } + if recovery.CleanedTaskAdmissions < 1 { + t.Fatalf("terminal recovery = %+v, want admission cleanup", recovery) + } + if err := first.pool.QueryRow(ctx, ` +SELECT + (SELECT count(*) FROM gateway_task_admissions WHERE task_id = $1::uuid), + (SELECT count(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`, + terminalResidue.ID, + ).Scan(&recoveredAdmissions, &recoveredLeases); err != nil { + t.Fatalf("count terminal admission residue: %v", err) + } + if recoveredAdmissions != 0 || recoveredLeases != 0 { + t.Fatalf("terminal task left admissions=%d leases=%d", recoveredAdmissions, recoveredLeases) + } } func TestWorkerCapacityAllocationAndFailover(t *testing.T) { diff --git a/apps/api/internal/store/rate_limits.go b/apps/api/internal/store/rate_limits.go index 83567f8..b830e4a 100644 --- a/apps/api/internal/store/rate_limits.go +++ b/apps/api/internal/store/rate_limits.go @@ -16,6 +16,7 @@ type RuntimeRecoveryResult struct { FailedAttempts int64 `json:"failedAttempts"` FailedTasks int64 `json:"failedTasks"` RequeuedAsyncTasks int64 `json:"requeuedAsyncTasks"` + CleanedTaskAdmissions int64 `json:"cleanedTaskAdmissions"` } var ErrConcurrencyLeaseLost = errors.New("concurrency lease lost") @@ -549,6 +550,8 @@ SET status = 'queued', locked_by = NULL, locked_at = NULL, heartbeat_at = NULL, + execution_token = NULL, + execution_lease_expires_at = NULL, next_run_at = now(), finished_at = NULL, updated_at = now() @@ -612,6 +615,24 @@ WHERE task.id = $1::uuid } } result.RequeuedAsyncTasks = int64(len(asyncTaskIDs)) + if len(asyncTaskIDs) > 0 { + tag, err = tx.Exec(ctx, ` +UPDATE gateway_concurrency_leases +SET released_at = now() +WHERE task_id = ANY($1::uuid[]) + AND released_at IS NULL`, asyncTaskIDs) + if err != nil { + return RuntimeRecoveryResult{}, err + } + result.ReleasedConcurrencyLeases += tag.RowsAffected() + tag, err = tx.Exec(ctx, ` +DELETE FROM gateway_task_admissions +WHERE task_id = ANY($1::uuid[])`, asyncTaskIDs) + if err != nil { + return RuntimeRecoveryResult{}, err + } + result.CleanedTaskAdmissions += tag.RowsAffected() + } taskRows, err := tx.Query(ctx, ` UPDATE gateway_tasks @@ -620,6 +641,8 @@ SET status = 'failed', error_code = 'server_restarted', error_message = 'task interrupted by service restart', remote_task_payload = '{}'::jsonb, + execution_token = NULL, + execution_lease_expires_at = NULL, finished_at = now(), updated_at = now() WHERE async_mode = false @@ -663,6 +686,26 @@ VALUES ( } } result.FailedTasks = int64(len(taskIDs)) + tag, err = tx.Exec(ctx, ` +UPDATE gateway_concurrency_leases lease +SET released_at = now() +FROM gateway_tasks task +WHERE lease.task_id = task.id + AND lease.released_at IS NULL + AND task.status IN ('succeeded', 'failed', 'cancelled')`) + if err != nil { + return RuntimeRecoveryResult{}, err + } + result.ReleasedConcurrencyLeases += tag.RowsAffected() + tag, err = tx.Exec(ctx, ` +DELETE FROM gateway_task_admissions admission +USING gateway_tasks task +WHERE admission.task_id = task.id + AND task.status IN ('succeeded', 'failed', 'cancelled')`) + if err != nil { + return RuntimeRecoveryResult{}, err + } + result.CleanedTaskAdmissions += tag.RowsAffected() return result, tx.Commit(ctx) } diff --git a/apps/api/internal/store/worker_registry.go b/apps/api/internal/store/worker_registry.go index 839b621..2f93403 100644 --- a/apps/api/internal/store/worker_registry.go +++ b/apps/api/internal/store/worker_registry.go @@ -171,3 +171,105 @@ WHERE instance_id = $1`, strings.TrimSpace(instanceID)) } return nil } + +// RecoverOrphanedAsyncRiverJobs retries River jobs whose owning worker is no +// longer alive after the gateway execution lease has already returned the task +// to queued. Normal long-running jobs remain protected by their running task +// state and the active worker heartbeat. +func (s *Store) RecoverOrphanedAsyncRiverJobs( + ctx context.Context, + workerStaleAfter time.Duration, + limit int, +) (int64, error) { + if workerStaleAfter < workerHeartbeatStaleAfter { + workerStaleAfter = workerHeartbeatStaleAfter + } + if limit <= 0 || limit > 1000 { + limit = 100 + } + tx, err := s.pool.Begin(ctx) + if err != nil { + return 0, err + } + defer tx.Rollback(ctx) + if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended('gateway-river-orphan-recovery', 0))`); err != nil { + return 0, err + } + var recovered int64 + if err := tx.QueryRow(ctx, ` +WITH orphaned AS MATERIALIZED ( + SELECT job.id AS job_id, + task.id AS task_id, + job.attempt + FROM river_job job + JOIN gateway_tasks task ON task.river_job_id = job.id + WHERE job.queue = 'gateway_tasks' + AND job.kind = 'gateway_task_run' + AND job.state = 'running' + AND job.attempted_at <= now() - $1::interval + AND task.async_mode = true + AND task.status = 'queued' + AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now()) + AND NOT EXISTS ( + SELECT 1 + FROM gateway_worker_instances worker + WHERE worker.status = 'active' + AND worker.heartbeat_at > now() - $1::interval + AND COALESCE(job.attempted_by[array_length(job.attempted_by, 1)], '') + LIKE worker.instance_id || '-exec-%' + ) + ORDER BY job.attempted_at ASC, job.id ASC + LIMIT $2 + FOR UPDATE OF job SKIP LOCKED +), +released_leases AS ( + UPDATE gateway_concurrency_leases lease + SET released_at = now() + FROM orphaned + WHERE lease.task_id = orphaned.task_id + AND lease.released_at IS NULL +), +cleared_admissions AS ( + DELETE FROM gateway_task_admissions admission + USING orphaned + WHERE admission.task_id = orphaned.task_id +), +reset_tasks AS ( + UPDATE gateway_tasks task + SET locked_by = NULL, + locked_at = NULL, + heartbeat_at = NULL, + execution_token = NULL, + execution_lease_expires_at = NULL, + next_run_at = now(), + updated_at = now() + FROM orphaned + WHERE task.id = orphaned.task_id +), +recovered_jobs AS ( + UPDATE river_job job + SET errors = array_append( + job.errors, + jsonb_build_object( + 'at', now(), + 'attempt', orphaned.attempt, + 'error', 'Orphaned gateway job recovered after worker heartbeat expired', + 'trace', '' + ) + ), + finalized_at = NULL, + scheduled_at = now(), + state = 'retryable' + FROM orphaned + WHERE job.id = orphaned.job_id + RETURNING job.id +) +SELECT count(*)::bigint +FROM recovered_jobs`, workerStaleAfter.String(), limit).Scan(&recovered); err != nil { + return 0, err + } + if err := tx.Commit(ctx); err != nil { + return 0, err + } + return recovered, nil +}