From ba9866cf0a1c1e9ec7dd58dfeb1e791481b27f02 Mon Sep 17 00:00:00 2001 From: Richard Markiewicz Date: Fri, 21 Aug 2026 10:14:41 -0400 Subject: [PATCH] fix(jetsocat): flush JMUX messages without waiting on a timer --- Cargo.lock | 5 + crates/jmux-proxy/src/lib.rs | 74 +++++-- jetsocat/src/listener.rs | 11 +- jetsocat/src/utils.rs | 33 ++- testsuite/Cargo.toml | 7 + testsuite/tests/jmux_flow_control.rs | 301 +++++++++++++++++++++++++++ testsuite/tests/main.rs | 1 + 7 files changed, 411 insertions(+), 21 deletions(-) create mode 100644 testsuite/tests/jmux_flow_control.rs diff --git a/Cargo.lock b/Cargo.lock index f2fa1f931..b508500bb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7512,10 +7512,15 @@ dependencies = [ "anyhow", "assert_cmd", "base64 0.23.1", + "bytes 1.12.1", "dynosaur", "escargot", "expect-test", "fastrand", + "http-body-util", + "hyper 1.10.1", + "hyper-util", + "jmux-proxy", "libsql", "mcp-proxy", "network-scanner", diff --git a/crates/jmux-proxy/src/lib.rs b/crates/jmux-proxy/src/lib.rs index 7373fce85..50b262aa1 100644 --- a/crates/jmux-proxy/src/lib.rs +++ b/crates/jmux-proxy/src/lib.rs @@ -41,7 +41,21 @@ pub use self::event::{EventOutcome, TrafficEvent, TransportProtocol}; const MAXIMUM_PACKET_SIZE_IN_BYTES: u16 = 4 * 1024; // 4 kiB const WINDOW_ADJUSTMENT_THRESHOLD: u32 = 4 * 1024; // 4 kiB -const JMUX_FLUSH_DELAY: core::time::Duration = core::time::Duration::from_millis(10); + +/// Backstop for flushing buffered messages, measured from the first unflushed byte. +/// +/// Never reset by later messages, so a steady stream cannot postpone a flush indefinitely. +/// Tokio's timer granularity is around a millisecond, so this bounds latency rather than +/// providing fine control; the drain check below is what keeps latency low in practice. +const JMUX_FLUSH_COALESCING_WINDOW: core::time::Duration = core::time::Duration::from_millis(1); + +/// Minimum spacing between flushes triggered by the send queue running dry. +/// +/// Flushing every time the queue drains is ideal for latency but ruinous for throughput: a +/// relay's producer is paced by the network, so under a bulk transfer the queue drains +/// constantly and each drain would write out a partial buffer. Spacing those flushes lets +/// bulk traffic keep filling the write buffer while a lone message still goes out promptly. +const JMUX_FLUSH_MIN_SPACING: core::time::Duration = core::time::Duration::from_micros(50); // The JMUX channel will require at most `MAXIMUM_PACKET_SIZE_IN_BYTES × JMUX_MESSAGE_CHANNEL_SIZE` bytes to be kept alive. const JMUX_MESSAGE_MPSC_CHANNEL_SIZE: usize = 512; @@ -385,28 +399,58 @@ impl JmuxSenderTask { let mut jmux_writer = tokio::io::BufWriter::with_capacity(16 * 1024, jmux_writer); let mut buf = bytes::BytesMut::new(); let mut needs_flush = false; - let flush_timer = tokio::time::sleep(JMUX_FLUSH_DELAY); - tokio::pin!(flush_timer); + // `None` until the first flush, so the first message out is never held back. + let mut last_flush: Option = None; + let flush_deadline = tokio::time::sleep(JMUX_FLUSH_COALESCING_WINDOW); + tokio::pin!(flush_deadline); - loop { + 'outer: loop { tokio::select! { msg = msg_to_send_rx.recv() => { - let Some(msg) = msg else { + let Some(mut msg) = msg else { break; }; - trace!(?msg, "Send channel message"); + // Write out everything already queued before considering a flush, so that + // bursts are coalesced into as few writes as possible. + // INVARIANT: `msg` always holds a message that has not been encoded yet. + loop { + trace!(?msg, "Send channel message"); + + buf.clear(); + msg.encode(&mut buf)?; + + jmux_writer.write_all(&buf).await?; + + if !needs_flush { + flush_deadline + .as_mut() + .reset(tokio::time::Instant::now() + JMUX_FLUSH_COALESCING_WINDOW); + needs_flush = true; + } + + match msg_to_send_rx.try_recv() { + Ok(next) => msg = next, + Err(mpsc::error::TryRecvError::Empty) => break, + Err(mpsc::error::TryRecvError::Disconnected) => break 'outer, + } + } - buf.clear(); - msg.encode(&mut buf)?; + // The queue ran dry, so there is nothing left to batch with: flush, unless + // a flush just happened and more traffic is plainly still flowing. + let flushed_recently = last_flush + .is_some_and(|instant| instant.elapsed() < JMUX_FLUSH_MIN_SPACING); - jmux_writer.write_all(&buf).await?; - needs_flush = true; - flush_timer.as_mut().reset(tokio::time::Instant::now() + JMUX_FLUSH_DELAY); + if !flushed_recently { + jmux_writer.flush().await?; + needs_flush = false; + last_flush = Some(tokio::time::Instant::now()); + } } - _ = flush_timer.as_mut(), if needs_flush => { + _ = flush_deadline.as_mut(), if needs_flush => { jmux_writer.flush().await?; needs_flush = false; + last_flush = Some(tokio::time::Instant::now()); } } } @@ -1196,6 +1240,12 @@ impl StreamResolverTask { for socket_addr in socket_addrs { match TcpStream::connect(socket_addr).await { Ok(stream) => { + // Nagle's algorithm is deliberately left enabled here. Disabling it + // costs about 25% of bulk throughput, because `DataWriterTask` writes + // every ~4 kiB chunk straight to this socket with no buffering, so + // each one would go out as its own undersized segment. Buffering + // those writes first would make `TCP_NODELAY` affordable. + // Update channel with resolved target IP and connect time. channel.target_ip = Some(socket_addr.ip()); channel.connect_at = SystemTime::now(); diff --git a/jetsocat/src/listener.rs b/jetsocat/src/listener.rs index 2b96d998f..ef5165e09 100644 --- a/jetsocat/src/listener.rs +++ b/jetsocat/src/listener.rs @@ -288,7 +288,16 @@ where loop { match listener.accept().await { - Ok((stream, addr)) => processor(stream, addr), + Ok((stream, addr)) => { + // Disable Nagle's algorithm: the traffic relayed here is shaped by the + // application on the other side, and delaying a sub-MSS segment until the + // previous one is acknowledged only adds latency to its round trips. + if let Err(error) = stream.set_nodelay(true) { + warn!(%error, %addr, "Couldn’t set TCP_NODELAY on accepted stream"); + } + + processor(stream, addr) + } Err(error) => { error!(%error, "Couldn’t accept next TCP stream"); break; diff --git a/jetsocat/src/utils.rs b/jetsocat/src/utils.rs index a475fe7cf..228dcf9fb 100644 --- a/jetsocat/src/utils.rs +++ b/jetsocat/src/utils.rs @@ -25,6 +25,25 @@ async fn resolve_dest_addr(dest_addr: DestAddr) -> anyhow::Result { } } +/// Connects to `addr` with Nagle's algorithm disabled. +/// +/// jetsocat relays traffic whose write pattern is dictated by its peers, so holding back a +/// sub-MSS segment until the previous one is acknowledged only adds latency to every round +/// trip crossing the pipe. Coalescing already happens upstream, where JMUX messages are +/// batched before being written out. +pub(crate) async fn connect_nodelay(addr: A) -> std::io::Result +where + A: tokio::net::ToSocketAddrs, +{ + let stream = TcpStream::connect(addr).await?; + + if let Err(error) = stream.set_nodelay(true) { + warn!(%error, "Couldn’t set TCP_NODELAY"); + } + + Ok(stream) +} + macro_rules! impl_tcp_connect { ($req_addr:expr, $proxy_cfg:expr, $output_ty:ty, | $stream:ident | $operation:block) => {{ use proxy_socks::{Socks4Stream, Socks5Stream}; @@ -34,15 +53,14 @@ macro_rules! impl_tcp_connect { ty: ProxyType::Socks4, addr: proxy_addr, }) => { - let $stream = - Socks4Stream::connect(TcpStream::connect(proxy_addr).await?, $req_addr, "jetsocat").await?; + let $stream = Socks4Stream::connect(connect_nodelay(proxy_addr).await?, $req_addr, "jetsocat").await?; $operation.await } Some(ProxyConfig { ty: ProxyType::Socks5, addr: proxy_addr, }) => { - let $stream = Socks5Stream::connect(TcpStream::connect(proxy_addr).await?, $req_addr).await?; + let $stream = Socks5Stream::connect(connect_nodelay(proxy_addr).await?, $req_addr).await?; $operation.await } Some(ProxyConfig { @@ -50,11 +68,11 @@ macro_rules! impl_tcp_connect { addr: proxy_addr, }) => { // unknown SOCKS version, try SOCKS5 first and then SOCKS4 - match Socks5Stream::connect(TcpStream::connect(&proxy_addr).await?, &$req_addr).await { + match Socks5Stream::connect(connect_nodelay(&proxy_addr).await?, &$req_addr).await { Ok($stream) => $operation.await, Err(_) => { let $stream = - Socks4Stream::connect(TcpStream::connect(proxy_addr).await?, $req_addr, "jetsocat").await?; + Socks4Stream::connect(connect_nodelay(proxy_addr).await?, $req_addr, "jetsocat").await?; $operation.await } } @@ -67,14 +85,13 @@ macro_rules! impl_tcp_connect { ty: ProxyType::Https, addr: proxy_addr, }) => { - let $stream = - proxy_http::ProxyStream::connect(TcpStream::connect(proxy_addr).await?, $req_addr).await?; + let $stream = proxy_http::ProxyStream::connect(connect_nodelay(proxy_addr).await?, $req_addr).await?; $operation.await } None => { let dest_addr = resolve_dest_addr($req_addr.to_dest_addr().with_context(|| "invalid target address")?).await?; - let $stream = TcpStream::connect(dest_addr).await?; + let $stream = connect_nodelay(dest_addr).await?; $operation.await } }; diff --git a/testsuite/Cargo.toml b/testsuite/Cargo.toml index 895e12531..2680cec35 100644 --- a/testsuite/Cargo.toml +++ b/testsuite/Cargo.toml @@ -38,6 +38,13 @@ mcp-proxy.path = "../crates/mcp-proxy" network-scanner = { path = "../crates/network-scanner", features = ["test-utils"] } network-scanner-proto = { path = "../crates/network-scanner-proto" } rstest = "0.25" +jmux-proxy = { path = "../crates/jmux-proxy" } +# `test-util` provides the paused clock used to assert flush behavior deterministically. +tokio = { version = "1", features = ["test-util"] } +hyper = { version = "1", features = ["server", "client", "http1", "http2"] } +hyper-util = { version = "0.1", features = ["tokio"] } +http-body-util = "0.1" +bytes = "1" serde_json = "1" sysevent.path = "../crates/sysevent" tempfile = "3" diff --git a/testsuite/tests/jmux_flow_control.rs b/testsuite/tests/jmux_flow_control.rs new file mode 100644 index 000000000..3000b679f --- /dev/null +++ b/testsuite/tests/jmux_flow_control.rs @@ -0,0 +1,301 @@ +//! Checks that JMUX does not penalize round trips. +//! +//! HTTP/2 gates an upload on its per-stream flow control window: the client may only have +//! `SETTINGS_INITIAL_WINDOW_SIZE` bytes of DATA in flight before it must wait for a +//! `WINDOW_UPDATE` to come back. RFC 9113 puts the default at 65535 bytes, and a server that +//! never raises it turns a large upload into a long sequence of round trips rather than one +//! continuous stream. HTTP/1.1 has no such gate and streams the body in one go. +//! +//! That makes an HTTP/2 upload a sensitive probe for latency added *per round trip* by +//! anything relaying the connection. A JMUX sender that waits on a timer before flushing +//! small messages (a `WINDOW_UPDATE` never fills a write buffer on its own) turns every one +//! of those round trips into a stall, while leaving bulk HTTP/1.1 transfers untouched. +//! +//! These tests pin that behavior down: same payload, same JMUX pipe, HTTP/1.1 versus HTTP/2. + +use core::time::Duration; +use std::net::SocketAddr; +use std::time::Instant; + +use bytes::Bytes; +use http_body_util::{BodyExt as _, Empty, Full}; +use hyper_util::rt::{TokioExecutor, TokioIo}; +use jmux_proxy::{DestinationUrl, JmuxApiRequest, JmuxApiResponse, JmuxConfig, JmuxProxy}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::{mpsc, oneshot}; + +/// The HTTP/2 default initial window from RFC 9113, which a server has to explicitly raise. +/// Left at the default, it makes an upload round trip every 64 KiB. +/// +/// Servers really do leave it there: vCenter is one, and Broadcom documents raising it as the +/// remedy for slow Content Library uploads over high-latency links. +/// See . +const STREAM_WINDOW: u32 = 64 * 1024; + +/// Kept comfortably above `STREAM_WINDOW` so the per-stream window stays the binding limit. +const CONNECTION_WINDOW: u32 = 1024 * 1024; + +const PAYLOAD_SIZE: usize = 8 * 1024 * 1024; + +/// How much slower the same upload may be through JMUX than straight to the server. +/// +/// Expressed as a ratio rather than a wall-clock budget on purpose. An absolute budget has to +/// be loose enough for slow CI, which makes it too loose to catch the regression: with a 10 ms +/// per-flush delay over the 128 round trips this payload takes, the old sender only needs to +/// lose 10-20 ms per round trip to blow past any budget generous enough to be safe, and on a +/// platform where the timer fires closer to its nominal delay it could sneak under. Scaling +/// against the direct measurement normalizes for machine speed instead. +/// +/// Observed ratios: ~2.5x with the current sender, ~70x with the 10 ms idle-flush timer. +const MAX_JMUX_OVERHEAD_FACTOR: u32 = 8; + +/// Bound for the HTTP/1.1 control case, which is not round-trip gated and so is not the +/// sensitive measurement. It only needs to be loose enough not to flake. +const H1_BUDGET: Duration = Duration::from_secs(2); + +#[derive(Copy, Clone, PartialEq, Eq, Debug)] +enum Proto { + Http1, + Http2, +} + +/// Spawns a server that drains request bodies and replies with an empty 200. +async fn spawn_server(proto: Proto) -> SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + tokio::spawn(async move { + loop { + let (stream, _) = listener.accept().await.unwrap(); + let io = TokioIo::new(stream); + + tokio::spawn(async move { + let service = hyper::service::service_fn(|req: hyper::Request| async move { + let mut body = req.into_body(); + + // Consume the body as it arrives. This is what releases flow control + // credit back to the peer, so it must not buffer the whole thing. + while let Some(frame) = body.frame().await { + frame?; + } + + Ok::<_, hyper::Error>(hyper::Response::new(Empty::::new())) + }); + + match proto { + Proto::Http1 => hyper::server::conn::http1::Builder::new() + .serve_connection(io, service) + .await + .map_err(|error| format!("{error}")), + Proto::Http2 => hyper::server::conn::http2::Builder::new(TokioExecutor::new()) + .initial_stream_window_size(STREAM_WINDOW) + .initial_connection_window_size(CONNECTION_WINDOW) + .serve_connection(io, service) + .await + .map_err(|error| format!("{error}")), + } + }); + } + }); + + addr +} + +/// Uploads `PAYLOAD_SIZE` bytes to `addr` and returns how long it took. +async fn upload(addr: SocketAddr, proto: Proto) -> Duration { + let io = TokioIo::new(TcpStream::connect(addr).await.unwrap()); + + let request = hyper::Request::builder() + .method("POST") + .uri(format!("http://{addr}/upload")) + .body(Full::new(Bytes::from(vec![0u8; PAYLOAD_SIZE]))) + .unwrap(); + + let started_at = Instant::now(); + + let response = match proto { + Proto::Http1 => { + let (mut sender, conn) = hyper::client::conn::http1::handshake(io).await.unwrap(); + tokio::spawn(conn); + sender.send_request(request).await.unwrap() + } + Proto::Http2 => { + let (mut sender, conn) = hyper::client::conn::http2::handshake(TokioExecutor::new(), io) + .await + .unwrap(); + tokio::spawn(conn); + sender.send_request(request).await.unwrap() + } + }; + + assert!(response.status().is_success(), "upload failed: {}", response.status()); + response.into_body().collect().await.unwrap(); + + started_at.elapsed() +} + +/// Runs a JMUX proxy pair and returns a local address forwarding to `target` through it. +/// +/// This mirrors the deployed topology — jetsocat on one end exposing a local listener, the +/// Gateway on the other end connecting out to the target — with the two ends wired together +/// by a loopback TCP connection standing in for the WebSocket pipe. +async fn spawn_jmux_forward(target: SocketAddr) -> SocketAddr { + let pipe_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let pipe_addr = pipe_listener.local_addr().unwrap(); + let dialing = tokio::spawn(async move { TcpStream::connect(pipe_addr).await.unwrap() }); + let (gateway_end, _) = pipe_listener.accept().await.unwrap(); + let client_end = dialing.await.unwrap(); + + // The end that accepts channels and connects out to the target. + let (reader, writer) = gateway_end.into_split(); + tokio::spawn( + JmuxProxy::new(Box::new(reader), Box::new(writer)) + .with_config(JmuxConfig::permissive()) + .run(), + ); + + // The end that opens channels on behalf of local connections. + let (api_request_tx, api_request_rx) = mpsc::channel(16); + let (reader, writer) = client_end.into_split(); + tokio::spawn( + JmuxProxy::new(Box::new(reader), Box::new(writer)) + .with_requester_api(api_request_rx) + .run(), + ); + + let local_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let local_addr = local_listener.local_addr().unwrap(); + let destination_url = DestinationUrl::new("tcp", &target.ip().to_string(), target.port()); + + tokio::spawn(async move { + loop { + let (stream, _) = local_listener.accept().await.unwrap(); + let api_request_tx = api_request_tx.clone(); + let destination_url = destination_url.clone(); + + tokio::spawn(async move { + let (api_response_tx, api_response_rx) = oneshot::channel(); + + api_request_tx + .send(JmuxApiRequest::OpenChannel { + destination_url, + api_response_tx, + }) + .await + .unwrap(); + + match api_response_rx.await.unwrap() { + JmuxApiResponse::Success { id } => { + api_request_tx + .send(JmuxApiRequest::Start { + id, + stream, + leftover: None, + }) + .await + .unwrap(); + } + JmuxApiResponse::Failure { id, reason_code } => { + panic!("channel {id} failed to open: {reason_code}") + } + } + }); + } + }); + + local_addr +} + +/// An HTTP/2 upload is gated on a 64 KiB window, so it round trips 128 times for this payload. +/// Any per-round-trip latency introduced by JMUX shows up here, multiplied. +#[tokio::test(flavor = "multi_thread")] +async fn http2_upload_through_jmux_is_not_round_trip_penalized() { + let server_addr = spawn_server(Proto::Http2).await; + let forward_addr = spawn_jmux_forward(server_addr).await; + + let direct = upload(server_addr, Proto::Http2).await; + let through_jmux = upload(forward_addr, Proto::Http2).await; + + println!("http2 direct={direct:?} through_jmux={through_jmux:?}"); + + let budget = direct * MAX_JMUX_OVERHEAD_FACTOR; + + assert!( + through_jmux < budget, + "HTTP/2 upload took {through_jmux:?} through JMUX versus {direct:?} direct, over the \ + {MAX_JMUX_OVERHEAD_FACTOR}x budget of {budget:?}; JMUX is likely delaying flow \ + control updates" + ); +} + +/// The HTTP/1.1 counterpart streams the body without gating, so it stays fast even when JMUX +/// delays small messages. Keeping it here documents *why* the HTTP/2 case is the sensitive one: +/// a regression that only this test catches is a round-trip regression, not a bandwidth one. +#[tokio::test(flavor = "multi_thread")] +async fn http1_upload_through_jmux_matches_direct() { + let server_addr = spawn_server(Proto::Http1).await; + let forward_addr = spawn_jmux_forward(server_addr).await; + + let direct = upload(server_addr, Proto::Http1).await; + let through_jmux = upload(forward_addr, Proto::Http1).await; + + println!("http1 direct={direct:?} through_jmux={through_jmux:?}"); + + assert!( + through_jmux < H1_BUDGET, + "HTTP/1.1 upload through JMUX took {through_jmux:?}, over the {H1_BUDGET:?} budget \ + (direct took {direct:?})" + ); +} + +/// Pins the sender's flush behavior directly, without depending on wall-clock timing. +/// +/// With the clock paused, tokio advances virtual time only when every task is idle. A sender +/// that parks on a timer before flushing therefore shows up as virtual time elapsing between +/// queueing a message and it reaching the pipe, on any machine and at any speed. A sender that +/// flushes once its queue is drained shows zero. +/// +/// This is the deterministic counterpart to the HTTP/2 test above: that one proves the +/// end-to-end effect on realistic traffic, this one fails for exactly one reason. +#[tokio::test(start_paused = true)] +async fn sender_flushes_without_advancing_the_clock() { + use tokio::io::AsyncReadExt as _; + + let (near, mut far) = tokio::io::duplex(64 * 1024); + let (near_reader, near_writer) = tokio::io::split(near); + let (api_request_tx, api_request_rx) = mpsc::channel(1); + + tokio::spawn( + JmuxProxy::new(Box::new(near_reader), Box::new(near_writer)) + .with_requester_api(api_request_rx) + .run(), + ); + + // Opening a channel queues a single small message, which is exactly the shape of traffic + // that never fills the sender's write buffer on its own. + let (api_response_tx, _api_response_rx) = oneshot::channel(); + api_request_tx + .send(JmuxApiRequest::OpenChannel { + destination_url: DestinationUrl::new("tcp", "127.0.0.1", 1), + api_response_tx, + }) + .await + .unwrap(); + + let started_at = tokio::time::Instant::now(); + + let mut buf = [0u8; 128]; + let read = tokio::time::timeout(Duration::from_secs(5), far.read(&mut buf)) + .await + .expect("sender never flushed the CHANNEL OPEN") + .unwrap(); + + let waited = started_at.elapsed(); + + assert!(read > 0, "sender flushed an empty write"); + assert!( + waited < Duration::from_millis(1), + "sender held the message for {waited:?} of virtual time before flushing; it is \ + waiting on a timer rather than flushing once its queue is drained" + ); +} diff --git a/testsuite/tests/main.rs b/testsuite/tests/main.rs index 5d8b1e6c5..a0fbbc54c 100644 --- a/testsuite/tests/main.rs +++ b/testsuite/tests/main.rs @@ -3,6 +3,7 @@ #![allow(clippy::print_stderr, reason = "test code uses print for diagnostics")] mod cli; +mod jmux_flow_control; mod mcp_proxy; mod network_scanner; mod sysevent;