diff --git a/crates/jmux-proxy/Cargo.toml b/crates/jmux-proxy/Cargo.toml index 76a96670e..b3db36dd0 100644 --- a/crates/jmux-proxy/Cargo.toml +++ b/crates/jmux-proxy/Cargo.toml @@ -33,4 +33,4 @@ bitvec = "1.0" # Test infrastructure for socket manipulation. socket2 = "0.5" # Enable extra tokio features. -tokio = { version = "1.52", features = ["rt-multi-thread"] } +tokio = { version = "1.52", features = ["rt-multi-thread", "test-util"] } diff --git a/crates/jmux-proxy/src/lib.rs b/crates/jmux-proxy/src/lib.rs index 4c5c7d00f..b2a91f03b 100644 --- a/crates/jmux-proxy/src/lib.rs +++ b/crates/jmux-proxy/src/lib.rs @@ -43,7 +43,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; @@ -428,28 +442,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()); } } } @@ -1313,6 +1357,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 + // write can leave an undersized tail instead of coalescing it with the + // next chunk. Buffering those writes first would make `TCP_NODELAY` + // affordable. channel.target_ip = Some(socket_addr.ip()); return Ok(Box::new(stream)); } @@ -1400,3 +1450,90 @@ fn is_really_an_error(original_error: &(dyn std::error::Error + 'static)) -> boo true } + +#[cfg(test)] +mod sender_tests { + use bytes::BytesMut; + use tokio::io::AsyncReadExt as _; + + use super::*; + + fn open_message(id: u32) -> (Message, BytesMut) { + let message = Message::open( + LocalChannelId::from(id), + MAXIMUM_PACKET_SIZE_IN_BYTES, + DestinationUrl::new("tcp", "127.0.0.1", 1), + ); + let mut encoded = BytesMut::new(); + message.encode(&mut encoded).expect("encode message"); + (message, encoded) + } + + async fn assert_immediate_flush( + msg_to_send_tx: &mpsc::Sender, + reader: &mut tokio::io::DuplexStream, + id: u32, + ) { + let (message, expected) = open_message(id); + msg_to_send_tx.send(message).await.expect("queue message"); + let started_at = tokio::time::Instant::now(); + let mut actual = vec![0; expected.len()]; + tokio::time::timeout(core::time::Duration::from_secs(5), reader.read_exact(&mut actual)) + .await + .expect("sender never flushed the message") + .expect("read flushed message"); + + assert_eq!(actual, expected); + assert_eq!( + started_at.elapsed(), + core::time::Duration::ZERO, + "sender waited on its flush deadline" + ); + } + + #[tokio::test(start_paused = true)] + async fn sender_flushes_a_drained_queue_with_bounded_coalescing() { + let (writer, mut reader) = tokio::io::duplex(1024); + let (msg_to_send_tx, msg_to_send_rx) = mpsc::channel(1); + let sender_task = tokio::spawn( + JmuxSenderTask { + jmux_writer: writer, + msg_to_send_rx, + } + .run(), + ); + + assert_immediate_flush(&msg_to_send_tx, &mut reader, 1).await; + + let (message, expected) = open_message(2); + msg_to_send_tx.send(message).await.expect("queue message"); + let mut actual = vec![0; expected.len()]; + let read = reader.read_exact(&mut actual); + tokio::pin!(read); + let before_deadline = JMUX_FLUSH_COALESCING_WINDOW + .checked_sub(core::time::Duration::from_nanos(1)) + .expect("coalescing window is nonzero"); + + tokio::select! { + biased; + result = &mut read => panic!("sender flushed within the minimum spacing: {result:?}"), + () = tokio::time::sleep(before_deadline) => {} + } + + tokio::time::advance(core::time::Duration::from_nanos(1)).await; + tokio::time::timeout(core::time::Duration::from_secs(5), &mut read) + .await + .expect("sender did not flush at the coalescing deadline") + .expect("read coalesced message"); + assert_eq!(actual, expected); + + tokio::time::advance(JMUX_FLUSH_MIN_SPACING).await; + assert_immediate_flush(&msg_to_send_tx, &mut reader, 3).await; + + drop(msg_to_send_tx); + sender_task + .await + .expect("sender task panicked") + .expect("sender task failed"); + } +} 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/pipe.rs b/jetsocat/src/pipe.rs index 94f0960ca..3a21df456 100644 --- a/jetsocat/src/pipe.rs +++ b/jetsocat/src/pipe.rs @@ -197,6 +197,10 @@ pub async fn open_pipe(mode: PipeMode, proxy_cfg: Option) -> Result info!(%peer_addr, "Accepted peer"); + if let Err(error) = socket.set_nodelay(true) { + warn!(%error, %peer_addr, "Couldn’t set TCP_NODELAY on accepted stream"); + } + Ok(Pipe { name: "tcp-listener", stream: Box::new(socket), @@ -326,6 +330,10 @@ pub async fn open_pipe(mode: PipeMode, proxy_cfg: Option) -> Result info!(%peer_addr, "Accepted peer"); + if let Err(error) = socket.set_nodelay(true) { + warn!(%error, %peer_addr, "Couldn’t set TCP_NODELAY on accepted stream"); + } + let ws = accept_async(socket) .await .with_context(|| "WebSocket handshake failed")?; 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/tests/cli/jetsocat.rs b/testsuite/tests/cli/jetsocat.rs index 571a25fc2..565e349b7 100644 --- a/testsuite/tests/cli/jetsocat.rs +++ b/testsuite/tests/cli/jetsocat.rs @@ -1406,6 +1406,118 @@ async fn mcp_proxy_terminated_on_broken_pipe() { ); } +/// Relays windowed data through the jetsocat CLI and measures how quickly credits return. +async fn run_jmux_flow_control_case(use_websocket: bool) -> Duration { + use std::time::Instant; + + use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _}; + use tokio::net::{TcpListener, TcpStream}; + use tokio::time::timeout; + + const WINDOW_SIZE: usize = 64 * 1024; + const MEASURED_WINDOWS: usize = 256; + const CREDIT: u8 = 1; + const HANG_TIMEOUT: Duration = Duration::from_secs(20); + + let ports = find_unused_ports(2); + let jmux_server_port = ports[0]; + let proxy_port = ports[1]; + + let target_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let target_port = target_listener.local_addr().unwrap().port(); + let target_task = tokio::spawn(async move { + let (mut stream, _) = target_listener.accept().await.unwrap(); + let mut window = vec![0; WINDOW_SIZE]; + + // Model receiver-driven flow control: each complete window releases one byte of credit. + for _ in 0..=MEASURED_WINDOWS { + stream.read_exact(&mut window).await.unwrap(); + stream.write_all(&[CREDIT]).await.unwrap(); + } + }); + + let jmux_pipe = if use_websocket { + format!("ws-listen://127.0.0.1:{jmux_server_port}") + } else { + format!("tcp-listen://127.0.0.1:{jmux_server_port}") + }; + let mut jmux_server = jetsocat_tokio_cmd() + .env( + "JETSOCAT_ARGS", + format!("jmux-proxy {jmux_pipe} --allow-all --no-proxy"), + ) + .kill_on_drop(true) + .spawn() + .expect("failed to start JMUX server"); + + wait_for_port_bound(jmux_server_port).await.expect("JMUX server ready"); + + let peer_pipe = if use_websocket { + format!("ws://127.0.0.1:{jmux_server_port}") + } else { + format!("tcp://127.0.0.1:{jmux_server_port}") + }; + let mut jmux_client = jetsocat_tokio_cmd() + .env( + "JETSOCAT_ARGS", + format!("jmux-proxy {peer_pipe} tcp-listen://127.0.0.1:{proxy_port}/127.0.0.1:{target_port} --no-proxy"), + ) + .kill_on_drop(true) + .spawn() + .expect("failed to start JMUX client"); + + wait_for_port_bound(proxy_port).await.expect("JMUX client proxy ready"); + + let transfer = timeout(HANG_TIMEOUT, async { + let mut stream = TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap(); + let window = vec![0; WINDOW_SIZE]; + let mut credit = [0]; + + // Warm up the tunnel so connection setup is excluded from the measurement. + stream.write_all(&window).await.unwrap(); + stream.read_exact(&mut credit).await.unwrap(); + assert_eq!(credit, [CREDIT]); + + let started_at = Instant::now(); + + for _ in 0..MEASURED_WINDOWS { + stream.write_all(&window).await.unwrap(); + stream.read_exact(&mut credit).await.unwrap(); + assert_eq!(credit, [CREDIT]); + } + + started_at.elapsed() + }) + .await; + + let _ = jmux_client.start_kill(); + let _ = jmux_server.start_kill(); + let _ = jmux_client.wait().await; + let _ = jmux_server.wait().await; + + let elapsed = transfer.expect("flow-controlled transfer timed out"); + target_task.await.expect("target server task panicked"); + elapsed +} + +#[tokio::test] +async fn jmux_flow_control_credits_are_not_delayed() { + const TRANSFER_BUDGET: Duration = Duration::from_secs(2); + + for use_websocket in [false, true] { + let elapsed = run_jmux_flow_control_case(use_websocket).await; + let transport = if use_websocket { "WebSocket" } else { "TCP" }; + + println!("{transport} JMUX flow-controlled transfer completed in {elapsed:?}"); + + assert!( + elapsed < TRANSFER_BUDGET, + "{transport} JMUX took {elapsed:?} to relay flow-controlled traffic, exceeding the \ + {TRANSFER_BUDGET:?} budget" + ); + } +} + /// SOCKS5 client → SOCKS5 listener → JMUX tunnel → TCP echo server. #[rstest] #[tokio::test]