From 38f87b6970d23eba4c8ee3ccfe70bb9064c5f1fd Mon Sep 17 00:00:00 2001 From: wangbo Date: Thu, 30 Jul 2026 18:25:00 +0800 Subject: [PATCH] =?UTF-8?q?fix(queue):=20=E9=98=B2=E6=AD=A2=E8=BF=87?= =?UTF-8?q?=E6=9C=9F=E6=89=A7=E8=A1=8C=E5=8D=A0=E4=BD=8D=E9=98=BB=E5=A1=9E?= =?UTF-8?q?=E4=B8=8E=E9=87=8D=E5=A4=8D=E6=8F=90=E4=BA=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 队列恢复后，仍在运行标记中的 River job 可能保留 waiting admission，形成全局 FIFO 队头阻塞；同时旧执行在租约被接管后仍可能创建 attempt 或切换到 submitting。 新增 stale admission 自动让出机制，并用当前 execution token 对 attempt 创建和上游提交状态切换做 fencing。任务、River job 和结算状态均不在让出流程中改写。 验证：go vet ./...；env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1；真实 PostgreSQL 集成测试覆盖 stale admission 让出与旧 token 提交拦截。 --- apps/api/internal/runner/queue_client_test.go | 103 +++++++++++++++++- apps/api/internal/runner/queue_worker.go | 14 +++ apps/api/internal/runner/service.go | 18 ++- .../store/billing_v2_integration_test.go | 44 ++++++++ apps/api/internal/store/runtime_types.go | 1 + apps/api/internal/store/tasks_runtime.go | 63 ++++++++++- apps/api/internal/store/worker_registry.go | 71 ++++++++++++ 7 files changed, 306 insertions(+), 8 deletions(-) diff --git a/apps/api/internal/runner/queue_client_test.go b/apps/api/internal/runner/queue_client_test.go index a254c92..6750392 100644 --- a/apps/api/internal/runner/queue_client_test.go +++ b/apps/api/internal/runner/queue_client_test.go @@ -73,7 +73,7 @@ WHERE id = $1`, queued.RiverJobID).Scan(&attemptedByCount); err != nil { } } -func TestRecoverOrphanedAsyncRiverJobRequiresInactiveWorker(t *testing.T) { +func TestRecoverOrphanedAsyncRiverJobAndYieldStaleAdmission(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") @@ -114,14 +114,15 @@ func TestRecoverOrphanedAsyncRiverJobRequiresInactiveWorker(t *testing.T) { orphaned := createQueuedTask("orphaned") protected := createQueuedTask("protected") + yielded := createQueuedTask("yielded") 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}) + _, _ = db.Pool().Exec(cleanupCtx, `DELETE FROM river_job WHERE id = ANY($1::bigint[])`, []int64{orphaned.RiverJobID, protected.RiverJobID, yielded.RiverJobID}) + _, _ = db.Pool().Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = ANY($1::uuid[])`, []string{orphaned.ID, protected.ID, yielded.ID}) }) if _, err := db.Pool().Exec(ctx, ` @@ -145,17 +146,109 @@ SET state = 'running', CASE id WHEN $1::bigint THEN $3 WHEN $2::bigint THEN $4 + WHEN $5::bigint THEN $4 END ]::text[] -WHERE id = ANY($5::bigint[])`, +WHERE id = ANY($6::bigint[])`, orphaned.RiverJobID, protected.RiverJobID, staleWorkerID+"-exec-1-test", activeWorkerID+"-exec-1-test", - []int64{orphaned.RiverJobID, protected.RiverJobID}, + yielded.RiverJobID, + []int64{orphaned.RiverJobID, protected.RiverJobID, yielded.RiverJobID}, ); err != nil { t.Fatalf("mark River jobs running: %v", err) } + var platformID, platformModelID string + if err := db.Pool().QueryRow(ctx, ` +SELECT platform.id::text, model.id::text +FROM integration_platforms platform +JOIN platform_models model ON model.platform_id = platform.id +ORDER BY platform.priority ASC, model.created_at ASC +LIMIT 1`).Scan(&platformID, &platformModelID); err != nil { + t.Fatalf("load admission binding: %v", err) + } + if _, err := db.Pool().Exec(ctx, ` +UPDATE gateway_tasks +SET execution_lease_expires_at = now() - interval '1 minute' +WHERE id = $1::uuid`, + yielded.ID, + ); err != nil { + t.Fatalf("expire yielded task execution lease: %v", err) + } + if _, err := db.Pool().Exec(ctx, ` +UPDATE gateway_tasks +SET execution_lease_expires_at = now() + interval '5 minutes' +WHERE id = $1::uuid`, + protected.ID, + ); err != nil { + t.Fatalf("renew protected task execution lease: %v", err) + } + if _, err := db.Pool().Exec(ctx, ` +INSERT INTO gateway_task_admissions ( + task_id, platform_id, platform_model_id, queue_key, mode, status, + priority, enqueued_at, wait_deadline_at +) +VALUES ( + $1::uuid, $2::uuid, $3::uuid, 'integration-test', 'async', 'waiting', + 100, now() - interval '1 minute', now() + interval '10 minutes' +)`, + yielded.ID, + platformID, + platformModelID, + ); err != nil { + t.Fatalf("seed stale async admission: %v", err) + } + if _, err := db.Pool().Exec(ctx, ` +INSERT INTO gateway_task_admissions ( + task_id, platform_id, platform_model_id, queue_key, mode, status, + priority, enqueued_at, wait_deadline_at +) +VALUES ( + $1::uuid, $2::uuid, $3::uuid, 'integration-test', 'async', 'waiting', + 100, now(), now() + interval '10 minutes' +)`, + protected.ID, + platformID, + platformModelID, + ); err != nil { + t.Fatalf("seed protected async admission: %v", err) + } + yieldedCount, err := db.YieldStaleAsyncTaskAdmissions(ctx, 30*time.Second, 10) + if err != nil { + t.Fatalf("yield stale async task admission: %v", err) + } + if yieldedCount != 1 { + t.Fatalf("yielded admissions=%d, want 1", yieldedCount) + } + var yieldedAdmissions, protectedAdmissions int + if err := db.Pool().QueryRow(ctx, ` +SELECT + count(*) FILTER (WHERE task_id = $1::uuid), + count(*) FILTER (WHERE task_id = $2::uuid) +FROM gateway_task_admissions +WHERE task_id = ANY($3::uuid[])`, + yielded.ID, + protected.ID, + []string{yielded.ID, protected.ID}, + ).Scan(&yieldedAdmissions, &protectedAdmissions); err != nil { + t.Fatalf("count yielded admissions: %v", err) + } + var yieldedState string + if err := db.Pool().QueryRow(ctx, ` +SELECT state::text +FROM river_job +WHERE id = $1`, yielded.RiverJobID).Scan(&yieldedState); err != nil { + t.Fatalf("read yielded River job state: %v", err) + } + if yieldedAdmissions != 0 || protectedAdmissions != 1 || yieldedState != "running" { + t.Fatalf( + "yielded/protected admissions=%d/%d River state=%s, want 0/1/running", + yieldedAdmissions, + protectedAdmissions, + yieldedState, + ) + } recovered, err := db.RecoverOrphanedAsyncRiverJobs(ctx, 30*time.Second, 10) if err != nil { diff --git a/apps/api/internal/runner/queue_worker.go b/apps/api/internal/runner/queue_worker.go index f3449a1..e765c39 100644 --- a/apps/api/internal/runner/queue_worker.go +++ b/apps/api/internal/runner/queue_worker.go @@ -589,6 +589,20 @@ func (s *Service) recoverOrphanedAsyncRiverJobs(ctx context.Context) { "cleanedTaskAdmissions", runtimeRecovery.CleanedTaskAdmissions, ) } + yieldedAdmissions, err := s.store.YieldStaleAsyncTaskAdmissions( + ctx, + orphanedRiverJobWorkerStaleAfter, + orphanedRiverJobRecoveryBatchSize, + ) + if err != nil { + if ctx.Err() == nil { + s.logger.Warn("yield stale async task admissions failed", "error", err) + } + return + } + if yieldedAdmissions > 0 { + s.logger.Warn("stale async task admissions yielded", "count", yieldedAdmissions) + } recovered, err := s.store.RecoverOrphanedAsyncRiverJobs( ctx, orphanedRiverJobWorkerStaleAfter, diff --git a/apps/api/internal/runner/service.go b/apps/api/internal/runner/service.go index 22a766c..0c3102b 100644 --- a/apps/api/internal/runner/service.go +++ b/apps/api/internal/runner/service.go @@ -1206,6 +1206,7 @@ func (s *Service) runCandidate( attemptID, err := s.store.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{ TaskID: task.ID, + ExecutionToken: task.ExecutionToken, AttemptNo: attemptNo, PlatformID: candidate.PlatformID, PlatformModelID: candidate.PlatformModelID, @@ -1334,7 +1335,13 @@ func (s *Service) runCandidate( if strings.TrimSpace(task.RemoteTaskID) != "" { submissionStatus = "response_received" markUpstreamSubmissionStarted(ctx) - if err := s.store.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, submissionStatus); err != nil { + if err := s.store.SetAttemptUpstreamSubmissionStatusForExecution( + ctx, + attemptID, + task.ID, + task.ExecutionToken, + submissionStatus, + ); err != nil { return clients.Response{}, fmt.Errorf("restore upstream submission status: %w", err) } } @@ -1345,7 +1352,13 @@ func (s *Service) runCandidate( if submissionStatus == status { return nil } - if err := s.store.SetAttemptUpstreamSubmissionStatus(context.WithoutCancel(ctx), attemptID, status); err != nil { + if err := s.store.SetAttemptUpstreamSubmissionStatusForExecution( + context.WithoutCancel(ctx), + attemptID, + task.ID, + task.ExecutionToken, + status, + ); err != nil { return err } submissionStatus = status @@ -1883,6 +1896,7 @@ func (s *Service) recordFailedAttempt(ctx context.Context, input failedAttemptRe metrics := mergeMetrics(append([]map[string]any{baseMetrics, failure}, input.ExtraMetrics...)...) attemptID, err := s.store.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{ TaskID: input.Task.ID, + ExecutionToken: input.Task.ExecutionToken, AttemptNo: attemptNo, PlatformID: platformID, PlatformModelID: platformModelID, diff --git a/apps/api/internal/store/billing_v2_integration_test.go b/apps/api/internal/store/billing_v2_integration_test.go index 30ed972..088a275 100644 --- a/apps/api/internal/store/billing_v2_integration_test.go +++ b/apps/api/internal/store/billing_v2_integration_test.go @@ -48,6 +48,15 @@ func TestTaskIdempotencyAndExecutionLease(t *testing.T) { if err := db.RenewTaskExecutionLease(ctx, created.Task.ID, firstToken, 5*time.Minute); err != nil { t.Fatalf("renew first lease: %v", err) } + firstAttemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{ + TaskID: created.Task.ID, + ExecutionToken: firstToken, + AttemptNo: 1, + Status: "running", + }) + if err != nil { + t.Fatalf("create first owned attempt: %v", err) + } if _, err := db.pool.Exec(ctx, `UPDATE gateway_tasks SET execution_lease_expires_at=now()-interval '1 second' WHERE id=$1::uuid`, created.Task.ID); err != nil { t.Fatal(err) } @@ -55,6 +64,41 @@ func TestTaskIdempotencyAndExecutionLease(t *testing.T) { if _, err := db.ClaimTaskExecution(ctx, created.Task.ID, secondToken, 5*time.Minute); err != nil { t.Fatalf("take over expired lease: %v", err) } + if err := db.SetAttemptUpstreamSubmissionStatusForExecution( + ctx, + firstAttemptID, + created.Task.ID, + firstToken, + "submitting", + ); !errors.Is(err, ErrTaskExecutionLeaseLost) { + t.Fatalf("stale attempt submission fence error=%v", err) + } + if _, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{ + TaskID: created.Task.ID, + ExecutionToken: firstToken, + AttemptNo: 2, + Status: "running", + }); !errors.Is(err, ErrTaskExecutionLeaseLost) { + t.Fatalf("stale attempt creation fence error=%v", err) + } + secondAttemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{ + TaskID: created.Task.ID, + ExecutionToken: secondToken, + AttemptNo: 2, + Status: "running", + }) + if err != nil { + t.Fatalf("create current owned attempt: %v", err) + } + if err := db.SetAttemptUpstreamSubmissionStatusForExecution( + ctx, + secondAttemptID, + created.Task.ID, + secondToken, + "submitting", + ); err != nil { + t.Fatalf("current attempt submission fence: %v", err) + } if _, err := db.FinishTaskFailure(ctx, FinishTaskFailureInput{TaskID: created.Task.ID, ExecutionToken: firstToken, Code: "old_worker", Message: "old"}); !errors.Is(err, ErrTaskExecutionLeaseLost) { t.Fatalf("old worker terminal error=%v", err) } diff --git a/apps/api/internal/store/runtime_types.go b/apps/api/internal/store/runtime_types.go index 42635af..fca24bb 100644 --- a/apps/api/internal/store/runtime_types.go +++ b/apps/api/internal/store/runtime_types.go @@ -263,6 +263,7 @@ type ConcurrencyLease struct { type CreateTaskAttemptInput struct { TaskID string + ExecutionToken string AttemptNo int PlatformID string PlatformModelID string diff --git a/apps/api/internal/store/tasks_runtime.go b/apps/api/internal/store/tasks_runtime.go index 77d52aa..b59f244 100644 --- a/apps/api/internal/store/tasks_runtime.go +++ b/apps/api/internal/store/tasks_runtime.go @@ -702,12 +702,19 @@ WHERE id = $1::uuid }) } -func (s *Store) SetAttemptUpstreamSubmissionStatus(ctx context.Context, attemptID string, status string) error { +func validateAttemptUpstreamSubmissionStatus(status string) error { switch status { case "not_submitted", "submitting", "response_received": default: return fmt.Errorf("invalid upstream submission status %q", status) } + return nil +} + +func (s *Store) SetAttemptUpstreamSubmissionStatus(ctx context.Context, attemptID string, status string) error { + if err := validateAttemptUpstreamSubmissionStatus(status); err != nil { + return err + } _, err := s.pool.Exec(ctx, ` UPDATE gateway_task_attempts SET upstream_submission_status = $2, @@ -716,6 +723,42 @@ WHERE id = $1::uuid`, attemptID, status) return err } +func (s *Store) SetAttemptUpstreamSubmissionStatusForExecution( + ctx context.Context, + attemptID string, + taskID string, + executionToken string, + status string, +) error { + if err := validateAttemptUpstreamSubmissionStatus(status); err != nil { + return err + } + tag, err := s.pool.Exec(ctx, ` +UPDATE gateway_task_attempts attempt +SET upstream_submission_status = $4, + upstream_submission_updated_at = now() +FROM gateway_tasks task +WHERE attempt.id = $1::uuid + AND attempt.task_id = task.id + AND attempt.task_id = $2::uuid + AND attempt.status = 'running' + AND task.status = 'running' + AND task.execution_token = $3::uuid + AND task.execution_lease_expires_at > now()`, + attemptID, + taskID, + executionToken, + status, + ) + if err != nil { + return err + } + if tag.RowsAffected() != 1 { + return ErrTaskExecutionLeaseLost + } + return nil +} + func (s *Store) CancelQueuedTask(ctx context.Context, taskID string, message string) (GatewayTask, bool, error) { message = strings.TrimSpace(message) if message == "" { @@ -980,6 +1023,24 @@ func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptIn } defer rollbackTransaction(tx) + if strings.TrimSpace(input.ExecutionToken) != "" { + var ownsExecution bool + if err := tx.QueryRow(ctx, ` +SELECT status = 'running' + AND execution_token = $2::uuid + AND execution_lease_expires_at > now() +FROM gateway_tasks +WHERE id = $1::uuid +FOR UPDATE`, + input.TaskID, + input.ExecutionToken, + ).Scan(&ownsExecution); err != nil { + return "", err + } + if !ownsExecution { + return "", ErrTaskExecutionLeaseLost + } + } var attemptID string err = tx.QueryRow(ctx, ` INSERT INTO gateway_task_attempts ( diff --git a/apps/api/internal/store/worker_registry.go b/apps/api/internal/store/worker_registry.go index 96f067c..87a055e 100644 --- a/apps/api/internal/store/worker_registry.go +++ b/apps/api/internal/store/worker_registry.go @@ -226,6 +226,77 @@ WHERE instance_id = $1`, strings.TrimSpace(instanceID)) return nil } +// YieldStaleAsyncTaskAdmissions removes waiting FIFO entries whose River job +// is still marked running even though the task no longer owns a live execution +// lease. The job and task remain untouched: a surviving worker can recreate +// the admission, while unrelated queue followers are no longer head-of-line +// blocked behind stale runtime ownership. +func (s *Store) YieldStaleAsyncTaskAdmissions( + ctx context.Context, + staleAfter time.Duration, + limit int, +) (int64, error) { + if staleAfter < workerHeartbeatStaleAfter { + staleAfter = workerHeartbeatStaleAfter + } + if limit <= 0 || limit > 1000 { + limit = 100 + } + tx, err := s.pool.Begin(ctx) + if err != nil { + return 0, err + } + defer rollbackTransaction(tx) + if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended('gateway-stale-admission-yield', 0))`); err != nil { + return 0, err + } + var yielded int64 + if err := tx.QueryRow(ctx, ` +WITH stale AS MATERIALIZED ( + SELECT admission.task_id + FROM gateway_task_admissions admission + JOIN gateway_tasks task ON task.id = admission.task_id + JOIN river_job job ON job.id = task.river_job_id + WHERE admission.status = 'waiting' + AND admission.mode = 'async' + 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 job.queue = 'gateway_tasks' + AND job.kind = 'gateway_task_run' + AND job.state = 'running' + AND job.attempted_at <= now() - $1::interval + ORDER BY admission.priority ASC, admission.enqueued_at ASC, admission.task_id ASC + LIMIT $2 +), +locked AS MATERIALIZED ( + SELECT task_id, + pg_advisory_xact_lock(hashtextextended('task-admission:' || task_id::text, 0)) + FROM stale +), +deleted AS ( + DELETE FROM gateway_task_admissions admission + USING locked + WHERE admission.task_id = locked.task_id + AND admission.status = 'waiting' + AND admission.mode = 'async' + RETURNING admission.task_id +) +SELECT count(*)::bigint +FROM deleted`, staleAfter.String(), limit).Scan(&yielded); err != nil { + return 0, err + } + if yielded > 0 { + if _, err := tx.Exec(ctx, `SELECT pg_notify('gateway_task_admission', '*')`); err != nil { + return 0, err + } + } + if err := tx.Commit(ctx); err != nil { + return 0, err + } + return yielded, 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