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 提交拦截。
2026-07-30 18:25:00 +08:00
parent fa818e9ebf
commit 38f87b6970
7 changed files with 306 additions and 8 deletions
@@ -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