feat(worker): 实现集群限流与自适应负载
保留平台模型 RPM、TPM 和并发策略语义，增加 PostgreSQL 集群级租约、饱和候选重选和多平台自动负载，避免突发任务固定等待首个平台。\n\n新增 Worker 实时负载采样、自适应 active/heavy 容量、心跳与管理端指标，并扩展本地 acceptance runner，覆盖三 Worker、同模型三平台 2/4/6 并发和 48 个带图视频突发任务。\n\n验证：go test ./...、go vet ./...、PostgreSQL 跨 Store 集成测试、gofmt、bash -n、ShellCheck 及本地集群 provider-burst 验收通过；48/48 成功，无越限、重复提交、重复计费、重复回调或终态资源泄漏。
This commit is contained in:
@@ -98,23 +98,168 @@ WHERE scope_type = 'platform_model'
|
||||
}
|
||||
|
||||
var active int64
|
||||
var storedLimit float64
|
||||
if err := first.Pool().QueryRow(ctx, `
|
||||
SELECT COUNT(*)
|
||||
SELECT COUNT(*), COALESCE(MAX(limit_value), 0)::float8
|
||||
FROM gateway_concurrency_leases
|
||||
WHERE scope_type = 'platform_model'
|
||||
AND scope_key = $1
|
||||
AND released_at IS NULL
|
||||
AND expires_at > now()`, scopeKey).Scan(&active); err != nil {
|
||||
AND expires_at > now()`, scopeKey).Scan(&active, &storedLimit); err != nil {
|
||||
t.Fatalf("count active leases: %v", err)
|
||||
}
|
||||
if successes.Load() != 64 || active != 64 {
|
||||
t.Fatalf("successful reservations=%d active leases=%d, want exactly 64", successes.Load(), active)
|
||||
if successes.Load() != 64 || active != 64 || storedLimit != 64 {
|
||||
t.Fatalf("successful reservations=%d active leases=%d stored limit=%.0f, want exactly 64", successes.Load(), active, storedLimit)
|
||||
}
|
||||
if peak.Load() > 64 {
|
||||
t.Fatalf("active lease peak=%d, want <=64", peak.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrencyLeaseTimestampStartsAtReservationStatement(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 concurrency lease PostgreSQL integration tests")
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
defer cancel()
|
||||
database, err := Connect(ctx, databaseURL)
|
||||
if err != nil {
|
||||
t.Fatalf("connect store: %v", err)
|
||||
}
|
||||
defer database.Close()
|
||||
|
||||
scopeKey := "statement-clock-" + time.Now().UTC().Format("20060102150405.000000000")
|
||||
taskIDs := createLeaseTestTasks(t, ctx, database, 1, scopeKey)
|
||||
defer deleteLeaseTestTasks(t, database, taskIDs)
|
||||
|
||||
tx, err := database.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
t.Fatalf("begin reservation transaction: %v", err)
|
||||
}
|
||||
defer rollbackTransaction(tx)
|
||||
var transactionStartedAt time.Time
|
||||
if err := tx.QueryRow(ctx, `SELECT now()`).Scan(&transactionStartedAt); err != nil {
|
||||
t.Fatalf("read transaction start: %v", err)
|
||||
}
|
||||
time.Sleep(1100 * time.Millisecond)
|
||||
lease, err := reserveConcurrencyLease(ctx, tx, taskIDs[0], "", RateLimitReservation{
|
||||
ScopeType: "platform_model",
|
||||
ScopeKey: scopeKey,
|
||||
Metric: "concurrent",
|
||||
Limit: 1,
|
||||
Amount: 1,
|
||||
LeaseTTLSeconds: 30,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("reserve concurrency lease: %v", err)
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
t.Fatalf("commit reservation transaction: %v", err)
|
||||
}
|
||||
|
||||
var acquiredAt, expiresAt time.Time
|
||||
if err := database.pool.QueryRow(ctx, `
|
||||
SELECT acquired_at, expires_at
|
||||
FROM gateway_concurrency_leases
|
||||
WHERE id = $1::uuid`, lease.ID).Scan(&acquiredAt, &expiresAt); err != nil {
|
||||
t.Fatalf("read lease timestamps: %v", err)
|
||||
}
|
||||
if elapsed := acquiredAt.Sub(transactionStartedAt); elapsed < time.Second {
|
||||
t.Fatalf("lease acquired_at advanced by %s, want at least 1s after transaction start", elapsed)
|
||||
}
|
||||
if ttl := expiresAt.Sub(acquiredAt); ttl != 30*time.Second {
|
||||
t.Fatalf("lease ttl=%s, want 30s", ttl)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCounterWindowReservationIsAtomicAcrossPools(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 rate limit PostgreSQL integration tests")
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second)
|
||||
defer cancel()
|
||||
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()
|
||||
|
||||
tests := []struct {
|
||||
metric string
|
||||
limit float64
|
||||
amount float64
|
||||
wantSuccesses int64
|
||||
}{
|
||||
{metric: "rpm", limit: 37, amount: 1, wantSuccesses: 37},
|
||||
{metric: "tpm_total", limit: 1_000, amount: 25, wantSuccesses: 40},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.metric, func(t *testing.T) {
|
||||
scopeKey := "atomic-" + test.metric + "-" + time.Now().UTC().Format("20060102150405.000000000")
|
||||
taskIDs := createLeaseTestTasks(t, ctx, first, 128, scopeKey)
|
||||
defer deleteLeaseTestTasks(t, first, taskIDs)
|
||||
|
||||
var successes atomic.Int64
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, len(taskIDs))
|
||||
for index, taskID := range taskIDs {
|
||||
wg.Add(1)
|
||||
go func(index int, taskID string) {
|
||||
defer wg.Done()
|
||||
target := first
|
||||
if index%2 == 1 {
|
||||
target = second
|
||||
}
|
||||
_, err := target.ReserveRateLimits(ctx, taskID, "", []RateLimitReservation{{
|
||||
ScopeType: "platform_model",
|
||||
ScopeKey: scopeKey,
|
||||
Metric: test.metric,
|
||||
Limit: test.limit,
|
||||
Amount: test.amount,
|
||||
WindowSeconds: 3600,
|
||||
}})
|
||||
if err == nil {
|
||||
successes.Add(1)
|
||||
return
|
||||
}
|
||||
if !errors.Is(err, ErrRateLimited) {
|
||||
errs <- err
|
||||
}
|
||||
}(index, taskID)
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
t.Fatalf("unexpected reservation error: %v", err)
|
||||
}
|
||||
|
||||
var current float64
|
||||
if err := first.Pool().QueryRow(ctx, `
|
||||
SELECT COALESCE(MAX(used_value + reserved_value), 0)::float8
|
||||
FROM gateway_rate_limit_counters
|
||||
WHERE scope_type = 'platform_model'
|
||||
AND scope_key = $1
|
||||
AND metric = $2`, scopeKey, test.metric).Scan(¤t); err != nil {
|
||||
t.Fatalf("read %s counter: %v", test.metric, err)
|
||||
}
|
||||
if successes.Load() != test.wantSuccesses {
|
||||
t.Fatalf("successful %s reservations=%d, want exactly %d", test.metric, successes.Load(), test.wantSuccesses)
|
||||
}
|
||||
wantCurrent := float64(test.wantSuccesses) * test.amount
|
||||
if current != wantCurrent || current > test.limit {
|
||||
t.Fatalf("%s current=%.0f, want %.0f and <= %.0f", test.metric, current, wantCurrent, test.limit)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrencyLeaseRenewalExtendsAndReleases(t *testing.T) {
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
|
||||
Reference in New Issue
Block a user