From 11ec057e6a039444c473e92c9b8b8fcf783ad8b4 Mon Sep 17 00:00:00 2001 From: wangbo Date: Fri, 31 Jul 2026 09:14:13 +0800 Subject: [PATCH] =?UTF-8?q?perf(worker):=20=E6=B6=88=E9=99=A4=E8=BF=90?= =?UTF-8?q?=E8=A1=8C=E6=80=81=E8=AE=A1=E6=95=B0=E8=B7=A8=E5=9F=8E=E7=83=AD?= =?UTF-8?q?=E8=A1=8C=E9=98=BB=E5=A1=9E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 将仅用于候选负载均衡的 runtime_client_states 分配和释放写入改为事务级异步提交，避免同步复制期间持续持有单候选热行锁。\n\n运行态计数仍由本地 WAL 持久化，并在运行时恢复中按 running attempt 重建；任务、账务、租约和状态提交不受影响。新增 PostgreSQL 并发与恢复集成测试。\n\n验证：go vet ./...；go test ./... -count=1；64 并发 PostgreSQL 集成测试；迁移安全检查；gofmt；git diff --check。 --- apps/api/internal/store/rate_limits.go | 17 +++ .../runtime_client_state_integration_test.go | 123 ++++++++++++++++++ apps/api/internal/store/tasks_runtime.go | 39 +++++- 3 files changed, 173 insertions(+), 6 deletions(-) create mode 100644 apps/api/internal/store/runtime_client_state_integration_test.go diff --git a/apps/api/internal/store/rate_limits.go b/apps/api/internal/store/rate_limits.go index baa006d..58d209e 100644 --- a/apps/api/internal/store/rate_limits.go +++ b/apps/api/internal/store/rate_limits.go @@ -717,6 +717,23 @@ WHERE admission.task_id = task.id return RuntimeRecoveryResult{}, err } result.CleanedTaskAdmissions += tag.RowsAffected() + if _, err := tx.Exec(ctx, ` +UPDATE runtime_client_states state +SET running_count = ( + SELECT count(*) + FROM gateway_task_attempts attempt + WHERE attempt.client_id = state.client_id + AND attempt.status = 'running' + ), + updated_at = now() +WHERE state.running_count IS DISTINCT FROM ( + SELECT count(*) + FROM gateway_task_attempts attempt + WHERE attempt.client_id = state.client_id + AND attempt.status = 'running' +)`); err != nil { + return RuntimeRecoveryResult{}, err + } return result, tx.Commit(ctx) } diff --git a/apps/api/internal/store/runtime_client_state_integration_test.go b/apps/api/internal/store/runtime_client_state_integration_test.go new file mode 100644 index 0000000..1805992 --- /dev/null +++ b/apps/api/internal/store/runtime_client_state_integration_test.go @@ -0,0 +1,123 @@ +package store + +import ( + "context" + "os" + "strings" + "sync" + "testing" + "time" + + "github.com/google/uuid" +) + +func TestRuntimeClientStateConcurrentUpdatesAndRecovery(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 runtime client state PostgreSQL integration tests") + } + ctx, cancel := context.WithTimeout(context.Background(), 60*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) + } + + suffix := strings.ReplaceAll(uuid.NewString(), "-", "") + platform, err := db.CreatePlatform(ctx, CreatePlatformInput{ + Provider: "runtime-state-test", + PlatformKey: "runtime-state-test-" + suffix, + Name: "Runtime State Test " + suffix, + AuthType: "none", + Status: "enabled", + }) + if err != nil { + t.Fatalf("create platform: %v", err) + } + candidate := RuntimeModelCandidate{ + PlatformID: platform.ID, + ClientID: platform.PlatformKey + ":image_generate:model", + QueueKey: platform.PlatformKey + ":image_generate:model", + Provider: platform.Provider, + ModelType: "image_generate", + } + t.Cleanup(func() { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cleanupCancel() + _, _ = db.pool.Exec(cleanupCtx, `DELETE FROM runtime_client_states WHERE client_id = $1`, candidate.ClientID) + _, _ = db.pool.Exec(cleanupCtx, `DELETE FROM integration_platforms WHERE id = $1::uuid`, platform.ID) + }) + + runConcurrent := func(operation func() error) { + t.Helper() + const operations = 64 + errs := make(chan error, operations) + var wg sync.WaitGroup + for index := 0; index < operations; index++ { + wg.Add(1) + go func() { + defer wg.Done() + errs <- operation() + }() + } + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + t.Fatalf("runtime client state operation: %v", err) + } + } + } + runningCount := func() int { + t.Helper() + var count int + if err := db.pool.QueryRow(ctx, ` +SELECT running_count +FROM runtime_client_states +WHERE client_id = $1`, candidate.ClientID).Scan(&count); err != nil { + t.Fatalf("read runtime client state: %v", err) + } + return count + } + + runConcurrent(func() error { + return db.RecordClientAssignment(ctx, candidate) + }) + if count := runningCount(); count != 64 { + t.Fatalf("running count after assignments = %d, want 64", count) + } + runConcurrent(func() error { + return db.RecordClientRelease(ctx, candidate.ClientID, "") + }) + if count := runningCount(); count != 0 { + t.Fatalf("running count after releases = %d, want 0", count) + } + if _, err := db.pool.Exec(ctx, ` +UPDATE runtime_client_states +SET running_count = 99 +WHERE client_id = $1`, candidate.ClientID); err != nil { + t.Fatalf("corrupt runtime client state: %v", err) + } + if _, err := db.RecoverInterruptedRuntimeState(ctx); err != nil { + t.Fatalf("recover runtime client state: %v", err) + } + if count := runningCount(); count != 0 { + t.Fatalf("recovered running count = %d, want 0", count) + } + var synchronousCommit string + if err := db.pool.QueryRow(ctx, `SHOW synchronous_commit`).Scan(&synchronousCommit); err != nil { + t.Fatalf("read synchronous_commit: %v", err) + } + if synchronousCommit != "on" { + t.Fatalf("session synchronous_commit = %q, want on", synchronousCommit) + } +} diff --git a/apps/api/internal/store/tasks_runtime.go b/apps/api/internal/store/tasks_runtime.go index fd594c0..6e048c1 100644 --- a/apps/api/internal/store/tasks_runtime.go +++ b/apps/api/internal/store/tasks_runtime.go @@ -2156,7 +2156,19 @@ ON CONFLICT (task_id, seq, callback_url) DO NOTHING`, } func (s *Store) RecordClientAssignment(ctx context.Context, candidate RuntimeModelCandidate) error { - _, err := s.pool.Exec(ctx, ` + tx, err := s.pool.Begin(ctx) + if err != nil { + return err + } + defer rollbackTransaction(tx) + // This counter only influences candidate load balancing and is rebuilt from + // active attempts during runtime recovery. Keeping a synchronous-replication + // commit open while updating one row per upstream client serializes every + // assignment for a hot model across regions. + if _, err := tx.Exec(ctx, `SET LOCAL synchronous_commit = off`); err != nil { + return err + } + if _, err := tx.Exec(ctx, ` INSERT INTO runtime_client_states ( client_id, platform_id, provider, method_name, queue_key, running_count, last_assigned_at ) @@ -2170,16 +2182,31 @@ SET running_count = runtime_client_states.running_count + 1, candidate.Provider, candidate.ModelType, candidate.QueueKey, - ) - return err + ); err != nil { + return err + } + return tx.Commit(ctx) } func (s *Store) RecordClientRelease(ctx context.Context, clientID string, lastError string) error { - _, err := s.pool.Exec(ctx, ` + tx, err := s.pool.Begin(ctx) + if err != nil { + return err + } + defer rollbackTransaction(tx) + // runtime_client_states is derived scheduler state; see + // RecordClientAssignment for why it must not hold its hot row lock through + // a cross-region synchronous replication wait. + if _, err := tx.Exec(ctx, `SET LOCAL synchronous_commit = off`); err != nil { + return err + } + if _, err := tx.Exec(ctx, ` UPDATE runtime_client_states SET running_count = GREATEST(running_count - 1, 0), last_error = NULLIF($2::text, ''), updated_at = now() -WHERE client_id = $1`, clientID, lastError) - return err +WHERE client_id = $1`, clientID, lastError); err != nil { + return err + } + return tx.Commit(ctx) }