perf(worker): 消除运行态计数跨城热行阻塞

将仅用于候选负载均衡的 runtime_client_states 分配和释放写入改为事务级异步提交，避免同步复制期间持续持有单候选热行锁。\n\n运行态计数仍由本地 WAL 持久化，并在运行时恢复中按 running attempt 重建；任务、账务、租约和状态提交不受影响。新增 PostgreSQL 并发与恢复集成测试。\n\n验证：go vet ./...；go test ./... -count=1；64 并发 PostgreSQL 集成测试；迁移安全检查；gofmt；git diff --check。
2026-07-31 09:14:13 +08:00
parent b0dac9f8ec
commit 11ec057e6a
3 changed files with 173 additions and 6 deletions
@@ -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)
 }
@@ -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)
+	}
+}
@@ -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)
 }