wangbo
335b21d1f8
将跨实例探测协调从持有会话级 advisory lock 改为带过期时间的数据库短租约，避免网络探测期间占用关键 PostgreSQL 连接。新增可配置的探测并发上限和真实 PostgreSQL 回归测试，确保租约生效期间连接池仍可全部获取。\n\n验证：go test ./... -count=1；go vet ./...；真实 PostgreSQL 17 集成测试；迁移安全检查。
63 lines
1.8 KiB
Go
63 lines
1.8 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
func TestExecutionPoolQueriesUseTimestampCutoff(t *testing.T) {
|
|
db := billingV2IntegrationStore(t)
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
if _, err := db.ListWorkers(ctx, time.Now()); err != nil {
|
|
t.Fatalf("list workers with timestamp cutoff: %v", err)
|
|
}
|
|
if _, err := db.ListCapacity(ctx, time.Now()); err != nil {
|
|
t.Fatalf("list capacity with timestamp cutoff: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestRouteProbeLeaseDoesNotPinDatabaseConnection(t *testing.T) {
|
|
db := billingV2IntegrationStore(t)
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
if err := db.UpsertExecutionPool(ctx, executionpool.ExecutionPool{ID: "probe-pool", State: executionpool.PoolActive}); err != nil {
|
|
t.Fatalf("upsert execution pool: %v", err)
|
|
}
|
|
token, claimed, err := db.TryClaimRouteProbe(ctx, "probe-pool", "route-profile", 15*time.Second)
|
|
if err != nil {
|
|
t.Fatalf("claim route probe: %v", err)
|
|
}
|
|
if !claimed || token == "" {
|
|
t.Fatal("route probe was not claimed")
|
|
}
|
|
|
|
connections := make([]*pgxpool.Conn, 0, db.Pool().Stat().MaxConns())
|
|
defer func() {
|
|
for _, connection := range connections {
|
|
connection.Release()
|
|
}
|
|
}()
|
|
for range db.Pool().Stat().MaxConns() {
|
|
connection, err := db.Pool().Acquire(ctx)
|
|
if err != nil {
|
|
t.Fatalf("acquire all pool connections while lease is active: %v", err)
|
|
}
|
|
connections = append(connections, connection)
|
|
}
|
|
for _, connection := range connections {
|
|
connection.Release()
|
|
}
|
|
connections = connections[:0]
|
|
|
|
if err := db.ReleaseRouteProbe(ctx, "probe-pool", "route-profile", token); err != nil {
|
|
t.Fatalf("release route probe: %v", err)
|
|
}
|
|
}
|