fix(queue): 防止过期执行占位阻塞与重复提交

队列恢复后,仍在运行标记中的 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 提交拦截。
This commit is contained in:
2026-07-30 18:25:00 +08:00
parent fa818e9ebf
commit 38f87b6970
7 changed files with 306 additions and 8 deletions
+98 -5
View File
@@ -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 {
+14
View File
@@ -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,
+16 -2
View File
@@ -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,
@@ -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)
}
+1
View File
@@ -263,6 +263,7 @@ type ConcurrencyLease struct {
type CreateTaskAttemptInput struct {
TaskID string
ExecutionToken string
AttemptNo int
PlatformID string
PlatformModelID string
+62 -1
View File
@@ -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 (
@@ -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