feat(runtime): multiplex TCP response streams - #11918
Conversation
|
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
Code reviewFull read of the diff (all 15 files) plus verification of call sites and lint behavior. OverviewOpt-in, versioned TCP response-mux: 4 persistent connections per worker-process/frontend pair, logical streams multiplexed by UUID, HTTP/2-like framing with per-stream + per-connection byte credits, cumulative connection ACKs, a priority lane for control frames, fair round-robin scheduling (quantum 8), and up-to-5 ms cross-stream write batching. Dedicated transport retained for rollback. Overall well engineered: strict frame validation, credit-conserving flow control, bounded-cardinality metrics, and a genuinely strong test suite. Correctness — one significant bug1.
Correctness — smaller issues2. Normal client cancellation now surfaces as a handler error. 3. One malformed prologue tears down the whole shared connection. In 4. Mixed-fleet rollout hazard. Enablement is per-process via env with no capability negotiation — the frontend unilaterally hands out mux connection info, and a worker with mux disabled hard-fails those requests. Rolling deploys with mixed config break requests. Deserves a prominent deploy note; the new env vars also aren't documented anywhere yet. 5. Handshake buffer discard is safe only by convention. Both sides call Code quality / style
Performance
Test coverageStrong: codec round-trips and rejections, config bounds, credit-cap invariants, cumulative-ACK monotonicity, waiter wakeups on close, End-bypasses-credit-exhaustion, cross-stream backpressure isolation, failure scoping with reconnect, 1,000 streams over exactly 4 connections. Gaps:
SecurityNo concerns beyond the existing transport's trust model. Handshake validates version and frontend UUID (misdirection protection, not authentication — same as the dedicated transport). Frame validation is strict in both codecs; the single VerdictArchitecture, flow-control accounting, and test discipline are excellent. Requesting changes for #1 (frame loss — fix + regression test); recommend addressing #2 and the clippy discrepancy before merge. #3–#5 and doc notes can be follow-ups. 🤖 Generated with Claude Code |
WalkthroughChangesTCP response multiplexing
Estimated code review effort: 5 (Critical) | ~120 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
lib/runtime/src/pipeline/network/tcp/mux/client.rs (1)
89-92: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winGate per-frame metrics behind config
per_frame_metrics_enabled()is hardcoded totrue, so the per-frame histogram/counter paths run unconditionally on the hot writer/reader/send paths. If this is meant to mirrorresponse_packet_metrics_enabled(), wire it to config instead of always paying the cost.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/runtime/src/pipeline/network/tcp/mux/client.rs` around lines 89 - 92, Update per_frame_metrics_enabled() to read the appropriate configuration value, mirroring response_packet_metrics_enabled(), instead of always returning true. Ensure the existing per-frame metric paths remain enabled only when that configuration flag is enabled.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@benchmarks/frontend/scripts/run_perf.sh`:
- Around line 764-771: Validate AIPERF_WORKERS_MAX and AIPERF_RECORD_PROCESSORS
as canonical non-negative integer strings before adding them to
_AIPERF_WORKER_ARGS or serializing them into config.json. Reject values such as
foo, 08, or other non-canonical inputs consistently for both CLI and environment
sources, including the corresponding handling around the alternate lines noted
in the comment.
---
Nitpick comments:
In `@lib/runtime/src/pipeline/network/tcp/mux/client.rs`:
- Around line 89-92: Update per_frame_metrics_enabled() to read the appropriate
configuration value, mirroring response_packet_metrics_enabled(), instead of
always returning true. Ensure the existing per-frame metric paths remain enabled
only when that configuration flag is enabled.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ee7ed488-6a16-40a3-92d2-339d5e1f8d94
📒 Files selected for processing (15)
benchmarks/frontend/scripts/run_perf.shlib/runtime/src/component/endpoint.rslib/runtime/src/config/environment_names.rslib/runtime/src/distributed.rslib/runtime/src/metrics.rslib/runtime/src/metrics/response_mux.rslib/runtime/src/pipeline/network.rslib/runtime/src/pipeline/network/egress/addressed_router.rslib/runtime/src/pipeline/network/egress/tcp_client.rslib/runtime/src/pipeline/network/ingress/push_handler.rslib/runtime/src/pipeline/network/tcp.rslib/runtime/src/pipeline/network/tcp/client.rslib/runtime/src/pipeline/network/tcp/mux.rslib/runtime/src/pipeline/network/tcp/mux/client.rslib/runtime/src/pipeline/network/tcp/server.rs
Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
c616b8e to
06d2282
Compare
github-actions
Bot
added
the
documentation
Signed-off-by: jthomson04 <jwillthomson19@gmail.com>
|
@coderabbitai full review |
✅ Action performedFull review finished. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@lib/runtime/src/pipeline/network/tcp/server.rs`:
- Around line 782-792: Bound the handshake prefix acquisition in the server
handshake flow around the stream.peek loop: use a single buffered read such as
read_exact or fill_buf and wrap it in the established handshake timeout
mechanism. Preserve the connection-closed error for incomplete input, while
ensuring peers that send only 1–4 bytes and stall cannot keep the loop polling
indefinitely.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e1c87ba6-0f8a-4848-a01b-9cd5aaaceb2b
📒 Files selected for processing (15)
docs/design-docs/request-plane.mdlib/runtime/src/component/endpoint.rslib/runtime/src/config/environment_names.rslib/runtime/src/distributed.rslib/runtime/src/metrics.rslib/runtime/src/metrics/response_mux.rslib/runtime/src/pipeline/network.rslib/runtime/src/pipeline/network/egress/addressed_router.rslib/runtime/src/pipeline/network/egress/tcp_client.rslib/runtime/src/pipeline/network/ingress/push_handler.rslib/runtime/src/pipeline/network/tcp.rslib/runtime/src/pipeline/network/tcp/client.rslib/runtime/src/pipeline/network/tcp/mux.rslib/runtime/src/pipeline/network/tcp/mux/client.rslib/runtime/src/pipeline/network/tcp/server.rs
| let mut prefix = [0_u8; 5]; | ||
| loop { | ||
| let read = stream.peek(&mut prefix).await?; | ||
| if read == 0 { | ||
| anyhow::bail!("connection closed before its handshake"); | ||
| } | ||
| if read == prefix.len() { | ||
| break; | ||
| } | ||
| tokio::task::yield_now().await; | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant area of the server implementation.
file='lib/runtime/src/pipeline/network/tcp/server.rs'
wc -l "$file"
sed -n '720,840p' "$file"
# Look for other handshake-prefix handling or timeout patterns in the repo.
rg -n "peek\\(&mut prefix\\)|HANDSHAKE|handshake timeout|yield_now\\(\\)" lib/runtime/src -SRepository: ai-dynamo/dynamo
Length of output: 6414
🌐 Web query:
tokio AsyncReadExt peek buffered data immediate return partial prefix busy loop yield_now semantics
💡 Result:
In Tokio, the peek operation—typically accessed via TcpStream::peek (or ReadHalf::peek)—is designed to retrieve data from a socket without removing it from the kernel's receive queue [1]. Regarding your specific technical concerns: Partial Prefix and Immediate Return: The peek method will return as soon as it can retrieve any amount of data up to the requested buffer size [1]. If data is already available in the socket's receive buffer, it returns immediately with that data [1]. If no data is currently available, it will asynchronously wait (suspend the task) until data arrives or the connection is closed [2]. It is not guaranteed to fill the entire provided buffer; it returns the number of bytes actually peeked [1]. Busy Loop and yield_now: Using a loop that repeatedly calls peek without awaiting it or reacting to its asynchronous nature can lead to a busy loop (consuming 100% CPU) [3]. To prevent this, you must await the peek call. If you are implementing a custom polling mechanism or need to ensure other tasks get a chance to run, tokio::task::yield_now() can be used to yield execution back to the runtime, though the runtime's scheduler does not guarantee that other tasks will run before the current task is polled again [3]. Common Pitfalls: A known issue exists where peek on certain platforms (e.g., Windows) may behave unexpectedly or block after a read because the underlying I/O state is not correctly re-registered when peek would have returned WouldBlock [2]. If you encounter hanging or inefficient behavior, wrapping your reader in a BufReader and using fill_buf()/consume() is the recommended, more robust alternative to manual peek operations [2][4]. BufReader provides a buffered, predictable interface that avoids the complexities of raw socket peek semantics [2]. It is important to remember that peek does not consume data [1]. If you rely on peek to make decisions (like framing), you must subsequently use read to advance the stream position; otherwise, subsequent read calls will return the same data, leading to infinite loops if not handled correctly [1].
Citations:
- 1: https://docs.rs/tokio/latest/tokio/net/tcp/struct.ReadHalf.html
- 2: peek will block after read tokio-rs/tokio#3789
- 3: https://docs.rs/tokio/latest/tokio/task/fn.yield_now.html
- 4: https://docs.rs/tokio/latest/src/tokio/io/util/async_buf_read_ext.rs.html
Bound the handshake prefix read peek() returns immediately once any bytes are buffered, so a peer that sends 1–4 bytes and then stalls can keep this loop alive indefinitely and make it repeatedly yield/poll. Add a handshake timeout, and prefer a single buffered read (read_exact/fill_buf) instead of repeated peeks.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@lib/runtime/src/pipeline/network/tcp/server.rs` around lines 782 - 792, Bound
the handshake prefix acquisition in the server handshake flow around the
stream.peek loop: use a single buffered read such as read_exact or fill_buf and
wrap it in the established handshake timeout mechanism. Preserve the
connection-closed error for incomplete input, while ensuring peers that send
only 1–4 bytes and stall cannot keep the loop polling indefinitely.
ReviewOverview: This hard-replaces the per-request dedicated worker→frontend response socket with the versioned I read the full diff and verified the load-bearing details against the codebase. Things that check out:
Issues are posted as inline comments. In severity order: (1) busy-wait on a partial handshake prefix in Design & operational risks
Test coverageGood breadth: codec round-trips/validation, config parsing (100 ms cap, malformed rejection), ordering under batching, slow-consumer isolation across streams on one connection, failure scoping + pool repair, 1,000 streams over exactly 4 connections, pipelined raw-wire handshake with malformed-prologue isolation, legacy-protocol rejection. Gaps worth adding:
VerdictImplementation quality is high — the concurrency design (single writer owning the socket, permits carried inside 🤖 Generated with Claude Code |
| if read == prefix.len() { | ||
| break; | ||
| } | ||
| tokio::task::yield_now().await; |
There was a problem hiding this comment.
Busy-wait on partial handshake prefix. peek returns Ready immediately whenever ≥1 byte is buffered, so a peer that sends 1–4 bytes and then stalls turns this loop into a hot spin (peek/yield_now forever), pinning CPU per such connection. Combined with the (pre-existing) lack of a handshake timeout, this is a cheap resource-exhaustion vector on an exposed port.
At minimum, wrap the prefix wait + handshake in a tokio::time::timeout; better, read_exact the 5 bytes and seed them into the FramedRead via read_buffer_mut() so there's no polling loop at all.
| pool | ||
| } | ||
|
|
||
| fn start_maintenance(pool: &Arc<Self>) { |
There was a problem hiding this comment.
Host pools are immortal and retry dead frontends forever. hosts never evicts, and this 1 s maintenance loop calls warm() on every under-filled pool indefinitely. A frontend that goes away permanently (rolling restarts mint a new server UUID and usually a new address) leaves a stale HostKey whose pool attempts a TCP connect roughly every second, forever — per stale frontend. Over weeks of frontend churn this accumulates unbounded connect churn, warn-log spam from warm(), and memory.
Suggest an idle policy: evict (or stop warming) a host pool after N minutes with zero active streams, or exponential backoff that gives up until a new stream references the host again.
| .map_err(|_| anyhow!("response mux urgent writer stopped")) | ||
| } | ||
|
|
||
| fn try_send_urgent(&self, frame: MuxFrame) { |
There was a problem hiding this comment.
Urgent-queue overflow escalates to whole-connection teardown. If the urgent queue (4,096) fills — e.g. a mass-cancellation burst where thousands of MuxResponseStreamSender::drops each enqueue a Reset while the writer is TCP-blocked — this fails the connection and kills every stream on it, and those kills produce further resets.
If the escalation is intentional (it is observable via connection_lost_streams_total), a comment here would help; also, the reader_task call sites are async and could use send_urgent(...).await instead, reserving the lossy path for Drop, which cannot await.
| }); | ||
| } | ||
|
|
||
| fn try_send_control( |
There was a problem hiding this comment.
Same fanout tradeoff as try_send_urgent on the worker side: one try_send overflow on a 4,096-deep channel cancels control_failed and fails the connection plus all of its streams. Fine as a deliberate bounded-admission simplification, but worth a comment stating that this blast radius is intended.
| }; | ||
| crate::metrics::response_mux::SETUP_SECONDS | ||
| .observe(pending.registered_at.elapsed().as_secs_f64()); | ||
| let mailbox_frames = pending.send_buffer_count.max( |
There was a problem hiding this comment.
Nit: data_plane_channel's doc comment still implies send_buffer_count sizes the response mailbox, but this makes it a floor, not a capacity (window / MUX_HEADER_LEN ≈ 12,483 by default). That floor is required for credit compliance — a credit-compliant worker must never see mailbox_full — so a one-line comment here would prevent someone "fixing" it back and reintroducing spurious resets.
|
|
||
| self.pump_response_stream(stream, &publisher, payload_codec) | ||
| .await; | ||
| if let Err(err) = publisher.finish().await { |
There was a problem hiding this comment.
Nit: the success path above still ignores the prologue result (let _result = publisher.send_prologue(None).await;). Connection-failure cases end up killed and are handled gracefully here, but a non-kill send failure only surfaces later as the misleading "response mux stream finished before its prologue". Worth logging the prologue error at its point of failure.
|
This PR is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 5 days. |
1 participant
Summary
tcp_response_mux_v1transport0enables opportunistic batching and the maximum is 100 msThis is a coordinated transport cutover: response connections do not fall back to the former dedicated protocol. Frontends and workers must be upgraded together.
Design
Protocol and batching
MuxCodechandles the entire connection.ConnectionHellocontains the protocol version and frontend UUID, and the frontend returns an emptyConnectionReady. Subsequent frames use a 21-byte header containing payload length, frame kind, and logical-stream UUID.Successful prologues and resets have empty payloads. Failed prologues carry UTF-8 error text.
PrologueandResetuse an urgent lane;DataandEnduse an ordered lane.Endis the sole completion acknowledgement and flushes preceding data for its stream.Data flushes at the earliest of:
DYN_TCP_RESPONSE_BATCH_INTERVAL_MS(default1ms)DYN_TCP_RESPONSE_BATCH_MAX_BYTES(default65536)DYN_TCP_RESPONSE_BATCH_MAX_FRAMES(default64)Backpressure and ownership
Each stream has eight ordered-frame permits and byte credits controlled by
DYN_TCP_RESPONSE_STREAM_WINDOW_BYTES(default262144). A connection-wide queued-byte semaphore bounds userspace memory. TCP supplies physical-connection backpressure; there is no additional connection credit window.The connection writer exclusively owns batching and socket writes. Producers submit concrete frames through bounded channels without trait-object or boxed-future dispatch. Frontend
Datarouting stays in a connection-localHashMap, so decoded data frames do not touch the global lifecycle map.Pending and active response entries share one lifecycle
DashMap, with pending-to-active replacement performed atomically. Host pools live for the worker-process lifetime, and one process-wide maintenance task warms new frontend pools to four connections and repairs failed connections.Performance qualification
Final simplification versus mandatory-mux baseline
Three interleaved, freshly restarted repetitions at 8 workers / concurrency 1,024, ISL 1,024, OSL 256, and mocker speedup 10:
06d22825e9All six runs completed without request errors, cancellations, resets, or connection-failure fanout. Additional very-large-scale qualification also preserved performance after the final simplification.
Packet reduction versus dedicated responses
Earlier 16-worker / concurrency 2,048 qualification of the 1 ms mux mode measured 0.04705 response segments/token and 0.04708 response writes/token, versus 1.04417 segments/token and an estimated 1.00762 writes/token for dedicated responses. That is 22.2x fewer response segments and 21.4x fewer response writes, with exactly 64 mux response sockets. The final simplification changed writes/token by only +1.88% relative to the mandatory-mux baseline.
For comparison, the 5 ms interval produced larger batches (35.8 frames/write versus 21.5 at 1 ms) and 37.1x fewer response segments than dedicated responses, quantifying the additional packet-reduction/latency tradeoff available through configuration.
Validation
cargo fmt --allgit diff --checkcargo test -p dynamo-runtime --lib pipeline::network::tcp -- --nocapture— 50 passedcargo test -p dynamo-runtime --lib metrics::response_mux -- --nocapture— 1 passedcargo test -p dynamo-runtime --lib -- --skip transport_resolution_falls_back_when_selected_instance_disappears— 480 passed, 2 ignored, 1 known unrelated hang skippedcargo clippy -p dynamo-runtime --all-targets -- -D warningspre-commit run --all-filesfern check— 0 errorsfern docs broken-linksThe checked-in
benchmarks/frontend/scripts/run_perf.shis unchanged by this PR.Summary by CodeRabbit
New Features
Documentation