From 387bf5842dac4387e66172a34d0e884e418bc01b Mon Sep 17 00:00:00 2001 From: wangbo Date: Fri, 31 Jul 2026 09:29:01 +0800 Subject: [PATCH] =?UTF-8?q?perf(worker):=20=E9=9A=94=E7=A6=BB=E5=BC=82?= =?UTF-8?q?=E6=AD=A5=E7=AD=89=E5=BE=85=E7=99=BB=E8=AE=B0=E4=B8=8E=E6=89=A7?= =?UTF-8?q?=E8=A1=8C=E9=94=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 为异步等待队列登记增加独立的进程内和 PostgreSQL advisory lock 域，使严格队列上限校验不再等待 Worker 执行槽事务的跨城同步提交。任务级锁保持不变，调度期间 active 与 waiting 总量守恒。\n\n新增 PostgreSQL 集成测试，验证执行 scope 锁被占用时等待登记仍可独立完成。\n\n验证：go vet ./...；go test ./... -count=1；隔离 PostgreSQL 集成测试；迁移安全检查；gofmt；git diff --check。 --- apps/api/internal/store/admission_lock.go | 8 ++ apps/api/internal/store/admission_queue.go | 18 ++- ...sion_registration_lock_integration_test.go | 116 ++++++++++++++++++ 3 files changed, 138 insertions(+), 4 deletions(-) create mode 100644 apps/api/internal/store/admission_registration_lock_integration_test.go diff --git a/apps/api/internal/store/admission_lock.go b/apps/api/internal/store/admission_lock.go index ad9d9c5..02d4bcb 100644 --- a/apps/api/internal/store/admission_lock.go +++ b/apps/api/internal/store/admission_lock.go @@ -114,6 +114,14 @@ func admissionOperationLockKeys(input TaskAdmissionInput) []string { return normalizedAdmissionLockKeys(keys) } +func queuedAdmissionOperationLockKeys(input TaskAdmissionInput) []string { + keys := []string{"task-admission:" + input.TaskID} + for _, scope := range normalizedAdmissionScopes(input.Scopes) { + keys = append(keys, admissionQueueRegistrationLockKey(scope)) + } + return normalizedAdmissionLockKeys(keys) +} + func tryAdmissionTransactionLock(ctx context.Context, tx pgx.Tx, key string) error { // The process-local lock limits this wait to one PostgreSQL connection per // API or Worker process. Let PostgreSQL queue those process representatives diff --git a/apps/api/internal/store/admission_queue.go b/apps/api/internal/store/admission_queue.go index 0392c0c..18a080c 100644 --- a/apps/api/internal/store/admission_queue.go +++ b/apps/api/internal/store/admission_queue.go @@ -136,7 +136,7 @@ func (s *Store) QueueTaskAdmissionWithHook( if input.Mode != "async" { return TaskAdmission{}, errors.New("queued admission registration requires asynchronous mode") } - return retryAdmissionOperation(ctx, admissionOperationLockKeys(input), func() (TaskAdmission, error) { + return retryAdmissionOperation(ctx, queuedAdmissionOperationLockKeys(input), func() (TaskAdmission, error) { return s.queueTaskAdmissionOnce(ctx, input, onQueued) }) } @@ -188,8 +188,14 @@ WHERE id = $1::uuid`, input.TaskID).Scan(&taskActive); err != nil { AdmissionScope{ScopeType: "user_group", ScopeKey: admission.UserGroupID}, ) } + // Waiting registration has its own cross-process lock domain. Registrations + // still serialize with each other for exact queue-limit checks, but they do + // not wait behind a Worker admission transaction holding execution scope + // locks through synchronous replication. Moving one task from waiting to + // active preserves active+waiting, so concurrent dispatch cannot make this + // capacity check over-admit. for _, scope := range normalizedAdmissionScopes(lockScopes) { - if err := tryAdmissionTransactionLock(ctx, tx, admissionLockKey(scope)); err != nil { + if err := tryAdmissionTransactionLock(ctx, tx, admissionQueueRegistrationLockKey(scope)); err != nil { return TaskAdmission{}, err } } @@ -567,7 +573,7 @@ func (s *Store) RebindWaitingTaskAdmission(ctx context.Context, input TaskAdmiss if err := validateTaskAdmissionInput(input); err != nil { return TaskAdmission{}, err } - return retryAdmissionOperation(ctx, admissionOperationLockKeys(input), func() (TaskAdmission, error) { + return retryAdmissionOperation(ctx, queuedAdmissionOperationLockKeys(input), func() (TaskAdmission, error) { return s.rebindWaitingTaskAdmissionOnce(ctx, input) }) } @@ -598,7 +604,7 @@ func (s *Store) rebindWaitingTaskAdmissionOnce(ctx context.Context, input TaskAd AdmissionScope{ScopeType: "user_group", ScopeKey: current.UserGroupID}, ) for _, scope := range normalizedAdmissionScopes(lockScopes) { - if err := tryAdmissionTransactionLock(ctx, tx, admissionLockKey(scope)); err != nil { + if err := tryAdmissionTransactionLock(ctx, tx, admissionQueueRegistrationLockKey(scope)); err != nil { return TaskAdmission{}, err } } @@ -1099,6 +1105,10 @@ func admissionLockKey(scope AdmissionScope) string { return fmt.Sprintf("task-admission:%d:%s:%d:%s", len(scope.ScopeType), scope.ScopeType, len(scope.ScopeKey), scope.ScopeKey) } +func admissionQueueRegistrationLockKey(scope AdmissionScope) string { + return fmt.Sprintf("task-admission-queue:%d:%s:%d:%s", len(scope.ScopeType), scope.ScopeType, len(scope.ScopeKey), scope.ScopeKey) +} + type admissionScopeState struct { Scope AdmissionScope Active float64 diff --git a/apps/api/internal/store/admission_registration_lock_integration_test.go b/apps/api/internal/store/admission_registration_lock_integration_test.go new file mode 100644 index 0000000..ce7cc96 --- /dev/null +++ b/apps/api/internal/store/admission_registration_lock_integration_test.go @@ -0,0 +1,116 @@ +package store + +import ( + "context" + "os" + "strings" + "testing" + "time" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" + "github.com/google/uuid" +) + +func TestQueuedAdmissionRegistrationDoesNotWaitForExecutionScopeLock(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 admission registration PostgreSQL integration tests") + } + ctx, cancel := context.WithTimeout(context.Background(), 30*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: "registration-lock-test", + PlatformKey: "registration-lock-test-" + suffix, + Name: "Registration Lock Test " + suffix, + AuthType: "none", + Status: "enabled", + }) + if err != nil { + t.Fatalf("create platform: %v", err) + } + modelName := "registration-lock-model-" + suffix + 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, modelName).Scan(&platformModelID); err != nil { + t.Fatalf("create platform model: %v", err) + } + task, err := first.CreateTask(ctx, CreateTaskInput{ + Kind: "images.generate", + Model: modelName, + RunMode: "production", + Async: true, + Request: map[string]any{"prompt": "registration lock isolation"}, + }, &auth.User{ID: "registration-lock-test-" + suffix, Source: "gateway"}) + if err != nil { + t.Fatalf("create task: %v", err) + } + t.Cleanup(func() { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cleanupCancel() + _, _ = first.pool.Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = $1::uuid`, task.ID) + _, _ = 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) + }) + scope := AdmissionScope{ + ScopeType: "platform_model", + ScopeKey: platformModelID, + ScopeName: modelName, + ConcurrentLimit: 24, + Amount: 1, + LeaseTTLSeconds: 120, + QueueLimit: 100, + MaxWaitSeconds: 600, + } + input := TaskAdmissionInput{ + TaskID: task.ID, + PlatformID: platform.ID, + PlatformModelID: platformModelID, + QueueKey: platform.PlatformKey + ":image_generate:" + modelName, + Mode: "async", + Priority: 100, + Scopes: []AdmissionScope{scope}, + } + + executionLockTx, err := first.pool.Begin(ctx) + if err != nil { + t.Fatalf("begin execution lock transaction: %v", err) + } + if err := tryAdmissionTransactionLock(ctx, executionLockTx, admissionLockKey(scope)); err != nil { + _ = executionLockTx.Rollback(ctx) + t.Fatalf("hold execution scope lock: %v", err) + } + registrationCtx, registrationCancel := context.WithTimeout(ctx, 2*time.Second) + _, registrationErr := second.QueueTaskAdmissionWithHook(registrationCtx, input, nil) + registrationCancel() + if rollbackErr := executionLockTx.Rollback(ctx); rollbackErr != nil { + t.Fatalf("release execution scope lock: %v", rollbackErr) + } + if registrationErr != nil { + t.Fatalf("queue registration blocked behind execution scope lock: %v", registrationErr) + } +}