fix(jetsocat): flush JMUX messages without waiting on a timer

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<T: AsyncWrite + Unpin + Send + 'static> JmuxSenderTask<T> {
         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<tokio::time::Instant> = 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();

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;

jetsocat/src/utils.rs

@@ -25,6 +25,25 @@ async fn resolve_dest_addr(dest_addr: DestAddr) -> anyhow::Result<SocketAddr> {
     }
 }
 
+/// 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<A>(addr: A) -> std::io::Result<TcpStream>
+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,27 +53,26 @@ 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 {
                 ty: ProxyType::Socks,
                 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
             }
         };

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"