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
@@ -0,0 +1,499 @@
+package store
+
+import (
+	"context"
+	"errors"
+	"os"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+)
+
+func TestDistributedTaskAdmissionQueueAcrossStores(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 task admission PostgreSQL integration tests")
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	applyOIDCJITTestMigrations(t, ctx, databaseURL)
+
+	first, err := Connect(ctx, databaseURL)
+	if err != nil {
+		t.Fatalf("connect first store: %v", err)
+	}
+	defer first.Close()
+	second, err := Connect(ctx, databaseURL)
+	if err != nil {
+		t.Fatalf("connect second store: %v", err)
+	}
+	defer second.Close()
+	var databaseName string
+	if err := first.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
+		t.Fatalf("read test database name: %v", err)
+	}
+	if !strings.Contains(strings.ToLower(databaseName), "test") {
+		t.Fatalf("refusing to use non-test database %q", databaseName)
+	}
+
+	suffix := strings.ReplaceAll(uuid.NewString(), "-", "")
+	platform, err := first.CreatePlatform(ctx, CreatePlatformInput{
+		Provider:    "queue-test",
+		PlatformKey: "queue-test-" + suffix,
+		Name:        "Queue Test " + suffix,
+		AuthType:    "none",
+		Status:      "enabled",
+	})
+	if err != nil {
+		t.Fatalf("create platform: %v", err)
+	}
+	var platformModelID string
+	if err := first.pool.QueryRow(ctx, `
+INSERT INTO platform_models (
+  platform_id, model_name, provider_model_name, model_alias, model_type,
+  display_name, pricing_mode, enabled
+)
+VALUES ($1::uuid, $2, $2, $2, '["image_generate"]'::jsonb, $2, 'inherit_discount', true)
+RETURNING id::text`, platform.ID, "queue-model-"+suffix).Scan(&platformModelID); err != nil {
+		t.Fatalf("create platform model: %v", err)
+	}
+	taskIDs := make([]string, 0, 4)
+	t.Cleanup(func() {
+		cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cleanupCancel()
+		for _, taskID := range taskIDs {
+			_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = $1::uuid`, taskID)
+		}
+		_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM platform_models WHERE id = $1::uuid`, platformModelID)
+		_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM integration_platforms WHERE id = $1::uuid`, platform.ID)
+		_, _ = first.pool.Exec(cleanupCtx, `DELETE FROM gateway_worker_instances WHERE instance_id LIKE $1`, "queue-test-"+suffix+"%")
+	})
+
+	createTask := func(async bool) GatewayTask {
+		t.Helper()
+		task, createErr := first.CreateTask(ctx, CreateTaskInput{
+			Kind:    "images.generate",
+			Model:   "queue-model-" + suffix,
+			RunMode: "production",
+			Async:   async,
+			Request: map[string]any{"prompt": "queue admission test"},
+		}, &auth.User{ID: "queue-test-user-" + suffix, Source: "gateway"})
+		if createErr != nil {
+			t.Fatalf("create task: %v", createErr)
+		}
+		taskIDs = append(taskIDs, task.ID)
+		return task
+	}
+	scope := AdmissionScope{
+		ScopeType:       "platform_model",
+		ScopeKey:        platformModelID,
+		ScopeName:       "queue-model",
+		ConcurrentLimit: 1,
+		Amount:          1,
+		LeaseTTLSeconds: 137,
+		QueueLimit:      2,
+		MaxWaitSeconds:  600,
+	}
+	inputFor := func(task GatewayTask, priority int, waiterID string) TaskAdmissionInput {
+		mode := "sync"
+		if task.AsyncMode {
+			mode = "async"
+		}
+		return TaskAdmissionInput{
+			TaskID:          task.ID,
+			PlatformID:      platform.ID,
+			PlatformModelID: platformModelID,
+			QueueKey:        "queue-test:" + suffix,
+			Mode:            mode,
+			Priority:        priority,
+			WaiterID:        waiterID,
+			Scopes:          []AdmissionScope{scope},
+		}
+	}
+
+	running := createTask(false)
+	result, err := first.TryTaskAdmission(ctx, inputFor(running, 100, "waiter-running"))
+	if err != nil || !result.Admitted || len(result.Leases) != 1 {
+		t.Fatalf("first task admission = %+v, err=%v", result, err)
+	}
+	reloadedLeases, err := second.ActiveTaskAdmissionLeases(ctx, running.ID)
+	if err != nil || len(reloadedLeases) != 1 || reloadedLeases[0].TTL != 137*time.Second {
+		t.Fatalf("reloaded admission leases = %+v, err=%v, want one 137s lease", reloadedLeases, err)
+	}
+	lowerPriority := createTask(false)
+	result, err = second.TryTaskAdmission(ctx, inputFor(lowerPriority, 20, "waiter-low"))
+	if err != nil || result.Admitted {
+		t.Fatalf("second task should wait: result=%+v err=%v", result, err)
+	}
+	higherPriorityAsync := createTask(true)
+	result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, ""))
+	if err != nil || result.Admitted {
+		t.Fatalf("third task should wait: result=%+v err=%v", result, err)
+	}
+	rejected := createTask(true)
+	_, err = second.TryTaskAdmission(ctx, inputFor(rejected, 30, ""))
+	var limitErr *RateLimitExceededError
+	if !errors.As(err, &limitErr) || limitErr.Metric != "queue_size" || limitErr.Reason != "queue_full" ||
+		limitErr.QueueDepth != 2 || limitErr.QueueLimit != 2 {
+		t.Fatalf("fourth task error = %#v, want queue_full depth 2/2", err)
+	}
+
+	var queueCounters int
+	if err := first.pool.QueryRow(ctx, `
+SELECT COUNT(*)
+FROM gateway_rate_limit_counters
+WHERE metric = 'queue_size'
+  AND scope_key = $1`, platformModelID).Scan(&queueCounters); err != nil {
+		t.Fatalf("count queue counters: %v", err)
+	}
+	if queueCounters != 0 {
+		t.Fatalf("queue_size wrote %d fixed-window counters", queueCounters)
+	}
+	var attempts int
+	if err := first.pool.QueryRow(ctx, `
+SELECT COUNT(*)
+FROM gateway_task_attempts
+WHERE task_id = ANY($1::uuid[])`, taskIDs).Scan(&attempts); err != nil {
+		t.Fatalf("count attempts: %v", err)
+	}
+	if attempts != 0 {
+		t.Fatalf("waiting admissions created %d task attempts", attempts)
+	}
+
+	if err := first.DeleteTaskAdmission(ctx, running.ID); err != nil {
+		t.Fatalf("release running admission: %v", err)
+	}
+	result, err = second.TryTaskAdmission(ctx, inputFor(lowerPriority, 20, "waiter-low"))
+	if err != nil {
+		t.Fatalf("try lower-priority task after release: %v", err)
+	}
+	if result.Admitted {
+		t.Fatal("lower-priority task bypassed higher-priority asynchronous waiter")
+	}
+	result, err = first.TryTaskAdmission(ctx, inputFor(higherPriorityAsync, 10, ""))
+	if err != nil || !result.Admitted {
+		t.Fatalf("higher-priority async task was not admitted first: result=%+v err=%v", result, err)
+	}
+	for _, taskID := range taskIDs {
+		_ = first.DeleteTaskAdmission(ctx, taskID)
+	}
+
+	concurrentTasks := []GatewayTask{createTask(false), createTask(true), createTask(false), createTask(true)}
+	type concurrentResult struct {
+		result TaskAdmissionResult
+		err    error
+	}
+	start := make(chan struct{})
+	results := make(chan concurrentResult, len(concurrentTasks))
+	var waitGroup sync.WaitGroup
+	for index, task := range concurrentTasks {
+		waitGroup.Add(1)
+		go func(index int, task GatewayTask) {
+			defer waitGroup.Done()
+			<-start
+			target := first
+			if index%2 == 1 {
+				target = second
+			}
+			waiterID := ""
+			if !task.AsyncMode {
+				waiterID = "concurrent-waiter-" + task.ID
+			}
+			admission, admissionErr := target.TryTaskAdmission(ctx, inputFor(task, 100, waiterID))
+			results <- concurrentResult{result: admission, err: admissionErr}
+		}(index, task)
+	}
+	close(start)
+	waitGroup.Wait()
+	close(results)
+	admittedCount := 0
+	waitingCount := 0
+	queueFullCount := 0
+	for item := range results {
+		if item.err == nil && item.result.Admitted {
+			admittedCount++
+			continue
+		}
+		if item.err == nil {
+			waitingCount++
+			continue
+		}
+		if errors.As(item.err, &limitErr) && limitErr.Metric == "queue_size" {
+			queueFullCount++
+			continue
+		}
+		t.Fatalf("unexpected concurrent admission result: %+v err=%v", item.result, item.err)
+	}
+	if admittedCount != 1 || waitingCount != 2 || queueFullCount != 1 {
+		t.Fatalf("concurrent outcomes admitted=%d waiting=%d queue_full=%d, want 1/2/1", admittedCount, waitingCount, queueFullCount)
+	}
+	var activeLeases, waitingAdmissions int
+	if err := first.pool.QueryRow(ctx, `
+SELECT COUNT(*)
+FROM gateway_concurrency_leases
+WHERE scope_type = 'platform_model'
+  AND scope_key = $1
+  AND released_at IS NULL
+  AND expires_at > now()`, platformModelID).Scan(&activeLeases); err != nil {
+		t.Fatalf("count active leases: %v", err)
+	}
+	if err := first.pool.QueryRow(ctx, `
+SELECT COUNT(*)
+FROM gateway_task_admissions
+WHERE platform_model_id = $1::uuid
+  AND status = 'waiting'`, platformModelID).Scan(&waitingAdmissions); err != nil {
+		t.Fatalf("count waiting admissions: %v", err)
+	}
+	if activeLeases != 1 || waitingAdmissions != 2 {
+		t.Fatalf("database bounds active=%d waiting=%d, want 1/2", activeLeases, waitingAdmissions)
+	}
+
+	var deadlineTaskID string
+	if err := first.pool.QueryRow(ctx, `
+SELECT task_id::text
+FROM gateway_task_admissions
+WHERE platform_model_id = $1::uuid
+  AND status = 'waiting'
+ORDER BY task_id
+LIMIT 1`, platformModelID).Scan(&deadlineTaskID); err != nil {
+		t.Fatalf("select deadline waiter: %v", err)
+	}
+	if _, err := first.pool.Exec(ctx, `
+UPDATE gateway_task_admissions
+SET enqueued_at = now() - interval '2 seconds',
+    wait_deadline_at = now() - interval '1 second'
+WHERE task_id = $1::uuid`, deadlineTaskID); err != nil {
+		t.Fatalf("expire queue deadline: %v", err)
+	}
+	reaped, err := second.ReapExpiredTaskAdmissions(ctx, 10)
+	if err != nil || reaped.ExpiredDeadlines < 1 {
+		t.Fatalf("deadline reap = %+v, err=%v", reaped, err)
+	}
+	var taskStatus, taskErrorCode string
+	if err := first.pool.QueryRow(ctx, `
+SELECT status, COALESCE(error_code, '')
+FROM gateway_tasks
+WHERE id = $1::uuid`, deadlineTaskID).Scan(&taskStatus, &taskErrorCode); err != nil {
+		t.Fatalf("read deadline task: %v", err)
+	}
+	if taskStatus != "failed" || taskErrorCode != "queue_timeout" {
+		t.Fatalf("deadline task status=%s code=%s, want failed/queue_timeout", taskStatus, taskErrorCode)
+	}
+
+	for _, taskID := range taskIDs {
+		_ = first.DeleteTaskAdmission(ctx, taskID)
+	}
+	leaseHolder := createTask(false)
+	if result, err := first.TryTaskAdmission(ctx, inputFor(leaseHolder, 100, "lease-holder")); err != nil || !result.Admitted {
+		t.Fatalf("admit lease holder: result=%+v err=%v", result, err)
+	}
+	disconnectedWaiter := createTask(false)
+	if result, err := second.TryTaskAdmission(ctx, inputFor(disconnectedWaiter, 100, "disconnected-waiter")); err != nil || result.Admitted {
+		t.Fatalf("queue disconnected waiter: result=%+v err=%v", result, err)
+	}
+	if _, err := first.pool.Exec(ctx, `
+UPDATE gateway_task_admissions
+SET waiter_lease_expires_at = now() - interval '1 second'
+WHERE task_id = $1::uuid`, disconnectedWaiter.ID); err != nil {
+		t.Fatalf("expire waiter lease: %v", err)
+	}
+	reaped, err = second.ReapExpiredTaskAdmissions(ctx, 10)
+	if err != nil || reaped.ExpiredWaiters < 1 {
+		t.Fatalf("waiter reap = %+v, err=%v", reaped, err)
+	}
+	if err := first.pool.QueryRow(ctx, `
+SELECT status, COALESCE(error_code, '')
+FROM gateway_tasks
+WHERE id = $1::uuid`, disconnectedWaiter.ID).Scan(&taskStatus, &taskErrorCode); err != nil {
+		t.Fatalf("read disconnected task: %v", err)
+	}
+	if taskStatus != "cancelled" || taskErrorCode != "client_disconnected" {
+		t.Fatalf("disconnected task status=%s code=%s, want cancelled/client_disconnected", taskStatus, taskErrorCode)
+	}
+	if err := first.DeleteTaskAdmission(ctx, leaseHolder.ID); err != nil {
+		t.Fatalf("release lease holder admission: %v", err)
+	}
+
+	preSubmission := createTask(false)
+	claimed, err := first.ClaimTaskPreparation(ctx, preSubmission.ID, uuid.NewString(), time.Minute)
+	if err != nil {
+		t.Fatalf("claim pre-submission task: %v", err)
+	}
+	result, err = first.TryTaskAdmission(ctx, inputFor(claimed, 100, "pre-submission"))
+	if err != nil || !result.Admitted {
+		t.Fatalf("admit pre-submission task: result=%+v err=%v", result, err)
+	}
+	var attemptID string
+	if err := first.pool.QueryRow(ctx, `
+INSERT INTO gateway_task_attempts (
+  task_id, attempt_no, platform_id, platform_model_id, queue_key, status,
+  upstream_submission_status
+)
+VALUES ($1::uuid, 1, $2::uuid, $3::uuid, $4, 'running', 'not_submitted')
+RETURNING id::text`,
+		preSubmission.ID, platform.ID, platformModelID, "queue-test:"+suffix,
+	).Scan(&attemptID); err != nil {
+		t.Fatalf("create pre-submission attempt: %v", err)
+	}
+	if _, err := first.pool.Exec(ctx, `
+INSERT INTO gateway_task_param_preprocessing_logs (
+  task_id, attempt_id, attempt_no, model_type
+)
+VALUES ($1::uuid, $2::uuid, 1, 'image_generate')`, preSubmission.ID, attemptID); err != nil {
+		t.Fatalf("create pre-submission log: %v", err)
+	}
+	cancelled, changed, err := first.CancelTaskBeforeUpstreamSubmission(
+		ctx,
+		preSubmission.ID,
+		claimed.ExecutionToken,
+		"client disconnected before upstream submission",
+	)
+	if err != nil || !changed || cancelled.Status != "cancelled" || cancelled.ErrorCode != "client_disconnected" {
+		t.Fatalf("cancel pre-submission task = %+v changed=%v err=%v", cancelled, changed, err)
+	}
+	var remainingAdmissions, remainingLeases, remainingAttempts, remainingPreprocessing int
+	if err := first.pool.QueryRow(ctx, `
+SELECT
+  (SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
+  (SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL),
+  (SELECT COUNT(*) FROM gateway_task_attempts WHERE task_id = $1::uuid),
+  (SELECT COUNT(*) FROM gateway_task_param_preprocessing_logs WHERE task_id = $1::uuid)`,
+		preSubmission.ID,
+	).Scan(&remainingAdmissions, &remainingLeases, &remainingAttempts, &remainingPreprocessing); err != nil {
+		t.Fatalf("read pre-submission cleanup: %v", err)
+	}
+	if remainingAdmissions != 0 || remainingLeases != 0 || remainingAttempts != 0 || remainingPreprocessing != 0 {
+		t.Fatalf(
+			"pre-submission cleanup admissions=%d leases=%d attempts=%d preprocessing=%d",
+			remainingAdmissions,
+			remainingLeases,
+			remainingAttempts,
+			remainingPreprocessing,
+		)
+	}
+
+	postSubmission := createTask(false)
+	claimed, err = first.ClaimTaskPreparation(ctx, postSubmission.ID, uuid.NewString(), time.Minute)
+	if err != nil {
+		t.Fatalf("claim post-submission task: %v", err)
+	}
+	if _, err := first.pool.Exec(ctx, `
+INSERT INTO gateway_task_attempts (
+  task_id, attempt_no, platform_id, platform_model_id, queue_key, status,
+  upstream_submission_status
+)
+VALUES ($1::uuid, 1, $2::uuid, $3::uuid, $4, 'running', 'submitting')`,
+		postSubmission.ID, platform.ID, platformModelID, "queue-test:"+suffix,
+	); err != nil {
+		t.Fatalf("create submitted attempt: %v", err)
+	}
+	_, changed, err = first.CancelTaskBeforeUpstreamSubmission(
+		ctx,
+		postSubmission.ID,
+		claimed.ExecutionToken,
+		"client disconnected after upstream submission",
+	)
+	if err != nil || changed {
+		t.Fatalf("post-submission cancellation changed=%v err=%v, want unchanged", changed, err)
+	}
+
+	atomicTask := createTask(true)
+	hookFailure := errors.New("synthetic admitted hook failure")
+	_, err = first.TryTaskAdmissionWithAdmittedHook(ctx, inputFor(atomicTask, 100, ""), func(pgx.Tx) error {
+		return hookFailure
+	})
+	if !errors.Is(err, hookFailure) {
+		t.Fatalf("admitted hook failure = %v, want synthetic failure", err)
+	}
+	var atomicAdmissions, atomicLeases int
+	if err := first.pool.QueryRow(ctx, `
+SELECT
+  (SELECT COUNT(*) FROM gateway_task_admissions WHERE task_id = $1::uuid),
+  (SELECT COUNT(*) FROM gateway_concurrency_leases WHERE task_id = $1::uuid AND released_at IS NULL)`,
+		atomicTask.ID,
+	).Scan(&atomicAdmissions, &atomicLeases); err != nil {
+		t.Fatalf("read rolled back admitted hook state: %v", err)
+	}
+	if atomicAdmissions != 0 || atomicLeases != 0 {
+		t.Fatalf("failed admitted hook left admissions=%d leases=%d", atomicAdmissions, atomicLeases)
+	}
+	const syntheticRiverJobID int64 = 987654321
+	result, err = first.TryTaskAdmissionWithAdmittedHook(ctx, inputFor(atomicTask, 100, ""), func(tx pgx.Tx) error {
+		_, hookErr := tx.Exec(ctx, `
+UPDATE gateway_tasks
+SET river_job_id = $2
+WHERE id = $1::uuid`, atomicTask.ID, syntheticRiverJobID)
+		return hookErr
+	})
+	if err != nil || !result.Admitted {
+		t.Fatalf("atomic admitted hook result=%+v err=%v", result, err)
+	}
+	var riverJobID int64
+	if err := first.pool.QueryRow(ctx, `
+SELECT river_job_id
+FROM gateway_tasks
+WHERE id = $1::uuid`, atomicTask.ID).Scan(&riverJobID); err != nil {
+		t.Fatalf("read atomic River job marker: %v", err)
+	}
+	if riverJobID != syntheticRiverJobID {
+		t.Fatalf("atomic River job marker=%d, want %d", riverJobID, syntheticRiverJobID)
+	}
+	if err := first.DeleteTaskAdmission(ctx, atomicTask.ID); err != nil {
+		t.Fatalf("release atomic admitted hook task: %v", err)
+	}
+}
+
+func TestWorkerCapacityAllocationAndFailover(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 worker registry PostgreSQL integration tests")
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+	applyOIDCJITTestMigrations(t, ctx, databaseURL)
+	db, err := Connect(ctx, databaseURL)
+	if err != nil {
+		t.Fatalf("connect store: %v", err)
+	}
+	defer db.Close()
+	var databaseName string
+	if err := db.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
+		t.Fatalf("read test database name: %v", err)
+	}
+	if !strings.Contains(strings.ToLower(databaseName), "test") {
+		t.Fatalf("refusing to use non-test database %q", databaseName)
+	}
+
+	prefix := "queue-test-" + strings.ReplaceAll(uuid.NewString(), "-", "")
+	firstID := prefix + "-a"
+	secondID := prefix + "-b"
+	defer db.pool.Exec(context.Background(), `DELETE FROM gateway_worker_instances WHERE instance_id = ANY($1::text[])`, []string{firstID, secondID})
+
+	first, err := db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
+	if err != nil || first.Allocated != 5 || first.ActiveInstances != 1 {
+		t.Fatalf("first allocation = %+v, err=%v", first, err)
+	}
+	second, err := db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: secondID, DesiredCapacity: 5})
+	if err != nil || second.Allocated != 2 || second.ActiveInstances != 2 {
+		t.Fatalf("second allocation = %+v, err=%v", second, err)
+	}
+	first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
+	if err != nil || first.Allocated != 3 || first.ActiveInstances != 2 {
+		t.Fatalf("rebalanced first allocation = %+v, err=%v", first, err)
+	}
+	if _, err := db.pool.Exec(ctx, `
+UPDATE gateway_worker_instances
+SET heartbeat_at = now() - interval '16 seconds'
+WHERE instance_id = $1`, secondID); err != nil {
+		t.Fatalf("expire second heartbeat: %v", err)
+	}
+	first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
+	if err != nil || first.Allocated != 5 || first.ActiveInstances != 1 {
+		t.Fatalf("single-worker failover allocation = %+v, err=%v", first, err)
+	}
+}