wangbo
c28bf74230
保留平台模型 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 成功，无越限、重复提交、重复计费、重复回调或终态资源泄漏。
344 lines
10 KiB
Go
344 lines
10 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"os"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestConcurrencyLeaseReservationIsAtomicAcrossPools(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(), 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()
|
|
|
|
scopeKey := "atomic-" + time.Now().UTC().Format("20060102150405.000000000")
|
|
taskIDs := createLeaseTestTasks(t, ctx, first, 256, scopeKey)
|
|
defer deleteLeaseTestTasks(t, first, taskIDs)
|
|
|
|
var successes atomic.Int64
|
|
var peak atomic.Int64
|
|
monitorCtx, stopMonitor := context.WithCancel(ctx)
|
|
var monitorWG sync.WaitGroup
|
|
monitorWG.Add(1)
|
|
go func() {
|
|
defer monitorWG.Done()
|
|
ticker := time.NewTicker(2 * time.Millisecond)
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-monitorCtx.Done():
|
|
return
|
|
case <-ticker.C:
|
|
var active int64
|
|
if err := first.Pool().QueryRow(monitorCtx, `
|
|
SELECT COUNT(*)
|
|
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 {
|
|
for active > peak.Load() && !peak.CompareAndSwap(peak.Load(), active) {
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
|
|
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: "concurrent",
|
|
Limit: 64,
|
|
Amount: 1,
|
|
LeaseTTLSeconds: 30,
|
|
}})
|
|
if err == nil {
|
|
successes.Add(1)
|
|
return
|
|
}
|
|
if !errors.Is(err, ErrRateLimited) {
|
|
errs <- err
|
|
}
|
|
}(index, taskID)
|
|
}
|
|
wg.Wait()
|
|
stopMonitor()
|
|
monitorWG.Wait()
|
|
close(errs)
|
|
for err := range errs {
|
|
t.Fatalf("unexpected reservation error: %v", err)
|
|
}
|
|
|
|
var active int64
|
|
var storedLimit float64
|
|
if err := first.Pool().QueryRow(ctx, `
|
|
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, &storedLimit); err != nil {
|
|
t.Fatalf("count active leases: %v", err)
|
|
}
|
|
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 == "" {
|
|
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()
|
|
db, err := Connect(ctx, databaseURL)
|
|
if err != nil {
|
|
t.Fatalf("connect store: %v", err)
|
|
}
|
|
defer db.Close()
|
|
scopeKey := "renew-" + time.Now().UTC().Format("20060102150405.000000000")
|
|
taskIDs := createLeaseTestTasks(t, ctx, db, 1, scopeKey)
|
|
defer deleteLeaseTestTasks(t, db, taskIDs)
|
|
|
|
result, err := db.ReserveRateLimits(ctx, taskIDs[0], "", []RateLimitReservation{{
|
|
ScopeType: "platform_model",
|
|
ScopeKey: scopeKey,
|
|
Metric: "concurrent",
|
|
Limit: 1,
|
|
Amount: 1,
|
|
LeaseTTLSeconds: 2,
|
|
}})
|
|
if err != nil {
|
|
t.Fatalf("reserve short lease: %v", err)
|
|
}
|
|
time.Sleep(time.Second)
|
|
if err := db.RenewConcurrencyLeases(ctx, result.Leases); err != nil {
|
|
t.Fatalf("renew short lease: %v", err)
|
|
}
|
|
time.Sleep(1500 * time.Millisecond)
|
|
var active bool
|
|
if err := db.Pool().QueryRow(ctx, `
|
|
SELECT EXISTS (
|
|
SELECT 1 FROM gateway_concurrency_leases
|
|
WHERE id = $1::uuid AND released_at IS NULL AND expires_at > now()
|
|
)`, result.Leases[0].ID).Scan(&active); err != nil {
|
|
t.Fatalf("read renewed lease: %v", err)
|
|
}
|
|
if !active {
|
|
t.Fatal("renewed lease expired at its original TTL")
|
|
}
|
|
if err := db.ReleaseConcurrencyLeases(ctx, result.Leases); err != nil {
|
|
t.Fatalf("release renewed lease: %v", err)
|
|
}
|
|
}
|
|
|
|
func createLeaseTestTasks(t *testing.T, ctx context.Context, db *Store, count int, marker string) []string {
|
|
t.Helper()
|
|
rows, err := db.Pool().Query(ctx, `
|
|
INSERT INTO gateway_tasks (kind, run_mode, user_id, model, model_type, request, status, queue_key)
|
|
SELECT 'lease-test', 'simulation', $2, 'lease-test', 'text_generate', '{}'::jsonb, 'queued', $2
|
|
FROM generate_series(1, $1)
|
|
RETURNING id::text`, count, marker)
|
|
if err != nil {
|
|
t.Fatalf("create lease test tasks: %v", err)
|
|
}
|
|
defer rows.Close()
|
|
ids := make([]string, 0, count)
|
|
for rows.Next() {
|
|
var id string
|
|
if err := rows.Scan(&id); err != nil {
|
|
t.Fatalf("scan lease test task: %v", err)
|
|
}
|
|
ids = append(ids, id)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
t.Fatalf("create lease test tasks: %v", err)
|
|
}
|
|
return ids
|
|
}
|
|
|
|
func deleteLeaseTestTasks(t *testing.T, db *Store, taskIDs []string) {
|
|
t.Helper()
|
|
cleanupCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
if _, err := db.Pool().Exec(cleanupCtx, `DELETE FROM gateway_tasks WHERE id = ANY($1::uuid[])`, taskIDs); err != nil {
|
|
t.Errorf("delete lease test tasks: %v", err)
|
|
}
|
|
}
|