perf: pre-warm upstream realtime connection pool to cut session-establishment latency

[ishaan-berri]

Relevant issues

Latency optimization for the Rust ai-gateway realtime path. Re-benchmarked live on Render at 5000 calls / 500 concurrency, 10 gateway instances.

Pre-Submission checklist

• I have added meaningful tests
• My PR's scope is as isolated as possible; it only solves 1 specific problem (realtime session-establishment latency)
• Greptile review pending

Type

🚄 Infrastructure
🆕 New Feature

Changes

A latency optimization on the Rust ai-gateway realtime path.

The problem

The gateway's only added cost over direct-to-OpenAI is session establishment: on each client connect it dials a fresh upstream WS to OpenAI and blocks until session.created. Dial, first-audio and streaming add ~0.

What this does

A pre-warmed upstream connection pool keyed by (model, api_key, api_base). A background task keeps OpenAI realtime sockets already connected and already past session.created (buffered). On a client connect we take a warm socket and relay its buffered session.created instantly via a local handoff (a Vec::pop + liveness check — sub-millisecond, no network), then splice exactly as today. On a pool miss or dead warm socket we fall back to fresh-dial — the pool is never on the critical path for correctness, only latency.

Config: REALTIME_POOL_SIZE (default 4; 0 disables → fresh-dial) and REALTIME_POOL_MAX_IDLE_SECS (default 30).

Concurrent refill (makes it land at scale)

The replenisher originally dialed missing warm sockets sequentially, so a full refill cost needed × handshake (~needed × 350 ms). Under high connect rates the pool drained faster than it refilled — at 5000/500 only ~15% of connects hit the pool, so the median connect still paid the full handshake. This PR fires the refill dials concurrently (join_all), refilling in ~one handshake window. Result: pool hit rate went 15% → ~76% and median session 357 ms → 0 ms.

Replenisher safety (review follow-ups)

• Backoff on failing keys. A key whose warm-up dials all fail (invalid credentials / unreachable upstream) is put into exponential backoff (500 ms base, doubling, 30 s cap) instead of firing needed concurrent TLS dials every 250 ms tick — bounds connection attempts so a single misconfigured deployment can't exhaust upstream rate limits and degrade valid cold-path traffic. Resets on the first successful dial; its own lock keeps the request hot path (take) contention-free.
• MissedTickBehavior::Skip on the replenish interval, so a slow tick doesn't trigger catch-up bursts.
• UpstreamKey Debug redacts api_key ([REDACTED]) so the credential can't leak via logs.

Benchmark results (5000 calls / 500 concurrency, 10 instances)

All times in milliseconds. Each leg run twice; values pooled across both runs. Harness: ishaan-berri/litellm-realtime-bench (ws-bench), run from an 8-vCPU Render runner against OpenAI-direct and the deployed gateway. Phases per WebSocket connection:

• Dial — TCP + TLS + WS upgrade handshake
• Session — upgrade done → OpenAI session.created (the phase the pool removes)
• 1st-audio — response.create sent → first audio delta (OpenAI inference)
• Total — full wall-clock per connection

phase / metric	Direct OpenAI	Gateway pool OFF	Gateway pool ON
success	99.8%	99.8%	99.8%
session p50	7	367	0
session p95	18	692	767
session p99	29	847	1055
dial p50	276	85	158
1st-audio p50	440	459	664
total p50	816	952	1010
total p95	2152	1907	1970
total p99	2692	2386	2610

Direct vs Gateway (pool ON), with overhead and where the gateway is faster:

metric	Direct OpenAI	Gateway pool ON	Overhead (ms)	vs OpenAI
dial p50 (ms)	276	158	−118	faster
session p50 (ms)	7	0	−7	faster
1st-audio p50 (ms)	440	664	+224	slower¹
total p50 (ms)	816	1010	+194	slower¹
total p95 (ms)	2152	1970	−182	faster
total p99 (ms)	2692	2610	−82	faster

The gateway is faster than direct OpenAI on 4 of 6 metrics (negative overhead): dial and the total tail (p95/p99), because the runner + gateway are co-located in Render Oregon while OpenAI's edge is further, and the warm handoff makes session sub-millisecond. ¹ The two "slower" rows track a 1st-audio difference (440 → 664 ms) between the legs. 1st-audio is OpenAIs own inference time, which the gateway only relays — but the legs ran at different times, so this gap cannot yet be attributed to OpenAI-side variance vs a possible warm-socket penalty without an interleaved A/B (alternating direct/gateway connections within one run so OpenAI variance cancels). Defensible claim: session-establishment overhead is eliminated (−7 ms p50, sub-1ms warm handoff); end-to-end total is inconclusive pending that controlled run.

Gateway session overhead vs direct: pool OFF +360 ms p50 → pool ON −7 ms p50. The warm handoff is sub-millisecond, so the median connect's session phase collapses to ~0. ~76% of connects hit the pool; ~70% had session < 1 ms. The remaining ~24% miss and fresh-dial (~350 ms), which is the p95/p99 session tail — raise REALTIME_POOL_SIZE to shrink it.

Pool sizing: size to peak concurrent connects per instance, REALTIME_POOL_SIZE ≈ peak_concurrency / instances. This run used 64/instance for ~500 concurrency over 10 instances. The concurrent warm-up (≈64 dials/instance at boot) was accepted by OpenAI with no rate-limit errors.

Note: gateway dial is lower than direct and gateway total p95/p99 are near direct, because the runner + gateway are co-located in Render Oregon while OpenAI's edge is further. The gateway's only real overhead is the session phase, which the pool removes at the median.

Tests

Pool mechanics unit-tested against an in-process fake OpenAI realtime WS server: warm handoff relays the buffered session.created; pool miss falls back to fresh-dial; dead/stale warm socket is discarded; background replenisher tops up a registered key; a broken key enters backoff and is not re-dialed during the window; a healthy key never backs off. cargo test -p litellm-providers → all pool tests pass (concurrent refill + backoff included).

Docs in this PR

• src/routes/realtime/README.md — pooling design next to the code: handoff diagram, sizing rule, config knobs, caveats.
• benchmarks/realtime/README.md — the results table above + how to reproduce.

The load-gen harness lives in its own repo so it doesn't bloat litellm: https://github.com/ishaan-berri/litellm-realtime-bench (no keys hardcoded; bearer token via -key/$OPENAI_API_KEY).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ ishaan-jaff
❌ ishaan-berri
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[greptile-apps]

Greptile Summary

This PR adds a pre-warmed upstream WebSocket connection pool to the Rust ai-gateway's realtime route, eliminating the per-connect session.created handshake latency (~360 ms → ~0 ms at the p50). A background replenisher keeps a configurable number of already-connected OpenAI sockets ready; on a client connect the pool hands off a warm socket (sub-millisecond local Vec::pop) and falls back to fresh-dial on miss — the pool is never on the correctness path.

• Pool mechanics (realtime_pool.rs): concurrent refill via join_all, exponential backoff for broken keys, MissedTickBehavior::Skip, per-key backoff tracking, and non-blocking liveness check at handoff. Configurable via REALTIME_POOL_SIZE (default 4; 0 disables) and REALTIME_POOL_MAX_IDLE_SECS (default 30).
• Refactored realtime.rs: dial_upstream, splice, and read_event extracted so the pool reuses the same types and transforms as the cold path; realtime_warm relays the buffered session.created before splicing.
• Integration: pool registered at startup per deployment, threaded through AppState → mod.rs → service.rs; well-covered by unit tests against an in-process fake OpenAI WS server.

Confidence Score: 4/5

The fallback to fresh-dial on any pool miss means the pool cannot break realtime sessions for correctness, but one gap in warm_one — accepting any first JSON frame without checking the event type — can cause clients to receive a non-session.created event (e.g. an OpenAI application-level error) in place of the real one.

The overall design is sound and the cold-path fallback ensures the pool is never a correctness dependency. However, warm_one calls read_event and stores the result unconditionally in session_created without verifying event_type == "session.created". If OpenAI sends an application-level error event after accepting the WebSocket, warm_one succeeds, any_ok is set, backoff is cleared, and every subsequent client receives that error event where session.created is expected.

litellm-rust/crates/providers/src/realtime_pool.rs — specifically the warm_one function's handling of the first upstream event.

Important Files Changed

Filename	Overview
litellm-rust/crates/providers/src/realtime_pool.rs	New file implementing the pre-warmed connection pool. Exponential backoff, liveness checks, and concurrent refill are all present, but warm_one accepts any JSON first-frame as the session.created event without validating the event type — meaning application-level error events from OpenAI would be silently cached and handed to clients.
litellm-rust/crates/providers/src/realtime.rs	Refactored to expose dial_upstream, splice, and read_event so the pool can reuse them; the public realtime function is preserved with identical semantics, and the new realtime_warm correctly relays the buffered prelude before splicing.
litellm-rust/crates/ai-gateway/src/main.rs	Pool built and registered at startup; MissedTickBehavior::Skip is set, the disabled path is handled, and the weak-ref shutdown guard is correct.
litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs	Warm path added cleanly before the cold-dial fallback; pool is never on the correctness path — a miss or dead socket falls through to fresh-dial exactly as before.
litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs	Propagates the pool through the bridge closure; axum-layer changes are minimal and correct.
litellm-rust/crates/ai-gateway/src/state.rs	Adds realtime_pool: Arc<RealtimePool> to shared state; trivial and correct.
litellm-rust/crates/core/src/router/mod.rs	Adds a read-only deployments() accessor so main.rs can enumerate upstreams without coupling to internals; no behavior change.
litellm-rust/crates/providers/src/lib.rs	Adds pub mod realtime_pool; one-line change, correct.

Comments Outside Diff (1)

1. litellm-rust/crates/providers/src/realtime_pool.rs, line 1051-1062 (link)

   First event not validated as session.created

   warm_one reads the first text frame from OpenAI and stores it unconditionally in session_created, but read_event accepts any parseable JSON — it doesn't check event_type. If OpenAI sends an application-level error as the first frame (e.g. a rate-limit or model-not-found event accepted at the WS level) warm_one still returns Ok, any_ok becomes true, backoff is cleared, and every subsequent client on this key receives that error event relayed in place of session.created. A one-line guard after read_event would fix this: reject any event whose event_type != "session.created" by returning Err(CoreError::InvalidResponse(...)), so the socket is discarded and the failure is correctly counted toward backoff.

_{Reviews (3): Last reviewed commit: "docs: document the replenisher backoff i..." | Re-trigger Greptile}

[ishaan-berri]

@greptileai please re-review — addressed the replenisher backoff finding.

The background replenisher now puts a key into exponential backoff (500 ms base, doubling, 30 s cap) after a pass whose warm-up dials all fail (invalid credentials / unreachable upstream), instead of firing needed concurrent TLS dials against the upstream on every 250 ms tick. Backoff resets on the first successful dial; the backoff map has its own lock so the request hot path (take) never contends on it. Added two unit tests: a broken key enters backoff and is not re-dialed during the window, and a healthy key never backs off. fmt + clippy (-D warnings) + cargo test all green.

[ishaan-berri]

@greptileai review

[tin-berri]

1. Fix is_dead to not treat Ping/Pong as death (+ a regression test) — otherwise the optimization silently degrades at steady state.
2. Drop the clippy::too_many_arguments allow by bundling the upstream tx/rx into a struct — the repo doesn't suppress lints, it fixes the type.
3. Batch the refill push under one lock (folds together with replacing the any_ok mutable flag), and cap REALTIME_POOL_SIZE / add the idle-cost + ping caveats to the README.