From f33d6d64e09a5863e8deafa81f67ecea64ccbdf9 Mon Sep 17 00:00:00 2001 From: wangbo Date: Thu, 30 Jul 2026 19:15:19 +0800 Subject: [PATCH] =?UTF-8?q?fix(queue):=20=E5=9B=9E=E6=94=B6=E5=A4=B1?= =?UTF-8?q?=E6=95=88=E4=BB=BB=E5=8A=A1=E5=B9=B6=E9=98=BB=E6=AD=A2=E4=B8=8D?= =?UTF-8?q?=E7=A1=AE=E5=AE=9A=E6=8F=90=E4=BA=A4=E9=87=8D=E6=94=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Worker 进程心跳不能证明单个 River job goroutine 仍存活，改以任务执行租约作为回收所有权栅栏。\n\n排队任务若上一次 attempt 在租约中断时处于 submitting 或 response_received，则转入人工复核并生成 release 记录，不再重新提交上游；明确被上游拒绝的响应仍允许重试。\n\n验证：go vet ./...；env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1；真实 PostgreSQL 下相关 store 与 River 集成测试通过。 --- apps/api/internal/runner/queue_client_test.go | 18 +- .../store/billing_v2_integration_test.go | 121 +++++++++++++ apps/api/internal/store/tasks_runtime.go | 166 +++++++++++++----- apps/api/internal/store/worker_registry.go | 16 +- 4 files changed, 256 insertions(+), 65 deletions(-) diff --git a/apps/api/internal/runner/queue_client_test.go b/apps/api/internal/runner/queue_client_test.go index 6750392..d167bdc 100644 --- a/apps/api/internal/runner/queue_client_test.go +++ b/apps/api/internal/runner/queue_client_test.go @@ -254,18 +254,26 @@ WHERE id = $1`, yielded.RiverJobID).Scan(&yieldedState); err != nil { if err != nil { t.Fatalf("recover orphaned River jobs: %v", err) } - if recovered != 1 { - t.Fatalf("recovered jobs=%d, want 1", recovered) + if recovered != 2 { + t.Fatalf("recovered jobs=%d, want 2", recovered) } - var orphanedState, protectedState string + var orphanedState, protectedState, yieldedRecoveredState 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) + if err := db.Pool().QueryRow(ctx, `SELECT state::text FROM river_job WHERE id = $1`, yielded.RiverJobID).Scan(&yieldedRecoveredState); err != nil { + t.Fatalf("read yielded River job state: %v", err) + } + if orphanedState != "retryable" || protectedState != "running" || yieldedRecoveredState != "retryable" { + t.Fatalf( + "River states orphaned=%s protected=%s yielded=%s, want retryable/running/retryable", + orphanedState, + protectedState, + yieldedRecoveredState, + ) } recovered, err = db.RecoverOrphanedAsyncRiverJobs(ctx, 30*time.Second, 10) if err != nil || recovered != 0 { diff --git a/apps/api/internal/store/billing_v2_integration_test.go b/apps/api/internal/store/billing_v2_integration_test.go index 088a275..719c003 100644 --- a/apps/api/internal/store/billing_v2_integration_test.go +++ b/apps/api/internal/store/billing_v2_integration_test.go @@ -168,6 +168,127 @@ WHERE task_id=$1::uuid AND event_type='task.billing.review'`, created.ID).Scan(& } } +func TestQueuedPreparationDoesNotReplayInterruptedUpstreamSubmission(t *testing.T) { + db := billingV2IntegrationStore(t) + ctx := context.Background() + tenantID, gatewayUserID := seedWalletReservationUser(t, ctx, db) + + for _, submissionStatus := range []string{"submitting", "response_received"} { + t.Run(submissionStatus, func(t *testing.T) { + user := &auth.User{ + ID: "billing-queued-review-" + uuid.NewString(), + GatewayUserID: gatewayUserID, + GatewayTenantID: tenantID, + } + created, err := db.CreateTask(ctx, CreateTaskInput{ + Kind: "images.generations", Model: "billing-v2-model", RunMode: "production", + Request: map[string]any{"model": "billing-v2-model"}, + }, user) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID) + }) + + token := uuid.NewString() + if _, err := db.ClaimTaskExecution(ctx, created.ID, token, 5*time.Minute); err != nil { + t.Fatal(err) + } + attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{ + TaskID: created.ID, AttemptNo: 1, Status: "running", + }) + if err != nil { + t.Fatal(err) + } + if err := db.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, submissionStatus); err != nil { + t.Fatal(err) + } + if _, err := db.pool.Exec(ctx, ` +UPDATE gateway_task_attempts +SET status='failed', + retryable=true, + error_code='execution_lease_expired', + error_message='attempt execution lease expired', + finished_at=now() +WHERE id=$1::uuid`, attemptID); err != nil { + t.Fatal(err) + } + if _, err := db.pool.Exec(ctx, ` +UPDATE gateway_tasks +SET status='queued', + execution_token=NULL, + execution_lease_expires_at=NULL, + heartbeat_at=NULL, + locked_at=NULL, + updated_at=now() +WHERE id=$1::uuid`, created.ID); err != nil { + t.Fatal(err) + } + + if _, err := db.ClaimTaskPreparation(ctx, created.ID, uuid.NewString(), 5*time.Minute); !errors.Is(err, ErrTaskExecutionManualReview) { + t.Fatalf("queued ambiguous preparation error=%v", err) + } + review, err := db.GetTask(ctx, created.ID) + if err != nil { + t.Fatal(err) + } + if review.Status != "failed" || review.BillingStatus != "manual_review" || review.ErrorCode != "upstream_submission_unknown" { + t.Fatalf("manual review task=%+v", review) + } + }) + } +} + +func TestQueuedPreparationCanResumeAfterKnownRejectedResponse(t *testing.T) { + db := billingV2IntegrationStore(t) + ctx := context.Background() + _, gatewayUserID := seedWalletReservationUser(t, ctx, db) + user := &auth.User{ID: "billing-queued-known-response-" + uuid.NewString(), GatewayUserID: gatewayUserID} + created, err := db.CreateTask(ctx, CreateTaskInput{ + Kind: "images.generations", Model: "billing-v2-model", RunMode: "production", + Request: map[string]any{"model": "billing-v2-model"}, + }, user) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID) + }) + + if _, err := db.ClaimTaskExecution(ctx, created.ID, uuid.NewString(), 5*time.Minute); err != nil { + t.Fatal(err) + } + attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{ + TaskID: created.ID, AttemptNo: 1, Status: "running", + }) + if err != nil { + t.Fatal(err) + } + if err := db.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, "response_received"); err != nil { + t.Fatal(err) + } + if err := db.FinishTaskAttempt(ctx, FinishTaskAttemptInput{ + AttemptID: attemptID, Status: "failed", ErrorCode: "upstream_rejected", + }); err != nil { + t.Fatal(err) + } + if _, err := db.pool.Exec(ctx, ` +UPDATE gateway_tasks +SET status='queued', + execution_token=NULL, + execution_lease_expires_at=NULL, + heartbeat_at=NULL, + locked_at=NULL, + updated_at=now() +WHERE id=$1::uuid`, created.ID); err != nil { + t.Fatal(err) + } + if _, err := db.ClaimTaskPreparation(ctx, created.ID, uuid.NewString(), 5*time.Minute); err != nil { + t.Fatalf("known rejected response should remain retryable: %v", err) + } +} + func TestExpiredExecutionLeaseCanResumeAfterKnownRejectedResponse(t *testing.T) { db := billingV2IntegrationStore(t) ctx := context.Background() diff --git a/apps/api/internal/store/tasks_runtime.go b/apps/api/internal/store/tasks_runtime.go index b59f244..0b2edcb 100644 --- a/apps/api/internal/store/tasks_runtime.go +++ b/apps/api/internal/store/tasks_runtime.go @@ -304,58 +304,15 @@ FOR UPDATE`, taskID).Scan(&queuedReady, &runningExpired, &production, &hasGatewa if !queuedReady && !runningExpired { return ErrTaskExecutionLeaseUnavailable } - if runningExpired && production { - var submissionAmbiguous bool - if err := tx.QueryRow(ctx, ` -SELECT COALESCE(( - SELECT ( - (status = 'running' AND upstream_submission_status IN ('submitting', 'response_received')) - OR (status = 'failed' AND upstream_submission_status = 'submitting') - ) - FROM gateway_task_attempts - WHERE task_id = $1::uuid - ORDER BY attempt_no DESC, started_at DESC - LIMIT 1 -), false)`, taskID).Scan(&submissionAmbiguous); err != nil { + if (queuedReady || runningExpired) && production { + submissionAmbiguous, err := taskExecutionRequiresManualReviewTx(ctx, tx, taskID) + if err != nil { return err } if submissionAmbiguous { - if _, err := tx.Exec(ctx, ` -UPDATE gateway_tasks -SET status = 'failed', - billing_status = CASE WHEN $2 THEN 'manual_review' ELSE 'not_required' END, - billing_updated_at = now(), - error = NULL, - error_code = 'upstream_submission_unknown', - error_message = 'upstream submission result is unknown', - locked_by = NULL, - locked_at = NULL, - heartbeat_at = NULL, - execution_token = NULL, - execution_lease_expires_at = NULL, - remote_task_payload = '{}'::jsonb, - finished_at = now(), - updated_at = now() -WHERE id = $1::uuid`, taskID, hasGatewayUser); err != nil { + if err := markTaskExecutionManualReviewTx(ctx, tx, taskID, hasGatewayUser); err != nil { return err } - if hasGatewayUser { - payloadJSON, _ := json.Marshal(map[string]any{ - "taskId": taskID, "classification": "upstream_submission_unknown", - }) - if _, err := tx.Exec(ctx, ` -INSERT INTO settlement_outbox ( - task_id, event_type, action, amount, currency, pricing_snapshot, payload, - status, next_attempt_at, manual_review_reason -) -SELECT id, 'task.billing.review', 'release', reservation_amount, billing_currency, - pricing_snapshot, $2::jsonb, 'manual_review', now(), 'upstream_submission_unknown' -FROM gateway_tasks -WHERE id = $1::uuid -ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON)); err != nil { - return err - } - } manualReview = true return nil } @@ -389,11 +346,116 @@ RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second return task, nil } +func taskExecutionRequiresManualReviewTx(ctx context.Context, tx pgx.Tx, taskID string) (bool, error) { + var submissionAmbiguous bool + err := tx.QueryRow(ctx, ` +SELECT COALESCE(( + SELECT ( + (status = 'running' AND upstream_submission_status IN ('submitting', 'response_received')) + OR (status = 'failed' AND upstream_submission_status = 'submitting') + OR ( + status = 'failed' + AND upstream_submission_status = 'response_received' + AND error_code = 'execution_lease_expired' + ) + ) + FROM gateway_task_attempts + WHERE task_id = $1::uuid + ORDER BY attempt_no DESC, started_at DESC + LIMIT 1 +), false)`, taskID).Scan(&submissionAmbiguous) + return submissionAmbiguous, err +} + +func markTaskExecutionManualReviewTx( + ctx context.Context, + tx pgx.Tx, + taskID string, + hasGatewayUser bool, +) error { + if _, err := tx.Exec(ctx, ` +UPDATE gateway_tasks +SET status = 'failed', + billing_status = CASE WHEN $2 THEN 'manual_review' ELSE 'not_required' END, + billing_updated_at = now(), + error = NULL, + error_code = 'upstream_submission_unknown', + error_message = 'upstream submission result is unknown', + locked_by = NULL, + locked_at = NULL, + heartbeat_at = NULL, + execution_token = NULL, + execution_lease_expires_at = NULL, + remote_task_payload = '{}'::jsonb, + finished_at = now(), + updated_at = now() +WHERE id = $1::uuid`, taskID, hasGatewayUser); err != nil { + return err + } + if !hasGatewayUser { + return nil + } + payloadJSON, _ := json.Marshal(map[string]any{ + "taskId": taskID, "classification": "upstream_submission_unknown", + }) + if _, err := tx.Exec(ctx, ` +INSERT INTO settlement_outbox ( + task_id, event_type, action, amount, currency, pricing_snapshot, payload, + status, next_attempt_at, manual_review_reason +) +SELECT id, 'task.billing.review', 'release', reservation_amount, billing_currency, + pricing_snapshot, $2::jsonb, 'manual_review', now(), 'upstream_submission_unknown' +FROM gateway_tasks +WHERE id = $1::uuid +ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON)); err != nil { + return err + } + return nil +} + func (s *Store) ClaimTaskPreparation(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) (GatewayTask, error) { if leaseTTL <= 0 { leaseTTL = 5 * time.Minute } - task, err := scanGatewayTask(s.pool.QueryRow(ctx, ` + var task GatewayTask + manualReview := false + err := s.beginTransaction(ctx, func(tx pgx.Tx) error { + var queuedReady bool + var production bool + var hasGatewayUser bool + if err := tx.QueryRow(ctx, ` +SELECT status = 'queued' + AND next_run_at <= now() + AND ( + execution_token IS NULL + OR execution_lease_expires_at IS NULL + OR execution_lease_expires_at <= now() + ), + run_mode = 'production', + gateway_user_id IS NOT NULL +FROM gateway_tasks +WHERE id = $1::uuid +FOR UPDATE`, taskID).Scan(&queuedReady, &production, &hasGatewayUser); err != nil { + return err + } + if !queuedReady { + return ErrTaskExecutionLeaseUnavailable + } + if production { + submissionAmbiguous, err := taskExecutionRequiresManualReviewTx(ctx, tx, taskID) + if err != nil { + return err + } + if submissionAmbiguous { + if err := markTaskExecutionManualReviewTx(ctx, tx, taskID, hasGatewayUser); err != nil { + return err + } + manualReview = true + return nil + } + } + var err error + task, err = scanGatewayTask(tx.QueryRow(ctx, ` UPDATE gateway_tasks SET execution_token = $2::uuid, execution_lease_expires_at = now() + ($3::int * interval '1 second'), @@ -409,9 +471,17 @@ WHERE id = $1::uuid OR execution_lease_expires_at <= now() ) RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second))) + return err + }) if errors.Is(err, pgx.ErrNoRows) { return GatewayTask{}, ErrTaskExecutionLeaseUnavailable } + if err != nil { + return GatewayTask{}, err + } + if manualReview { + return GatewayTask{}, ErrTaskExecutionManualReview + } return task, err } diff --git a/apps/api/internal/store/worker_registry.go b/apps/api/internal/store/worker_registry.go index 87a055e..f075b87 100644 --- a/apps/api/internal/store/worker_registry.go +++ b/apps/api/internal/store/worker_registry.go @@ -297,10 +297,10 @@ FROM deleted`, staleAfter.String(), limit).Scan(&yielded); err != nil { 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 -// state and the active worker heartbeat. +// RecoverOrphanedAsyncRiverJobs retries stale River jobs after the gateway +// execution lease has already returned the task to queued. A process heartbeat +// does not prove that an individual River job goroutine is still alive; the +// task execution lease is the ownership fence. func (s *Store) RecoverOrphanedAsyncRiverJobs( ctx context.Context, workerStaleAfter time.Duration, @@ -335,14 +335,6 @@ WITH orphaned AS MATERIALIZED ( 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