feat(egress): pool shared HTTP clients on the egress hot path - #1718
Open
omrishiv wants to merge 2 commits into
Open
feat(egress): pool shared HTTP clients on the egress hot path#1718omrishiv wants to merge 2 commits into
omrishiv wants to merge 2 commits into
Conversation
Replace per-call httpx client construction on the egress data path with
process-lifetime, connection-pooled clients so repeated egress calls reuse
warm TCP+TLS connections instead of handshaking every request.
url_guard gains pooled accessors:
- shared_guarded_async_client(profile, verify): one SSRF-guarded client per
(profile, verify). The guard validates+pins per request before pool
checkout and pools by the pinned IP, so rebind-safety is preserved.
- shared_plain_async_client(): pooled plain client for in-cluster trusted
hops only (the registry egress-token vend).
- No shared mutable identity state: no default auth headers, and cookie
persistence is disabled (a Set-Cookie cannot bleed across requests/users).
- post_with_reconnect(): one transparent retry when a pooled keep-alive was
closed while idle (httpx does not auto-retry non-idempotent POSTs).
- aclose_shared_clients() / reset_shared_clients_for_tests().
Wire the pooled clients into the egress hot path (guard posture unchanged):
OBO + 3LO exchanges -> pooled guarded credentialed client; egress-token vend
-> pooled plain client; MCP-proxy egress stream -> pooled guarded client
(EGRESS_UPSTREAM_PROFILE), response-scoped and never closed per request
(non-egress keeps a per-call plain client); registry health checks -> shared
process-lived clients instead of per-cycle. Both app lifespans close the
pooled clients on shutdown (registry after the health loop is drained).
New settings across Docker/Terraform/Helm + config API:
EGRESS_HTTP_POOL_MAX_CONNECTIONS/MAX_KEEPALIVE/KEEPALIVE_EXPIRY_SECONDS/CONNECT_RETRIES.
New metric mcpgw_registry_egress_conn_reset_total{site}.
Tests: new tests/unit/utils/test_shared_http_clients.py; existing OBO/3LO/
vend/proxy tests migrated to the shared-client seam.
…rella chart
Review-driven follow-on to the egress connection-pooling change:
- Login OAuth callback (exchange_code_for_token + get_user_info) now reuses the
pooled PLAIN client instead of a per-call httpx.AsyncClient. Deliberately NOT
the HTTPS-only credentialed-OAuth guard: the callback token/userinfo endpoint
is the operator-configured login IdP -- Keycloak/PingFederate use the in-cluster
${KEYCLOAK_URL}/base URL (default http://), which the guard would reject and
break login. EGRESS_OAUTH_TRUSTED_IDP_HOSTS does not help (it relaxes the
private-IP block, not the HTTPS requirement). Target is static config, never
request-derived, so the plain client is correct; 5s timeout + per-request creds
preserved. Plain-client contract doc updated; test asserts the shared client is
not closed per request.
- post_with_reconnect narrowed to httpx.RemoteProtocolError only (the keep-alive
reuse reset); ConnectError is left to the transport retries= (removes the
redundant double-retry on connect). Tests updated + regression test added.
- Umbrella chart surfaces the four EGRESS_HTTP_POOL_* knobs under
registry.egressAuth / auth-server.egressAuth as operational tuning (rest of
egressAuth stays subchart-defaulted).
- docs/egress-http-client-pooling.md: document callback pooling + guard rationale.
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Every outbound call on the egress data path currently builds a fresh
httpxclient and tears it down per request, so each OBO exchange, 3LO
exchange/refresh, egress-token vend, MCP-proxy stream, and periodic health
check pays a full TCP + TLS handshake with no keep-alive reuse. This PR introduces
process-lifetime, connection-pooled clients so repeated egress calls reuse warm
connections, cutting handshake churn and latency on the single-worker asyncio
loop — the throughput ceiling of the authenticated data path.
The change is plumbing that preserves the existing security posture: the
SSRF guard, per-request validate+pin, and credential-handling are unchanged;
one path (the MCP-proxy egress stream) is left exactly as
mainalready had it(guarded via
EGRESS_UPSTREAM_PROFILE).Motivation
and the per-call client pattern re-handshakes on every request.
(default 300s) because the clients are rebuilt per cycle — the single biggest
handshake source under fan-out.
What changed
Core — pooled accessors (
registry/utils/url_guard.py)shared_guarded_async_client(profile, verify)— one SSRF-guarded client per(profile, verify).GuardedAsyncTransportvalidates + pins every requestbefore pool checkout; the pool is keyed by the pinned IP, so a rebound
hostname re-resolves to a new origin and never reuses a stale connection
(rebind-safe).
shared_plain_async_client()— pooled plain client for in-cluster,already-trusted hops only (the registry egress-token vend).
_NoStoreCookieJar— installed on every pooled client so aSet-Cookieisnever stored and can never be replayed onto a later or concurrent request
(concurrency-safe by construction; credentials never ride cookies here — this
is defense-in-depth).
post_with_reconnect()— one transparent retry when a pooled keep-alive wasclosed while idle (httpx does not auto-retry non-idempotent POSTs). Used only
on idempotent hops (OBO exchange, vend); see the 3LO note below.
aclose_shared_clients()/reset_shared_clients_for_tests().Call sites (guard posture unchanged)
auth_server/egress_obo.py) and 3LO exchange/refresh(
registry/egress_auth/oauth_engine.py) → pooled guarded credentialed client.auth_server/server.py) → pooled plain client.auth_server/server.py) → pooled guarded client(
EGRESS_UPSTREAM_PROFILE), response-scoped; the shared client is neverclosed per request. The non-egress stream keeps a per-call plain client
(closed in
finally).registry/health/service.py) → process-livedshared clients instead of per-cycle; initialize+probe to one server now reuse
a connection.
closes after the health loop is drained, so no in-flight health request hits a
closed client).
Config (three-surface parity)
Four new settings, wired through
registry/core/config.py, the config API(
config_routes.py), Docker (.env.example+ all 3 compose files), Terraform(
terraform/aws-ecs/…), and Helm (charts/{registry,auth-server}values +config maps +
reserved-env-names.txt), documented indocs/unified-parameter-reference.md:EGRESS_HTTP_POOL_MAX_CONNECTIONS100max_connectionsEGRESS_HTTP_POOL_MAX_KEEPALIVE20max_keepalive_connections(clamped ≤ max_connections)EGRESS_HTTP_POOL_KEEPALIVE_EXPIRY_SECONDS30EGRESS_HTTP_POOL_CONNECT_RETRIES1New metric
mcpgw_registry_egress_conn_reset_total{site}counts keep-alivereconnect retries (rising =
KEEPALIVE_EXPIRYset above an upstream idletimeout).
Security
(profile, verify)does not weakenvalidate+pin: pinning runs per request before pool checkout, and the pool key
is the pinned IP.
verifyis part of the key, so averify=Falseclient cannever be reused where verification is expected.
header/body bleed); no shared default identity headers; the no-store cookie
jar makes cookie handling stateless.
pool key is the pinned IP, two hostnames that both validate to the same public
IP can coalesce onto one TLS connection. This is safe — each request is
independently pinned, the
Hostheader is correct, and reaching host C overhost B's connection requires C to already resolve to that IP.
http2isdeliberately not enabled (it would coalesce far more aggressively).
oauth_engine._post_tokenhandlessingle-use/rotating grants (authorization_code / refresh_token), so a blind
re-POST after a connection reset could double-spend the grant. It is not
wrapped in
post_with_reconnect; a reset surfaces as a transient error therefresh worker retries safely.
proxy_pass_url) is unchanged bypooling — connection selection is per-request by the pinned target, so a
request for a changed address opens a new connection; it never reuses the old
upstream's keep-alive. (That threat is a registration-integrity concern,
orthogonal to this PR; this PR keeps the stream behind the SSRF guard.)
Testing
tests/unit/utils/test_shared_http_clients.py: sharing identity,verifykeying, guarded-transport/profile preservation, no-store cookie jar,
rebuild-after-close,
aclosecloses all,post_with_reconnectretry-once/no-retry/reraise, and a pin-to-IP coalescing test (two
hostnames → same IP origin, Host + SNI preserved).
seam (they now assert the credential-bearing POST flows through the guarded
client and fails closed), plus an assertion that the shared egress client is
not closed per request.
-W error::RuntimeWarning; all pre-commit gates green(ruff lint + format, mypy, bandit, detect-secrets, fast tests).
Backwards compatibility / rollout
vend budget, health per-call, proxy timeout), guard posture unchanged, no new
required config (all four settings default sensibly).
setting
EGRESS_HTTP_POOL_KEEPALIVE_EXPIRY_SECONDS=0to disable keep-alivereuse while keeping the pooled client objects. (Note:
max_connections=1isnot a rollback — it still keep-alives and serializes egress.)
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.