diff --git a/apps/api/internal/config/config.go b/apps/api/internal/config/config.go index 5079554..9ea9ef9 100644 --- a/apps/api/internal/config/config.go +++ b/apps/api/internal/config/config.go @@ -101,6 +101,7 @@ type Config struct { RouteProbeHotIntervalSeconds int RouteProbeColdIntervalSeconds int RouteProbeTimeoutMS int + RouteProbeConcurrency int WorkerAutoscalingEnabled bool CapacityOrchestratorAdapter string CapacityPoolsJSON string @@ -223,6 +224,7 @@ func Load() Config { RouteProbeHotIntervalSeconds: envIntValidated("AI_GATEWAY_ROUTE_PROBE_HOT_INTERVAL_SECONDS", 15), RouteProbeColdIntervalSeconds: envIntValidated("AI_GATEWAY_ROUTE_PROBE_COLD_INTERVAL_SECONDS", 60), 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", CapacityOrchestratorAdapter: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_CAPACITY_ORCHESTRATOR_ADAPTER", "kubernetes"))), 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) { 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 { return errors.New("AI_GATEWAY_WORKER_EXECUTION_SECRET must be at least 32 bytes in enforced routing mode") } diff --git a/apps/api/internal/runner/route_prober.go b/apps/api/internal/runner/route_prober.go index 416c3ac..18e0f38 100644 --- a/apps/api/internal/runner/route_prober.go +++ b/apps/api/internal/runner/route_prober.go @@ -27,6 +27,11 @@ func (s *Service) runRouteHealthProber(ctx context.Context) { ticker := time.NewTicker(time.Second) defer ticker.Stop() states := make(map[string]routeProbeState) + probeConcurrency := s.cfg.RouteProbeConcurrency + if probeConcurrency <= 0 { + probeConcurrency = 8 + } + probeSlots := make(chan struct{}, probeConcurrency) for { select { case <-ctx.Done(): @@ -58,19 +63,35 @@ func (s *Service) runRouteHealthProber(ctx context.Context) { if target.Hot { interval = time.Duration(s.cfg.RouteProbeHotIntervalSeconds) * time.Second } - states[profile.Key] = routeProbeState{next: now.Add(interval), lastRequestObserved: requestedAt} - go s.probeRouteTarget(ctx, target, profile) + select { + case probeSlots <- struct{}{}: + states[profile.Key] = routeProbeState{next: now.Add(interval), lastRequestObserved: requestedAt} + 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) { - 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 { 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, "") if err != nil { s.logRouteProbeFailure("prepare_client", err) diff --git a/apps/api/internal/store/execution_pools.go b/apps/api/internal/store/execution_pools.go index 65b4f72..f15a03c 100644 --- a/apps/api/internal/store/execution_pools.go +++ b/apps/api/internal/store/execution_pools.go @@ -222,6 +222,49 @@ WHERE expires_at > now()`) 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) { return s.listRouteHealth(ctx, "WHERE route_profile_key = $1", strings.TrimSpace(routeProfileKey)) } diff --git a/apps/api/internal/store/execution_pools_integration_test.go b/apps/api/internal/store/execution_pools_integration_test.go index 303364d..c94fc79 100644 --- a/apps/api/internal/store/execution_pools_integration_test.go +++ b/apps/api/internal/store/execution_pools_integration_test.go @@ -4,6 +4,9 @@ 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) { @@ -18,3 +21,42 @@ func TestExecutionPoolQueriesUseTimestampCutoff(t *testing.T) { 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) + } +} diff --git a/apps/api/migrations/0106_route_probe_leases.sql b/apps/api/migrations/0106_route_probe_leases.sql new file mode 100644 index 0000000..79e4a24 --- /dev/null +++ b/apps/api/migrations/0106_route_probe_leases.sql @@ -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); diff --git a/deploy/kubernetes/production/application-config.yaml b/deploy/kubernetes/production/application-config.yaml index bd96a76..6d4e7bf 100644 --- a/deploy/kubernetes/production/application-config.yaml +++ b/deploy/kubernetes/production/application-config.yaml @@ -52,6 +52,7 @@ data: AI_GATEWAY_ROUTE_PROBE_HOT_INTERVAL_SECONDS: "15" AI_GATEWAY_ROUTE_PROBE_COLD_INTERVAL_SECONDS: "60" AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS: "3000" + AI_GATEWAY_ROUTE_PROBE_CONCURRENCY: "8" AI_GATEWAY_WORKER_ENDPOINT_ALLOW_PRIVATE: "true" AI_GATEWAY_CAPACITY_ORCHESTRATOR_ADAPTER: kubernetes AI_GATEWAY_CAPACITY_POOLS: >-