BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf(worker): 隔离异步等待登记与执行锁

Browse Source 
为异步等待队列登记增加独立的进程内和 PostgreSQL advisory lock 域，使严格队列上限校验不再等待 Worker 执行槽事务的跨城同步提交。任务级锁保持不变，调度期间 active 与 waiting 总量守恒。\n\n新增 PostgreSQL 集成测试，验证执行 scope 锁被占用时等待登记仍可独立完成。\n\n验证：go vet ./...；go test ./... -count=1；隔离 PostgreSQL 集成测试；迁移安全检查；gofmt；git diff --check。
This commit is contained in:
wangbo
2026-07-31 09:29:01 +08:00
parent 11ec057e6a
commit 387bf5842d
3 changed files with 138 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/admission_lock.go
+8
View File
@@ -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
apps/api/internal/store/admission_queue.go
+14 -4
View File
@@ -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
apps/api/internal/store/admission_registration_lock_integration_test.go
+116
View File
@@ -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)
}
}