fix(routing): 避免路由探测耗尽数据库连接
将跨实例探测协调从持有会话级 advisory lock 改为带过期时间的数据库短租约,避免网络探测期间占用关键 PostgreSQL 连接。新增可配置的探测并发上限和真实 PostgreSQL 回归测试,确保租约生效期间连接池仍可全部获取。\n\n验证:go test ./... -count=1;go vet ./...;真实 PostgreSQL 17 集成测试;迁移安全检查。
This commit is contained in:
@@ -101,6 +101,7 @@ type Config struct {
|
|||||||
RouteProbeHotIntervalSeconds int
|
RouteProbeHotIntervalSeconds int
|
||||||
RouteProbeColdIntervalSeconds int
|
RouteProbeColdIntervalSeconds int
|
||||||
RouteProbeTimeoutMS int
|
RouteProbeTimeoutMS int
|
||||||
|
RouteProbeConcurrency int
|
||||||
WorkerAutoscalingEnabled bool
|
WorkerAutoscalingEnabled bool
|
||||||
CapacityOrchestratorAdapter string
|
CapacityOrchestratorAdapter string
|
||||||
CapacityPoolsJSON string
|
CapacityPoolsJSON string
|
||||||
@@ -223,6 +224,7 @@ func Load() Config {
|
|||||||
RouteProbeHotIntervalSeconds: envIntValidated("AI_GATEWAY_ROUTE_PROBE_HOT_INTERVAL_SECONDS", 15),
|
RouteProbeHotIntervalSeconds: envIntValidated("AI_GATEWAY_ROUTE_PROBE_HOT_INTERVAL_SECONDS", 15),
|
||||||
RouteProbeColdIntervalSeconds: envIntValidated("AI_GATEWAY_ROUTE_PROBE_COLD_INTERVAL_SECONDS", 60),
|
RouteProbeColdIntervalSeconds: envIntValidated("AI_GATEWAY_ROUTE_PROBE_COLD_INTERVAL_SECONDS", 60),
|
||||||
RouteProbeTimeoutMS: envIntValidated("AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS", 3000),
|
RouteProbeTimeoutMS: envIntValidated("AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS", 3000),
|
||||||
|
RouteProbeConcurrency: envIntValidated("AI_GATEWAY_ROUTE_PROBE_CONCURRENCY", 8),
|
||||||
WorkerAutoscalingEnabled: env("AI_GATEWAY_WORKER_AUTOSCALING_ENABLED", "false") == "true",
|
WorkerAutoscalingEnabled: env("AI_GATEWAY_WORKER_AUTOSCALING_ENABLED", "false") == "true",
|
||||||
CapacityOrchestratorAdapter: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_CAPACITY_ORCHESTRATOR_ADAPTER", "kubernetes"))),
|
CapacityOrchestratorAdapter: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_CAPACITY_ORCHESTRATOR_ADAPTER", "kubernetes"))),
|
||||||
CapacityPoolsJSON: strings.TrimSpace(env("AI_GATEWAY_CAPACITY_POOLS", "")),
|
CapacityPoolsJSON: strings.TrimSpace(env("AI_GATEWAY_CAPACITY_POOLS", "")),
|
||||||
@@ -345,6 +347,9 @@ func (c Config) Validate() error {
|
|||||||
if c.RouteProbeTimeoutMS != 0 && (c.RouteProbeTimeoutMS < 250 || c.RouteProbeTimeoutMS > 10000) {
|
if c.RouteProbeTimeoutMS != 0 && (c.RouteProbeTimeoutMS < 250 || c.RouteProbeTimeoutMS > 10000) {
|
||||||
return errors.New("AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS must be between 250 and 10000")
|
return errors.New("AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS must be between 250 and 10000")
|
||||||
}
|
}
|
||||||
|
if c.RouteProbeConcurrency != 0 && (c.RouteProbeConcurrency < 1 || c.RouteProbeConcurrency > 64) {
|
||||||
|
return errors.New("AI_GATEWAY_ROUTE_PROBE_CONCURRENCY must be between 1 and 64")
|
||||||
|
}
|
||||||
if strings.EqualFold(c.RoutingMode, "enforced") && len(c.WorkerExecutionSecret) < 32 {
|
if strings.EqualFold(c.RoutingMode, "enforced") && len(c.WorkerExecutionSecret) < 32 {
|
||||||
return errors.New("AI_GATEWAY_WORKER_EXECUTION_SECRET must be at least 32 bytes in enforced routing mode")
|
return errors.New("AI_GATEWAY_WORKER_EXECUTION_SECRET must be at least 32 bytes in enforced routing mode")
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -27,6 +27,11 @@ func (s *Service) runRouteHealthProber(ctx context.Context) {
|
|||||||
ticker := time.NewTicker(time.Second)
|
ticker := time.NewTicker(time.Second)
|
||||||
defer ticker.Stop()
|
defer ticker.Stop()
|
||||||
states := make(map[string]routeProbeState)
|
states := make(map[string]routeProbeState)
|
||||||
|
probeConcurrency := s.cfg.RouteProbeConcurrency
|
||||||
|
if probeConcurrency <= 0 {
|
||||||
|
probeConcurrency = 8
|
||||||
|
}
|
||||||
|
probeSlots := make(chan struct{}, probeConcurrency)
|
||||||
for {
|
for {
|
||||||
select {
|
select {
|
||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
@@ -58,19 +63,35 @@ func (s *Service) runRouteHealthProber(ctx context.Context) {
|
|||||||
if target.Hot {
|
if target.Hot {
|
||||||
interval = time.Duration(s.cfg.RouteProbeHotIntervalSeconds) * time.Second
|
interval = time.Duration(s.cfg.RouteProbeHotIntervalSeconds) * time.Second
|
||||||
}
|
}
|
||||||
|
select {
|
||||||
|
case probeSlots <- struct{}{}:
|
||||||
states[profile.Key] = routeProbeState{next: now.Add(interval), lastRequestObserved: requestedAt}
|
states[profile.Key] = routeProbeState{next: now.Add(interval), lastRequestObserved: requestedAt}
|
||||||
go s.probeRouteTarget(ctx, target, profile)
|
case <-ctx.Done():
|
||||||
|
return
|
||||||
|
default:
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
go func(target store.RouteProbeTarget, profile executionpool.RouteProfile) {
|
||||||
|
defer func() { <-probeSlots }()
|
||||||
|
s.probeRouteTarget(ctx, target, profile)
|
||||||
|
}(target, profile)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (s *Service) probeRouteTarget(ctx context.Context, target store.RouteProbeTarget, profile executionpool.RouteProfile) {
|
func (s *Service) probeRouteTarget(ctx context.Context, target store.RouteProbeTarget, profile executionpool.RouteProfile) {
|
||||||
leadership, locked, err := s.coordinationStore.TryAcquireRouteProbeLeadership(ctx, s.cfg.ExecutionPoolID, profile.Key)
|
leaseToken, locked, err := s.coordinationStore.TryClaimRouteProbe(ctx, s.cfg.ExecutionPoolID, profile.Key, time.Duration(s.cfg.RouteProbeTimeoutMS)*time.Millisecond+5*time.Second)
|
||||||
if err != nil || !locked {
|
if err != nil || !locked {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
defer leadership.Release()
|
defer func() {
|
||||||
|
releaseCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
if err := s.coordinationStore.ReleaseRouteProbe(releaseCtx, s.cfg.ExecutionPoolID, profile.Key, leaseToken); err != nil {
|
||||||
|
s.logRouteProbeFailure("release_lease", err)
|
||||||
|
}
|
||||||
|
}()
|
||||||
client, err := s.httpClientForCandidate(target.Candidate, false, "")
|
client, err := s.httpClientForCandidate(target.Candidate, false, "")
|
||||||
if err != nil {
|
if err != nil {
|
||||||
s.logRouteProbeFailure("prepare_client", err)
|
s.logRouteProbeFailure("prepare_client", err)
|
||||||
|
|||||||
@@ -222,6 +222,49 @@ WHERE expires_at > now()`)
|
|||||||
return items, rows.Err()
|
return items, rows.Err()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TryClaimRouteProbe coordinates probes without holding a database connection
|
||||||
|
// during DNS/TCP/TLS/HTTP I/O. The previous session advisory lock pinned one
|
||||||
|
// critical-pool connection per in-flight probe and could starve readiness.
|
||||||
|
func (s *Store) TryClaimRouteProbe(ctx context.Context, poolID, routeProfileKey string, ttl time.Duration) (string, bool, error) {
|
||||||
|
poolID = strings.TrimSpace(poolID)
|
||||||
|
routeProfileKey = strings.TrimSpace(routeProfileKey)
|
||||||
|
if poolID == "" || routeProfileKey == "" {
|
||||||
|
return "", false, errors.New("route probe pool and profile are required")
|
||||||
|
}
|
||||||
|
if ttl <= 0 {
|
||||||
|
ttl = 15 * time.Second
|
||||||
|
}
|
||||||
|
token := uuid.NewString()
|
||||||
|
var claimed string
|
||||||
|
err := s.pool.QueryRow(ctx, `
|
||||||
|
INSERT INTO gateway_route_probe_leases (
|
||||||
|
pool_id, route_profile_key, lease_token, expires_at, updated_at
|
||||||
|
)
|
||||||
|
VALUES ($1, $2, $3::uuid, now() + $4::interval, now())
|
||||||
|
ON CONFLICT (pool_id, route_profile_key) DO UPDATE
|
||||||
|
SET lease_token = EXCLUDED.lease_token,
|
||||||
|
expires_at = EXCLUDED.expires_at,
|
||||||
|
updated_at = now()
|
||||||
|
WHERE gateway_route_probe_leases.expires_at <= now()
|
||||||
|
RETURNING lease_token::text`, poolID, routeProfileKey, token, ttl.String()).Scan(&claimed)
|
||||||
|
if errors.Is(err, pgx.ErrNoRows) {
|
||||||
|
return "", false, nil
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return "", false, err
|
||||||
|
}
|
||||||
|
return claimed, true, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *Store) ReleaseRouteProbe(ctx context.Context, poolID, routeProfileKey, token string) error {
|
||||||
|
_, err := s.pool.Exec(ctx, `
|
||||||
|
DELETE FROM gateway_route_probe_leases
|
||||||
|
WHERE pool_id = $1
|
||||||
|
AND route_profile_key = $2
|
||||||
|
AND lease_token = $3::uuid`, strings.TrimSpace(poolID), strings.TrimSpace(routeProfileKey), strings.TrimSpace(token))
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
func (s *Store) ListRouteHealth(ctx context.Context, routeProfileKey string, _ time.Time) ([]executionpool.RouteHealth, error) {
|
func (s *Store) ListRouteHealth(ctx context.Context, routeProfileKey string, _ time.Time) ([]executionpool.RouteHealth, error) {
|
||||||
return s.listRouteHealth(ctx, "WHERE route_profile_key = $1", strings.TrimSpace(routeProfileKey))
|
return s.listRouteHealth(ctx, "WHERE route_profile_key = $1", strings.TrimSpace(routeProfileKey))
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -4,6 +4,9 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
|
||||||
|
"github.com/jackc/pgx/v5/pgxpool"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestExecutionPoolQueriesUseTimestampCutoff(t *testing.T) {
|
func TestExecutionPoolQueriesUseTimestampCutoff(t *testing.T) {
|
||||||
@@ -18,3 +21,42 @@ func TestExecutionPoolQueriesUseTimestampCutoff(t *testing.T) {
|
|||||||
t.Fatalf("list capacity with timestamp cutoff: %v", err)
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,11 @@
|
|||||||
|
CREATE TABLE IF NOT EXISTS gateway_route_probe_leases (
|
||||||
|
pool_id text NOT NULL REFERENCES gateway_execution_pools(pool_id) ON DELETE CASCADE,
|
||||||
|
route_profile_key text NOT NULL,
|
||||||
|
lease_token uuid NOT NULL,
|
||||||
|
expires_at timestamptz NOT NULL,
|
||||||
|
updated_at timestamptz NOT NULL DEFAULT now(),
|
||||||
|
PRIMARY KEY (pool_id, route_profile_key)
|
||||||
|
);
|
||||||
|
|
||||||
|
CREATE INDEX IF NOT EXISTS idx_gateway_route_probe_leases_expiry
|
||||||
|
ON gateway_route_probe_leases(expires_at);
|
||||||
@@ -52,6 +52,7 @@ data:
|
|||||||
AI_GATEWAY_ROUTE_PROBE_HOT_INTERVAL_SECONDS: "15"
|
AI_GATEWAY_ROUTE_PROBE_HOT_INTERVAL_SECONDS: "15"
|
||||||
AI_GATEWAY_ROUTE_PROBE_COLD_INTERVAL_SECONDS: "60"
|
AI_GATEWAY_ROUTE_PROBE_COLD_INTERVAL_SECONDS: "60"
|
||||||
AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS: "3000"
|
AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS: "3000"
|
||||||
|
AI_GATEWAY_ROUTE_PROBE_CONCURRENCY: "8"
|
||||||
AI_GATEWAY_WORKER_ENDPOINT_ALLOW_PRIVATE: "true"
|
AI_GATEWAY_WORKER_ENDPOINT_ALLOW_PRIVATE: "true"
|
||||||
AI_GATEWAY_CAPACITY_ORCHESTRATOR_ADAPTER: kubernetes
|
AI_GATEWAY_CAPACITY_ORCHESTRATOR_ADAPTER: kubernetes
|
||||||
AI_GATEWAY_CAPACITY_POOLS: >-
|
AI_GATEWAY_CAPACITY_POOLS: >-
|
||||||
|
|||||||
Reference in New Issue
Block a user