fix(queue): 回收失效任务并阻止不确定提交重放

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 集成测试通过。
This commit is contained in:
2026-07-30 19:15:19 +08:00
parent 9e06ed6162
commit f33d6d64e0
4 changed files with 256 additions and 65 deletions
+13 -5
View File
@@ -254,18 +254,26 @@ WHERE id = $1`, yielded.RiverJobID).Scan(&yieldedState); err != nil {
if err != nil { if err != nil {
t.Fatalf("recover orphaned River jobs: %v", err) t.Fatalf("recover orphaned River jobs: %v", err)
} }
if recovered != 1 { if recovered != 2 {
t.Fatalf("recovered jobs=%d, want 1", recovered) 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 { 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) 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 { 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) t.Fatalf("read protected River job state: %v", err)
} }
if orphanedState != "retryable" || protectedState != "running" { if err := db.Pool().QueryRow(ctx, `SELECT state::text FROM river_job WHERE id = $1`, yielded.RiverJobID).Scan(&yieldedRecoveredState); err != nil {
t.Fatalf("River states orphaned=%s protected=%s, want retryable/running", orphanedState, protectedState) 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) recovered, err = db.RecoverOrphanedAsyncRiverJobs(ctx, 30*time.Second, 10)
if err != nil || recovered != 0 { if err != nil || recovered != 0 {
@@ -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) { func TestExpiredExecutionLeaseCanResumeAfterKnownRejectedResponse(t *testing.T) {
db := billingV2IntegrationStore(t) db := billingV2IntegrationStore(t)
ctx := context.Background() ctx := context.Background()
+118 -48
View File
@@ -304,58 +304,15 @@ FOR UPDATE`, taskID).Scan(&queuedReady, &runningExpired, &production, &hasGatewa
if !queuedReady && !runningExpired { if !queuedReady && !runningExpired {
return ErrTaskExecutionLeaseUnavailable return ErrTaskExecutionLeaseUnavailable
} }
if runningExpired && production { if (queuedReady || runningExpired) && production {
var submissionAmbiguous bool submissionAmbiguous, err := taskExecutionRequiresManualReviewTx(ctx, tx, taskID)
if err := tx.QueryRow(ctx, ` if err != nil {
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 {
return err return err
} }
if submissionAmbiguous { if submissionAmbiguous {
if _, err := tx.Exec(ctx, ` if err := markTaskExecutionManualReviewTx(ctx, tx, taskID, hasGatewayUser); err != nil {
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 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 manualReview = true
return nil return nil
} }
@@ -389,11 +346,116 @@ RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second
return task, nil 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) { func (s *Store) ClaimTaskPreparation(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) (GatewayTask, error) {
if leaseTTL <= 0 { if leaseTTL <= 0 {
leaseTTL = 5 * time.Minute 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 UPDATE gateway_tasks
SET execution_token = $2::uuid, SET execution_token = $2::uuid,
execution_lease_expires_at = now() + ($3::int * interval '1 second'), execution_lease_expires_at = now() + ($3::int * interval '1 second'),
@@ -409,9 +471,17 @@ WHERE id = $1::uuid
OR execution_lease_expires_at <= now() OR execution_lease_expires_at <= now()
) )
RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second))) RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second)))
return err
})
if errors.Is(err, pgx.ErrNoRows) { if errors.Is(err, pgx.ErrNoRows) {
return GatewayTask{}, ErrTaskExecutionLeaseUnavailable return GatewayTask{}, ErrTaskExecutionLeaseUnavailable
} }
if err != nil {
return GatewayTask{}, err
}
if manualReview {
return GatewayTask{}, ErrTaskExecutionManualReview
}
return task, err return task, err
} }
+4 -12
View File
@@ -297,10 +297,10 @@ FROM deleted`, staleAfter.String(), limit).Scan(&yielded); err != nil {
return yielded, nil return yielded, nil
} }
// RecoverOrphanedAsyncRiverJobs retries River jobs whose owning worker is no // RecoverOrphanedAsyncRiverJobs retries stale River jobs after the gateway
// longer alive after the gateway execution lease has already returned the task // execution lease has already returned the task to queued. A process heartbeat
// to queued. Normal long-running jobs remain protected by their running task // does not prove that an individual River job goroutine is still alive; the
// state and the active worker heartbeat. // task execution lease is the ownership fence.
func (s *Store) RecoverOrphanedAsyncRiverJobs( func (s *Store) RecoverOrphanedAsyncRiverJobs(
ctx context.Context, ctx context.Context,
workerStaleAfter time.Duration, workerStaleAfter time.Duration,
@@ -335,14 +335,6 @@ WITH orphaned AS MATERIALIZED (
AND task.async_mode = true AND task.async_mode = true
AND task.status = 'queued' AND task.status = 'queued'
AND (task.execution_lease_expires_at IS NULL OR task.execution_lease_expires_at <= now()) 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 ORDER BY job.attempted_at ASC, job.id ASC
LIMIT $2 LIMIT $2
FOR UPDATE OF job SKIP LOCKED FOR UPDATE OF job SKIP LOCKED