feat(gateway-pools): load balancing - #7703
Conversation
|
💬 Discussion in Slack: #pr-review-infisical-7703-feat-gateway-pools-send-work-to-the-least-busy-gateway Posted by Review Police — reviews, comments, new commits, and CI failures will stream into this channel. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e954d8e86c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
|
| Filename | Overview |
|---|---|
| backend/src/ee/services/gateway-pool/gateway-pool-service.ts | Adds load-aware member selection, suspect exclusion, reservations, and bounded transport-only failover. |
| backend/src/lib/gateway-v2/gateway-load-tracker.ts | Adds Redis-backed reported load, per-pod channel counts, reservations, stale-data handling, and suspect markers. |
| backend/src/lib/gateway-v2/gateway-v2.ts | Tracks established channels and distinguishes pre-tunnel transport failures from target-side errors. |
| backend/src/ee/routes/v2/gateway-router.ts | Adds the authenticated load-report endpoint, but its schema metadata and receipt response violate the backend API guide. |
| backend/src/services/app-connection/app-connection-service.ts | Retries read-only connection validation through alternate pool members and reuses the successful member for non-retryable credential transitions. |
| backend/src/services/identity-kubernetes-auth/identity-kubernetes-auth-service.ts | Adds pool failover around gateway-mediated Kubernetes token review while preserving direct gateway fallback behavior. |
Reviews (1): Last reviewed commit: "fix(gateway-pools): address review, and ..." | Re-trigger Greptile
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 1 · PR risk: 0/10 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Workflows to automatically generate PRs for you. |
|
FINDING 1 Redis-backed load tracking governs selection for 500+ connection fan-out paths The repository migrates diverse connection providers to a centralized Evidence:
Agent Prompt: FINDING 2 AsyncLocalStorage attempt context decoupled from 2200+ deep cross-service execution chains The new retry logic in Evidence:
Agent Prompt: FINDING 3 Unified ALPN-based protocol multiplexing standardizes disparate tunneling execution paths Disparate connection providers for SSH, SQL, and HTTP are being refactored to use a unified Evidence:
Agent Prompt: |
bernie-g
changed the title
Adds gatewayId to TGatewayV2ConnectionDetails and threads it to the proxy layer, which until now received the gateway's mTLS material but no way to say which gateway it was talking to. Nothing consumes it yet. Call sites spread the connection details rather than re-listing their fields, so the type stays the single source of truth and a field added later cannot be silently dropped at 26 places. No behaviour change.
Replace random gateway-pool selection with least-connections, and route around a member whose tunnel cannot be established. Selection - Track open gateway channels per pod in Redis, summed at selection time alongside short-lived reservations covering the window between choosing a member and its channel opening. - Count channels inside setupRelayServer rather than around the caller's operation, since one operation can open many channels. - Pick the least-loaded member, breaking between two comparable members at random so concurrent pods reading the same counters do not stampede the same gateway. A tie band keeps the score meaningful for two-member pools instead of degrading to a coin flip. - Degrade to random whenever Redis is unavailable. Failover - runWithPoolFailover retries on another member, but only when no tunnel was ever established, so an operation that may have partially applied on the target is never replayed. Providers rewrap gateway errors in their own BadRequestError, so the signal travels through async-local storage scoped to a single attempt rather than on the exception. - Mark a failed member suspect for 60s, closing the window where a dead gateway still looks healthy to the heartbeat check. - Wire into Kubernetes auth login and app-connection validation. The platform-managed credential transition deliberately does not fail over and reuses the member validation succeeded on. Also folds the HSM connector's private selection copy onto the shared selector via a capability filter, and fixes two keystore mocks that had the wrong arity for setItemWithExpiry and a no-op decrementByOrDelete.
Review fixes - Only treat a failure as retryable when no tunnel was ever established. The relay error list is shared by every channel a proxy serves and is never cleared, so one transient setup failure would otherwise keep marking later target-side errors as safe to replay. - Clear that classification once any tunnel comes up, so an attempt whose first tunnel failed and was swallowed cannot be replayed after a later one reached the target. - Release a reservation when its channel opens rather than on a fixed timer, which double counted the member and expired early under a slow relay. - Read reservations and suspect marks from the Redis primary; they are written there and read back within the same sub-second window. - Tear down a channel whose caller gave up during the handshake, which otherwise stayed counted for the life of the pod. - Keep the BadRequest error name so the serialised error field does not change for callers; retry decisions use instanceof. - Reuse the member validation succeeded on for the platform-managed credential transition, which mutates the target and must not fail over. - Restore the 404 for an unknown pool, surface the real PKCS#11 error once HSM members run out, and never let load bookkeeping fail a selection. Load signal - The gateway now reports its own active channel count, which covers work the platform never proxied. A PAM CLI session dials the relay directly, so it was previously invisible and a saturated member scored zero. - A member too old to report is scored on platform-opened channels only, which is a lower number than a reporting peer's true occupancy, so a mixed pool gets no load awareness rather than a biased guess that would pull traffic onto un-upgraded gateways for a whole rollout. - Tie band is exact equality: a reservation is worth one, so any wider band cancelled it out and re-shortlisted the member just claimed.
Both retry decisions were inline boolean logic inside catch blocks, and neither could be reached from a test: they only trigger on a mixed sequence within one operation, where an early channel fails at setup and a later one reaches the target. Extracting them states the shared rule once (retry only when nothing can have reached the target) and makes both testable. Deleting either guard now fails tests, which was verified by mutation rather than assumed.
Choosing and claiming a pool member were a read then a separate write, so every concurrent selection saw the same minimum before any of them claimed it and they all routed to it. Both now happen in one Redis call; ties fall to the caller's order, which is shuffled, so the random tie-break stands. - Identity load reports run the same permission check heartbeat does, so a revoked identity cannot keep steering pool routing. - Load reports get their own limiter keyed by the reporting gateway. The shared IP-keyed write limiter would be exhausted by ~100 gateways behind one NAT address, and their entries would then go stale. - The endpoint returns the recorded state, documents its fields, and quotes identifiers in its errors. - The fallback keeps the suspect filter. It was scoped inside the try, so any failure in the load path re-admitted the member that just failed. - Channels opened since a gateway's last report count on top of it, instead of being invisible until the next one lands. - Drop chooseLeastLoadedGateway: the atomic claim supersedes it, and its tests were pinning behaviour production no longer runs. The failover suite's mocked tracker had no claimLeastLoaded, so it threw inside the try and only ever measured the random fallback. One test was passing for the wrong reason.
The page stated that selection was random and that there was no round-robin or weighted selection, which is now the opposite of what happens. Also documents that only failures occurring before the platform reaches the resource are retried, and that a pool containing a gateway too old to report its load falls back to random selection.
Hash tag the per-gateway keys so the claim script does not hit CROSSSLOT on Redis cluster, set the reservation TTL only on the first claim so a dead pod's reservations actually expire, guard the Lua script against a short occupancy list, and release published load on shutdown.
setupRelayServer can now open the tunnel during setup and hand it to the first client instead of dialing lazily, so an unreachable gateway fails where another member can still be tried. The tunnel is reused rather than probed and dropped, so the gateway sees one channel that becomes the session.
The rebase merged getPlatformConnectionDetailsByPoolId out of the service while leaving its caller in place. Restored, now routing through the load-aware selection instead of a random pick.
70850ba to
4562505
Compare
This call site enumerated the connection details by hand, so it stopped satisfying the type once gatewayId was added to it. Spreading matches the other withGatewayV2Proxy callers and survives the next field.
The harness carries an ephemeral port in relayHost, which the implementation forwards as the TLS servername, and Node 22.23 rejects a servername containing a colon instead of matching it against a SAN. Verification now runs against the connected host, which the certificate already covers.
A gateway too old to report can only be scored on the connections the platform itself opened, which misses whatever else it is carrying. Rather than trust that number on its own, any pool with a non-reporting member now picks at random. isPoolComparable becomes everyMemberReportsLoad: the old name read as a claim about the gateways rather than about where the numbers came from.
Scoring an unreported gateway from the platform's own view was the only reason the per-pod Redis hash existed, so the publish path, its debounce, refresh timer and shutdown cleanup all go with it, along with hashGetAll/hashDelete. Also drops TGatewayScore.score, which nothing read: the claim script fetches the reservation keys itself, so folding them in here meant reading them twice per selection. Its tests asserted on that field, so they now observe reservations through the claim script, which is the path that actually uses them.
Most explained mechanism the code already shows, or restated a point made a few lines above. Kept the non-obvious reasons in a clause each.
Named after the decision it gates rather than the criterion it checks, so the branch reads as least-busy versus random at the call site.
shutdown no longer writes to Redis, so it does not drop published counts.
3 participants
Context
Gateway pools picked a gateway at random, so one could end up carrying far more work than the others while its peers sat idle. Pools now send each new job to whichever gateway is least busy, and retry on another gateway when one can't be reached.
Needs the matching CLI change so gateways report how busy they are: Infisical/cli#370
Steps to verify the change
Type
Checklist
type(scope): short description(scope is optional, e.g.,fix: prevent crash on syncorfix(api): handle null response).