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。
This commit is contained in:
@@ -717,6 +717,23 @@ WHERE admission.task_id = task.id
|
|||||||
return RuntimeRecoveryResult{}, err
|
return RuntimeRecoveryResult{}, err
|
||||||
}
|
}
|
||||||
result.CleanedTaskAdmissions += tag.RowsAffected()
|
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)
|
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 {
|
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 (
|
INSERT INTO runtime_client_states (
|
||||||
client_id, platform_id, provider, method_name, queue_key, running_count, last_assigned_at
|
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.Provider,
|
||||||
candidate.ModelType,
|
candidate.ModelType,
|
||||||
candidate.QueueKey,
|
candidate.QueueKey,
|
||||||
)
|
); err != nil {
|
||||||
return err
|
return err
|
||||||
|
}
|
||||||
|
return tx.Commit(ctx)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (s *Store) RecordClientRelease(ctx context.Context, clientID string, lastError string) error {
|
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
|
UPDATE runtime_client_states
|
||||||
SET running_count = GREATEST(running_count - 1, 0),
|
SET running_count = GREATEST(running_count - 1, 0),
|
||||||
last_error = NULLIF($2::text, ''),
|
last_error = NULLIF($2::text, ''),
|
||||||
updated_at = now()
|
updated_at = now()
|
||||||
WHERE client_id = $1`, clientID, lastError)
|
WHERE client_id = $1`, clientID, lastError); err != nil {
|
||||||
return err
|
return err
|
||||||
|
}
|
||||||
|
return tx.Commit(ctx)
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user