feat(queue): 增加非文本模型分布式准入队列

使用 PostgreSQL 统一同步与异步非文本任务的并发准入、持久化等待和 Worker 容量分配，并将生产 API 与独立 Worker 角色拆分。补充策略管理、共享契约、OpenAPI、Kubernetes 双节点 Worker 清单及跨节点验收脚本；未默认启用任何生产 queue_size 策略。已在原基线完成 Go、前端、迁移、Shell、Kustomize 与长任务容量验收；合入最新主干后将重新执行发布门禁。
2026-07-29 16:15:43 +08:00
parent 2e5a90731b
commit 9e4fc7362d
55 changed files with 4565 additions and 237 deletions
@@ -389,6 +389,82 @@ RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second
 	return task, 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, `
+UPDATE gateway_tasks
+SET execution_token = $2::uuid,
+    execution_lease_expires_at = now() + ($3::int * interval '1 second'),
+    locked_at = now(),
+    heartbeat_at = now(),
+    updated_at = now()
+WHERE id = $1::uuid
+  AND 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()
+  )
+RETURNING `+gatewayTaskColumns, taskID, executionToken, int(leaseTTL/time.Second)))
+	if errors.Is(err, pgx.ErrNoRows) {
+		return GatewayTask{}, ErrTaskExecutionLeaseUnavailable
+	}
+	return task, err
+}
+
+func (s *Store) ReleaseTaskPreparation(ctx context.Context, taskID string, executionToken string) error {
+	tag, err := s.pool.Exec(ctx, `
+UPDATE gateway_tasks
+SET execution_token = NULL,
+    execution_lease_expires_at = NULL,
+    locked_at = NULL,
+    heartbeat_at = NULL,
+    updated_at = now()
+WHERE id = $1::uuid
+  AND status = 'queued'
+  AND execution_token = $2::uuid`, taskID, executionToken)
+	if err != nil {
+		return err
+	}
+	if tag.RowsAffected() != 1 {
+		return ErrTaskExecutionLeaseLost
+	}
+	return nil
+}
+
+func (s *Store) FailQueuedTask(ctx context.Context, taskID string, code string, message string) (GatewayTask, error) {
+	tag, err := s.pool.Exec(ctx, `
+UPDATE gateway_tasks
+SET status = 'failed',
+    error = NULL,
+    error_code = NULLIF($2, ''),
+    error_message = NULLIF($3, ''),
+    billing_status = CASE
+      WHEN run_mode <> 'production' OR gateway_user_id IS NULL THEN 'not_required'
+      ELSE 'released'
+    END,
+    billing_updated_at = now(),
+    locked_by = NULL,
+    locked_at = NULL,
+    heartbeat_at = NULL,
+    execution_token = NULL,
+    execution_lease_expires_at = NULL,
+    finished_at = now(),
+    updated_at = now()
+WHERE id = $1::uuid
+  AND status = 'queued'`, taskID, strings.TrimSpace(code), truncateUTF8Bytes(message, 2048))
+	if err != nil {
+		return GatewayTask{}, err
+	}
+	if tag.RowsAffected() != 1 {
+		return GatewayTask{}, ErrTaskExecutionFinished
+	}
+	return s.GetTask(ctx, taskID)
+}
+
 func (s *Store) RenewTaskExecutionLease(ctx context.Context, taskID string, executionToken string, leaseTTL time.Duration) error {
 	if leaseTTL <= 0 {
 		leaseTTL = 5 * time.Minute
@@ -399,7 +475,7 @@ SET execution_lease_expires_at = now() + ($3::int * interval '1 second'),
    heartbeat_at = now(),
    updated_at = now()
 WHERE id = $1::uuid
-  AND status = 'running'
+  AND status IN ('queued', 'running')
  AND execution_token = $2::uuid`, taskID, executionToken, int(leaseTTL/time.Second))
 	if err != nil {
 		return err
@@ -407,7 +483,7 @@ WHERE id = $1::uuid
 	if tag.RowsAffected() != 1 {
 		var stillRunning bool
 		if err := s.pool.QueryRow(ctx, `
-SELECT COALESCE((SELECT status = 'running' FROM gateway_tasks WHERE id = $1::uuid), false)`, taskID).Scan(&stillRunning); err != nil {
+SELECT COALESCE((SELECT status IN ('queued', 'running') FROM gateway_tasks WHERE id = $1::uuid), false)`, taskID).Scan(&stillRunning); err != nil {
 			return err
 		}
 		if !stillRunning {
@@ -428,7 +504,7 @@ SET status = 'running',
 	    heartbeat_at = now(),
 	    updated_at = now()
 WHERE id = $1::uuid
-  AND status = 'running'
+  AND status IN ('queued', 'running')
  AND execution_token = $2::uuid`, taskID, executionToken, modelType)
 	if err != nil {
 		return err
@@ -492,7 +568,7 @@ SET status = 'queued',
    error_message = NULL,
    updated_at = now()
 WHERE id = $1::uuid
-  AND status = 'running'
+  AND status IN ('queued', 'running')
  AND execution_token = $2::uuid
 RETURNING `+gatewayTaskColumns, taskID, executionToken, nextRunAt, strings.TrimSpace(queueKey)))
 }
@@ -566,7 +642,14 @@ func (s *Store) CancelQueuedTask(ctx context.Context, taskID string, message str
 		message = "任务已取消"
 	}
 	message = truncateUTF8Bytes(message, 2048)
-	tag, err := s.pool.Exec(ctx, `
+	var task GatewayTask
+	changed := false
+	err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
+		if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+taskID); err != nil {
+			return err
+		}
+		var err error
+		task, err = scanGatewayTask(tx.QueryRow(ctx, `
 UPDATE gateway_tasks
 SET status = 'cancelled',
    error = NULL,
@@ -576,22 +659,123 @@ SET status = 'cancelled',
    locked_by = NULL,
    locked_at = NULL,
    heartbeat_at = NULL,
+    execution_token = NULL,
+    execution_lease_expires_at = NULL,
    finished_at = now(),
    updated_at = now()
 WHERE id = $1::uuid
  AND status = 'queued'
-  AND COALESCE(remote_task_id, '') = ''`, taskID, message)
+  AND COALESCE(remote_task_id, '') = ''
+RETURNING `+gatewayTaskColumns, taskID, message))
+		if errors.Is(err, pgx.ErrNoRows) {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+		changed = true
+		if _, err := tx.Exec(ctx, `
+UPDATE gateway_concurrency_leases
+SET released_at = now()
+WHERE task_id = $1::uuid
+  AND released_at IS NULL`, taskID); err != nil {
+			return err
+		}
+		if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil {
+			return err
+		}
+		return notifyTaskAdmissionTx(ctx, tx, taskID)
+	})
 	if err != nil {
 		return GatewayTask{}, false, err
 	}
-	if tag.RowsAffected() == 0 {
-		return GatewayTask{}, false, nil
+	return task, changed, nil
+}
+
+func (s *Store) CancelTaskBeforeUpstreamSubmission(
+	ctx context.Context,
+	taskID string,
+	executionToken string,
+	message string,
+) (GatewayTask, bool, error) {
+	message = strings.TrimSpace(message)
+	if message == "" {
+		message = "client disconnected before upstream submission"
 	}
-	task, err := s.GetTask(ctx, taskID)
+	message = truncateUTF8Bytes(message, 2048)
+	var task GatewayTask
+	changed := false
+	err := pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
+		if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "task-admission:"+taskID); err != nil {
+			return err
+		}
+		var err error
+		task, err = scanGatewayTask(tx.QueryRow(ctx, `
+UPDATE gateway_tasks task
+SET status = 'cancelled',
+    error = NULL,
+    error_code = 'client_disconnected',
+    error_message = NULLIF($3::text, ''),
+    billing_status = CASE
+      WHEN run_mode <> 'production' OR gateway_user_id IS NULL THEN 'not_required'
+      ELSE 'released'
+    END,
+    billing_updated_at = now(),
+    locked_by = NULL,
+    locked_at = NULL,
+    heartbeat_at = NULL,
+    execution_token = NULL,
+    execution_lease_expires_at = NULL,
+    finished_at = now(),
+    updated_at = now()
+WHERE id = $1::uuid
+  AND status IN ('queued', 'running')
+  AND execution_token = NULLIF($2, '')::uuid
+  AND COALESCE(remote_task_id, '') = ''
+  AND NOT EXISTS (
+    SELECT 1
+    FROM gateway_task_attempts attempt
+    WHERE attempt.task_id = task.id
+      AND attempt.upstream_submission_status <> 'not_submitted'
+  )
+RETURNING `+gatewayTaskColumns, taskID, strings.TrimSpace(executionToken), message))
+		if errors.Is(err, pgx.ErrNoRows) {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+		changed = true
+		if _, err := tx.Exec(ctx, `
+DELETE FROM gateway_task_param_preprocessing_logs log
+USING gateway_task_attempts attempt
+WHERE log.attempt_id = attempt.id
+  AND attempt.task_id = $1::uuid
+  AND attempt.upstream_submission_status = 'not_submitted'`, taskID); err != nil {
+			return err
+		}
+		if _, err := tx.Exec(ctx, `
+DELETE FROM gateway_task_attempts
+WHERE task_id = $1::uuid
+  AND upstream_submission_status = 'not_submitted'`, taskID); err != nil {
+			return err
+		}
+		if _, err := tx.Exec(ctx, `
+UPDATE gateway_concurrency_leases
+SET released_at = now()
+WHERE task_id = $1::uuid
+  AND released_at IS NULL`, taskID); err != nil {
+			return err
+		}
+		if _, err := tx.Exec(ctx, `DELETE FROM gateway_task_admissions WHERE task_id = $1::uuid`, taskID); err != nil {
+			return err
+		}
+		return notifyTaskAdmissionTx(ctx, tx, taskID)
+	})
 	if err != nil {
-		return GatewayTask{}, true, err
+		return GatewayTask{}, false, err
 	}
-	return task, true, nil
+	return task, changed, nil
 }

 // CancelSubmittedTask records a confirmed upstream cancellation. Callers must
@@ -1506,7 +1690,7 @@ func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureIn
 	    finished_at = now(),
 	    updated_at = now()
 WHERE id = $1::uuid
-  AND status = 'running'
+  AND status IN ('queued', 'running')
  AND execution_token = $10::uuid`,
 			input.TaskID,
 			message,