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:
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -263,6 +263,7 @@ type ConcurrencyLease struct {
|
||||
|
||||
type CreateTaskAttemptInput struct {
|
||||
TaskID string
|
||||
ExecutionToken string
|
||||
AttemptNo int
|
||||
PlatformID string
|
||||
PlatformModelID string
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user