perf(durable-streams-rust): serve caught-up long-polls from the epoll reactor

.changeset/long-poll-reactor.md

@@ -0,0 +1,5 @@
+---
+'@electric-ax/durable-streams-server-rust': patch
+---
+
+Serve caught-up long-poll waits from the epoll reactor too, sharing the machinery with SSE: a blocked poll hands its socket to the reactor (freeing the connection task's future) and is handed back for keep-alive after the response. Linux only; other platforms keep the inline wait.

packages/durable-streams-rust/src/api.rs

@@ -48,6 +48,25 @@ pub struct SseReg {
     pub client_cursor: Option<u64>,
 }
 
+/// Caught-up long-poll registration: everything the reactor needs to wait at the
+/// tail and write the single long-poll response once data arrives (or on close /
+/// timeout). Constructed only on Linux, by `handle_long_poll` when the poll would
+/// block on a reactor-eligible stream (root, no tiering); see `Body::LongPollWait`.
+#[cfg(target_os = "linux")]
+pub struct LongPollReg {
+    pub st: Arc<crate::store::StreamState>,
+    /// Offset the poll is caught up at: the response delivers `from..tail` once
+    /// the tail advances past it.
+    pub from: u64,
+    pub client_cursor: Option<u64>,
+    /// Honour HTTP keep-alive: when set, the reactor hands the socket back to the
+    /// async connection loop after writing the response instead of closing it, so
+    /// a polling client's next request reuses the connection (no per-poll churn).
+    /// Filled in by the engine at hand-off — the handler can't see the
+    /// connection's keep-alive state.
+    pub keep_alive: bool,
+}
+
 /// A streamed (chunked) response body plus an abort signal.
 ///
 /// `rx` carries the body frames. `failed` is set by the producer if it ends the
@@ -128,6 +147,13 @@ pub enum Body {
     /// Pull-based streaming body driven inline on the connection task (SSE).
     /// No spawned task, no channel — see `EventSource`.
     Sse(Box<dyn EventSource>),
+    /// Linux: a caught-up long-poll that would block at the tail. The engine
+    /// hands the socket to the epoll reactor (freeing this connection task's
+    /// future for the up-to-30s wait, the dominant per-poll cost) instead of
+    /// awaiting inline; the reactor writes the single response when data arrives
+    /// or on close/timeout. Never serialized by `write_response`.
+    #[cfg(target_os = "linux")]
+    LongPollWait(LongPollReg),
     /// Byte ranges of data files plus optional framing (JSON `[` / `]`).
     /// Total body length is prefix + Σ segment.len + suffix.
     FileRange {
@@ -150,6 +176,8 @@ impl Body {
             Body::Full(b) => Some(b.len() as u64),
             Body::Channel(_) => None,
             Body::Sse(_) => None,
+            #[cfg(target_os = "linux")]
+            Body::LongPollWait(_) => None,
             Body::FileRange {
                 segments,
                 prefix,

packages/durable-streams-rust/src/engine_raw.rs

@@ -114,7 +114,36 @@ pub async fn drain(grace: std::time::Duration) {
     let _ = tokio::time::timeout(grace, sema.acquire_many(MAX_CONNECTIONS as u32)).await;
 }
 
+/// Resume an async connection loop on a socket the long-poll reactor handed back
+/// after writing a keep-alive response. Called from a reactor thread (NOT a tokio
+/// worker), so it spawns onto the runtime handle captured in `serve`. `leftover`
+/// is any data the client pipelined behind its poll. On any failure (no runtime —
+/// only in unit tests — or the fd can't be adopted) everything is dropped, which
+/// closes the socket and releases the permit.
+#[cfg(target_os = "linux")]
+pub(crate) fn resume_after_longpoll(
+    store: Arc<Store>,
+    std_stream: std::net::TcpStream,
+    leftover: BytesMut,
+    permit: tokio::sync::OwnedSemaphorePermit,
+) {
+    let Some(handle) = CONN_RUNTIME.get() else { return };
+    handle.spawn(async move {
+        if let Ok(stream) = TcpStream::from_std(std_stream) {
+            let _ = conn_loop(store, stream, leftover, permit).await;
+        }
+    });
+}
+
+/// Runtime handle captured at startup so the long-poll reactor (running on its
+/// own non-tokio threads) can hand a socket back to the async connection loop
+/// for keep-alive. See `resume_after_longpoll`.
+#[cfg(target_os = "linux")]
+static CONN_RUNTIME: std::sync::OnceLock<tokio::runtime::Handle> = std::sync::OnceLock::new();
+
 pub async fn serve(store: Arc<Store>, listener: TcpListener) {
+    #[cfg(target_os = "linux")]
+    let _ = CONN_RUNTIME.set(tokio::runtime::Handle::current());
     let conns = conn_limiter().clone();
     loop {
         let (stream, _) = match listener.accept().await {
@@ -178,26 +207,27 @@ pub async fn serve(store: Arc<Store>, listener: TcpListener) {
             // detached SSE streaming task (keeping the connection counted while
             // the big conn_loop future is freed); otherwise it is released when
             // conn_loop returns.
-            let _ = conn_loop(store, stream, permit).await;
+            let _ = conn_loop(store, stream, BytesMut::with_capacity(4 * 1024), permit).await;
         });
     }
 }
 
 async fn conn_loop(
     store: Arc<Store>,
     mut stream: TcpStream,
+    mut buf: BytesMut,
     permit: tokio::sync::OwnedSemaphorePermit,
 ) -> std::io::Result<()> {
     const BAD_REQUEST: &[u8] =
         b"HTTP/1.1 400 Bad Request\r\ncontent-length: 0\r\nconnection: close\r\n\r\n";
     const TOO_LARGE: &[u8] =
         b"HTTP/1.1 413 Payload Too Large\r\ncontent-length: 0\r\nconnection: close\r\n\r\n";
-    // Initial capacity for the per-connection read buffer. Small on purpose:
-    // it persists for the whole (possibly long-lived, idle) connection, and
-    // `read_buf` grows it on demand, so a request head/body larger than this
-    // still parses — it just triggers a reserve. 4 KiB covers a typical request
-    // head in one read while keeping idle SSE subscribers cheap.
-    let mut buf = BytesMut::with_capacity(4 * 1024);
+    // The per-connection read buffer is passed in (a fresh 4 KiB for a new
+    // connection; for a connection handed back from the long-poll reactor after
+    // keep-alive, any bytes the client pipelined behind its poll). Small on
+    // purpose: it persists for the whole (possibly long-lived, idle) connection,
+    // and `read_buf` grows it on demand, so a larger request head/body still
+    // parses — it just triggers a reserve.
     loop {
         // ---- read request head ----
         let head = match http1::read_head(&mut stream, &mut buf).await? {
@@ -331,6 +361,25 @@ async fn conn_loop(
             body,
         };
         let resp = handlers::handle(store.clone(), req).await;
+        // Long-poll (Linux): a caught-up poll on a reactor-eligible stream comes
+        // back as `LongPollWait` instead of having blocked inline. Hand the socket
+        // to the reactor — which parks for the wait without holding this conn_loop
+        // future, writes the single response, and (for keep-alive) hands the socket
+        // back so the next poll reuses the connection. `buf` carries any bytes the
+        // client pipelined behind the poll, fed back on resume.
+        #[cfg(target_os = "linux")]
+        let mut resp = resp;
+        #[cfg(target_os = "linux")]
+        if !is_head && matches!(resp.body, Body::LongPollWait(_)) {
+            let Body::LongPollWait(mut reg) =
+                std::mem::replace(&mut resp.body, Body::Empty)
+            else {
+                unreachable!()
+            };
+            reg.keep_alive = keep_alive;
+            crate::reactor::register_long_poll(stream, reg, buf, store, permit);
+            return Ok(());
+        }
         // SSE: hand the connection off to a DETACHED, minimal streaming task and
         // return. An SSE subscriber parks for up to SSE_MAX_DURATION; driving it
         // inline would keep this whole `conn_loop` future (sized to the largest
@@ -351,7 +400,7 @@ async fn conn_loop(
                 if let Some(reg) = src.reactor_reg() {
                     let mut head: Vec<u8> = Vec::with_capacity(256);
                     http1::write_head(&mut head, resp.status, &resp.headers, None, keep_alive);
-                    crate::sse_reactor::register(stream, head, reg, permit);
+                    crate::reactor::register(stream, head, reg, permit);
                     return Ok(());
                 }
             }
@@ -398,6 +447,10 @@ async fn write_response(
     }
 
     match resp.body {
+        // A caught-up long-poll on Linux is intercepted in `conn_loop` and handed
+        // to the reactor before reaching here, so this is never serialized inline.
+        #[cfg(target_os = "linux")]
+        Body::LongPollWait(_) => unreachable!("LongPollWait is handed to the reactor"),
         Body::Empty => {
             stream.write_all(&head).await?;
         }

packages/durable-streams-rust/src/handlers.rs

@@ -103,7 +103,7 @@ pub(crate) mod test_support {
     }
 }
 
-fn long_poll_timeout_dur() -> Duration {
+pub(crate) fn long_poll_timeout_dur() -> Duration {
     Duration::from_millis(LONG_POLL_TIMEOUT_MS.load(std::sync::atomic::Ordering::Relaxed))
 }
 
@@ -776,7 +776,7 @@ fn publish_durable_tail(st: &StreamState, tail: u64, wire: &Bytes) {
     st.tail_tx.send_replace(Tail { bytes: tail, closed });
     // Wake any reactor-served subscribers of this stream (no-op when none).
     #[cfg(target_os = "linux")]
-    crate::sse_reactor::wake_stream(st);
+    crate::reactor::wake_stream(st);
 }
 
 // ---------- POST (append) ----------
@@ -1155,7 +1155,7 @@ async fn handle_append_inner(store: Arc<Store>, req: Req, path: String) -> (Resp
         };
         st.tail_tx.send_replace(Tail { bytes: tail, closed: true });
         #[cfg(target_os = "linux")]
-        crate::sse_reactor::wake_stream(&st);
+        crate::reactor::wake_stream(&st);
     } else {
         st.schedule_meta_flush();
     }
@@ -1355,7 +1355,7 @@ where
     };
     st.tail_tx.send_replace(Tail { bytes: tail, closed });
     #[cfg(target_os = "linux")]
-    crate::sse_reactor::wake_stream(&st);
+    crate::reactor::wake_stream(&st);
 
     // Shared response builder: captures `st`, `producer` by ref.
     let make_ok = || {
@@ -1768,6 +1768,26 @@ async fn handle_long_poll(
         return long_poll_close(t0.bytes, cursor);
     }
 
+    // Caught up at the live tail with the stream still open: this poll must wait.
+    // On Linux, hand the wait to the epoll reactor (which frees this connection
+    // task's future for the up-to-30s park — the dominant per-poll cost) when the
+    // stream is reactor-eligible; the reactor writes the single response on
+    // data/close/timeout and hands the socket back for keep-alive. Cold/forked/
+    // tiered streams and non-Linux fall through to the inline wait below.
+    #[cfg(target_os = "linux")]
+    if reactor_eligible_long_poll(&st) {
+        // Status is a placeholder: the engine intercepts `LongPollWait` before
+        // serialization and the reactor builds the real response head itself.
+        let mut resp = Resp::new(200);
+        resp.body = Body::LongPollWait(crate::api::LongPollReg {
+            st,
+            from,
+            client_cursor,
+            keep_alive: false,
+        });
+        return resp;
+    }
+
     // Wait for new data / closure / timeout.
     let mut rx = st.tail_tx.subscribe();
     let deadline = Instant::now() + long_poll_timeout_dur();
@@ -1806,8 +1826,22 @@ async fn long_poll_data(
     hot: bool,
     cache_hit: &mut bool,
 ) -> Resp {
-    let cursor = compute_cursor(client_cursor);
     let body = read_range_body(st, from, t.bytes, hot, "long-poll", cache_hit).await;
+    long_poll_data_resp(st, from, t, client_cursor, body)
+}
+
+/// Build the `200` response for a long-poll that delivered `from..t.bytes`, given
+/// an already-materialized body. Sync and free of any read, so it is shared by
+/// the inline path (async body via `read_range_body`) and the reactor (sync body
+/// materialized by hand) — both must emit byte-identical headers.
+pub(crate) fn long_poll_data_resp(
+    st: &StreamState,
+    from: u64,
+    t: Tail,
+    client_cursor: Option<u64>,
+    body: Body,
+) -> Resp {
+    let cursor = compute_cursor(client_cursor);
     let mut b = ResponseBuilder::new(200)
         .h("content-type", st.config.content_type.clone())
         .h(H_NEXT_OFFSET, format_offset(t.bytes))
@@ -1821,7 +1855,26 @@ async fn long_poll_data(
     b.body(body)
 }
 
-fn long_poll_close(tail: u64, cursor: u64) -> Resp {
+/// Wrap the wire bytes of a delivered long-poll range (`from..tail`) as a
+/// fixed-length body, matching `read_range_body`'s framing: JSON streams strip
+/// the trailing record comma and wrap the records in `[ ]`; others pass through.
+/// Used by the reactor, which materializes the bytes itself rather than producing
+/// a `Body::FileRange`.
+#[cfg(target_os = "linux")]
+pub(crate) fn long_poll_full_body(is_json: bool, data: Bytes) -> Body {
+    if is_json {
+        let inner = &data[..data.len().saturating_sub(1)];
+        let mut out = BytesMut::with_capacity(inner.len() + 2);
+        out.put_u8(b'[');
+        out.put_slice(inner);
+        out.put_u8(b']');
+        Body::Full(out.freeze())
+    } else {
+        Body::Full(data)
+    }
+}
+
+pub(crate) fn long_poll_close(tail: u64, cursor: u64) -> Resp {
     ResponseBuilder::new(204)
         .h(H_NEXT_OFFSET, format_offset(tail))
         .h(H_CURSOR, cursor.to_string())
@@ -1831,7 +1884,7 @@ fn long_poll_close(tail: u64, cursor: u64) -> Resp {
         .body(empty())
 }
 
-fn long_poll_timeout(tail: u64, cursor: u64, closed: bool) -> Resp {
+pub(crate) fn long_poll_timeout(tail: u64, cursor: u64, closed: bool) -> Resp {
     let mut b = ResponseBuilder::new(204)
         .h(H_NEXT_OFFSET, format_offset(tail))
         .h(H_CURSOR, cursor.to_string())
@@ -1843,6 +1896,18 @@ fn long_poll_timeout(tail: u64, cursor: u64, closed: bool) -> Resp {
     b.body(empty())
 }
 
+/// Whether a caught-up long-poll on `st` can wait in the epoll reactor. Mirrors
+/// the SSE live-tail gate (`SseSource::reactor_reg`): the reactor serves the live
+/// data file by `pread`, so a forked (`parent`) or tiered (`blobstore`) stream
+/// stays on the inline path. No `file_base` check is needed — a caught-up poll's
+/// `from == tail >= file_base`, and the bytes it waits for are appended to the
+/// live file (a rare concurrent compaction past `from` is handled at produce time
+/// by ending the connection so the client retries through the cold path).
+#[cfg(target_os = "linux")]
+fn reactor_eligible_long_poll(st: &StreamState) -> bool {
+    st.parent.is_none() && st.blobstore.is_none()
+}
+
 // ---------- SSE ----------
 
 #[derive(Clone, Copy)]

packages/durable-streams-rust/src/main.rs

@@ -4,7 +4,7 @@ mod engine_raw;
 mod handlers;
 mod http1;
 #[cfg(target_os = "linux")]
-mod sse_reactor;
+mod reactor;
 mod store;
 mod telemetry;
 mod tier;
@@ -353,7 +353,7 @@ fn main() {
                 // Close reactor-served SSE subscribers first so their permits are
                 // released and `drain` doesn't wait out the full grace period.
                 #[cfg(target_os = "linux")]
-                sse_reactor::shutdown();
+                reactor::shutdown();
                 engine_raw::drain(std::time::Duration::from_secs(25)).await;
                 telemetry_guard.shutdown();
             }