diff --git a/src/conn.rs b/src/conn.rs index 2658c42..3bedc32 100644 --- a/src/conn.rs +++ b/src/conn.rs @@ -174,11 +174,23 @@ pub enum RespBody { /// body. pub struct StreamBody { pub rx: smarm::Receiver>, + /// If set, the connection actor's pump waits for each chunk with + /// `recv_timeout(interval)` instead of a bare `recv`, and on expiry + /// writes `payload` as its own chunk and keeps waiting. This is how + /// SSE keep-alive comments work: liveness comes from the ping (a dead + /// client is detected when a heartbeat write stalls past + /// `write_timeout`), not from any request clock. + pub heartbeat: Option<(std::time::Duration, Vec)>, } impl StreamBody { pub fn new(rx: smarm::Receiver>) -> Self { - Self { rx } + Self { rx, heartbeat: None } + } + + pub fn with_heartbeat(mut self, interval: std::time::Duration, payload: Vec) -> Self { + self.heartbeat = Some((interval, payload)); + self } } diff --git a/src/conn_actor.rs b/src/conn_actor.rs index 36ea5f1..cb521d0 100644 --- a/src/conn_actor.rs +++ b/src/conn_actor.rs @@ -657,24 +657,45 @@ fn pump_stream( chunked: bool, write_timeout: Duration, ) -> io::Result<()> { + let write_chunk = |payload: &[u8]| -> io::Result<()> { + let deadline = Instant::now() + write_timeout; + if chunked { + let mut framed = Vec::with_capacity(payload.len() + 20); + framed.extend_from_slice(format!("{:x}\r\n", payload.len()).as_bytes()); + framed.extend_from_slice(payload); + framed.extend_from_slice(b"\r\n"); + write_all(fd, &framed, deadline) + } else { + write_all(fd, payload, deadline) + } + }; + loop { - match stream.rx.recv() { - Ok(chunk) => { + // With a heartbeat configured (SSE), the wait between chunks is the + // heartbeat interval; expiry emits the ping and keeps waiting. A + // ping write that stalls past write_timeout errors out below — + // that is the dead-client detector. Without one, plain recv(): + // stoppable by the draining registry either way. + let msg = match &stream.heartbeat { + Some((interval, payload)) => match stream.rx.recv_timeout(*interval) { + Ok(chunk) => Some(chunk), + Err(smarm::RecvTimeoutError::Timeout) => { + write_chunk(payload)?; + continue; + } + Err(smarm::RecvTimeoutError::Disconnected) => None, + }, + None => stream.rx.recv().ok(), + }; + + match msg { + Some(chunk) => { if chunk.is_empty() { continue; } - let deadline = Instant::now() + write_timeout; - if chunked { - let mut framed = Vec::with_capacity(chunk.len() + 20); - framed.extend_from_slice(format!("{:x}\r\n", chunk.len()).as_bytes()); - framed.extend_from_slice(&chunk); - framed.extend_from_slice(b"\r\n"); - write_all(fd, &framed, deadline)?; - } else { - write_all(fd, &chunk, deadline)?; - } + write_chunk(&chunk)?; } - Err(_) => { + None => { // All senders dropped: end of stream. if chunked { write_all(fd, b"0\r\n\r\n", Instant::now() + write_timeout)?; diff --git a/src/lib.rs b/src/lib.rs index e85d03e..17acbc0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -26,11 +26,13 @@ pub mod net; pub mod conn_actor; pub mod conn_registry; pub mod serve; +pub mod sse; // Re-exports — what most users want at the crate root. pub use conn::{Assigns, Body, Conn, HeaderMap, HttpVersion, Method, Params, RespBody, StreamBody}; pub use plug::{Next, Pipeline, Plug}; pub use router::Router; +pub use sse::{EventSender, SseClosed}; pub use serve::{ serve, serve_with, serve_with_shutdown, shutdown_handle, Config, Handle, ShutdownSignal, }; diff --git a/src/sse.rs b/src/sse.rs new file mode 100644 index 0000000..ea599c2 --- /dev/null +++ b/src/sse.rs @@ -0,0 +1,178 @@ +//! Server-Sent Events sugar over `RespBody::Stream`. +//! +//! An SSE response is just a streaming body with the right headers and a +//! heartbeat: `Conn::sse()` wires all of it and hands back an +//! [`EventSender`] for a producer actor to feed. +//! +//! ```no_run +//! use urus::{Conn, Next, Pipeline, Router}; +//! use std::time::Duration; +//! +//! let pipeline = Pipeline::new().plug(Router::new().get( +//! "/events", +//! |c: Conn, _n: Next| { +//! let (c, events) = c.sse(); +//! smarm::spawn(move || { +//! let mut i = 0; +//! loop { +//! if events.send("tick", &i.to_string()).is_err() { +//! return; // client gone; conn actor dropped the stream +//! } +//! i += 1; +//! smarm::sleep(Duration::from_secs(1)); +//! } +//! }); +//! c +//! }, +//! )); +//! ``` +//! +//! Liveness: the connection actor emits a `: keep-alive` comment chunk +//! whenever [`HEARTBEAT_INTERVAL`] passes with no event (see +//! `StreamBody::heartbeat`). A client that went away is detected when a +//! write — event or heartbeat — stalls past `write_timeout`; the conn +//! actor then drops the stream and the producer's next [`EventSender`] +//! call returns `Err(SseClosed)`. There is no request clock on an SSE +//! response: `request_timeout` covers only the read phase, by design. + +use crate::conn::{Conn, RespBody, StreamBody}; + +use std::time::Duration; + +/// Default heartbeat: one comment line per this interval of silence. +pub const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15); + +/// What the heartbeat puts on the wire — an SSE comment, invisible to +/// `EventSource` consumers. +pub(crate) const HEARTBEAT_PAYLOAD: &[u8] = b": keep-alive\n\n"; + +/// The stream has ended: the connection actor dropped the receiving end +/// (client disconnect, write timeout, or server shutdown). Producers +/// should exit when they see this. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SseClosed; + +impl std::fmt::Display for SseClosed { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("SSE stream closed") + } +} +impl std::error::Error for SseClosed {} + +/// Producer handle for an SSE response. Clonable — multiple producers may +/// feed one stream; the stream ends when the LAST clone drops. +#[derive(Clone)] +pub struct EventSender { + tx: smarm::Sender>, +} + +impl EventSender { + /// Send a named event: `event: {event}` + `data:` line(s). + pub fn send(&self, event: &str, data: &str) -> Result<(), SseClosed> { + self.tx + .send(format_event(Some(event), data)) + .map_err(|_| SseClosed) + } + + /// Send an unnamed event (`data:` line(s) only) — what `EventSource` + /// surfaces as a plain `message`. + pub fn data(&self, data: &str) -> Result<(), SseClosed> { + self.tx + .send(format_event(None, data)) + .map_err(|_| SseClosed) + } + + /// Send a comment line. Invisible to `EventSource`; useful for + /// application-level pings beyond the built-in heartbeat. + pub fn comment(&self, text: &str) -> Result<(), SseClosed> { + let mut out = Vec::with_capacity(text.len() + 4); + out.extend_from_slice(b": "); + out.extend_from_slice(text.as_bytes()); + out.extend_from_slice(b"\n\n"); + self.tx.send(out).map_err(|_| SseClosed) + } +} + +/// Wire format for one event. Multi-line `data` becomes one `data:` line +/// per line (the SSE framing for embedded newlines); the blank line +/// dispatches the event. +pub(crate) fn format_event(event: Option<&str>, data: &str) -> Vec { + let mut out = Vec::with_capacity(data.len() + 32); + if let Some(e) = event { + out.extend_from_slice(b"event: "); + out.extend_from_slice(e.as_bytes()); + out.push(b'\n'); + } + for line in data.split('\n') { + out.extend_from_slice(b"data: "); + out.extend_from_slice(line.as_bytes()); + out.push(b'\n'); + } + out.push(b'\n'); + out +} + +impl Conn { + /// Turn this response into a Server-Sent Events stream with the + /// default [`HEARTBEAT_INTERVAL`]. Returns the `Conn` (return it from + /// the handler) and the [`EventSender`] to hand to a producer actor. + pub fn sse(self) -> (Self, EventSender) { + self.sse_with_heartbeat(HEARTBEAT_INTERVAL) + } + + /// [`Conn::sse`] with a custom heartbeat interval. + pub fn sse_with_heartbeat(mut self, interval: Duration) -> (Self, EventSender) { + let (tx, rx) = smarm::channel::>(); + if self.status.is_none() { + self.status = Some(200); + } + self.resp_headers.set("content-type", "text/event-stream"); + self.resp_headers.set("cache-control", "no-cache"); + self.resp_body = RespBody::Stream( + StreamBody::new(rx).with_heartbeat(interval, HEARTBEAT_PAYLOAD.to_vec()), + ); + (self, EventSender { tx }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn format_named_event() { + assert_eq!( + format_event(Some("tick"), "42"), + b"event: tick\ndata: 42\n\n".to_vec() + ); + } + + #[test] + fn format_unnamed_event() { + assert_eq!(format_event(None, "hi"), b"data: hi\n\n".to_vec()); + } + + #[test] + fn format_multiline_data() { + assert_eq!( + format_event(None, "a\nb"), + b"data: a\ndata: b\n\n".to_vec() + ); + } + + #[test] + fn sse_sets_headers_and_stream_body() { + let (conn, _events) = Conn::new().sse(); + assert_eq!(conn.status, Some(200)); + assert_eq!(conn.resp_headers.get("content-type"), Some("text/event-stream")); + assert_eq!(conn.resp_headers.get("cache-control"), Some("no-cache")); + match &conn.resp_body { + RespBody::Stream(s) => { + let (interval, payload) = s.heartbeat.as_ref().expect("heartbeat set"); + assert_eq!(*interval, HEARTBEAT_INTERVAL); + assert_eq!(payload.as_slice(), HEARTBEAT_PAYLOAD); + } + other => panic!("expected Stream body, got {other:?}"), + } + } +} diff --git a/tests/integration.rs b/tests/integration.rs index 55c8995..0f6b2df 100644 --- a/tests/integration.rs +++ b/tests/integration.rs @@ -905,3 +905,65 @@ fn chunked_request_stall_killed_at_request_timeout() { assert!(resp.is_empty(), "expected silent close, got: {:?}", String::from_utf8_lossy(&resp)); assert!(start.elapsed() < Duration::from_secs(3), "close took too long"); } + +// --------------------------------------------------------------------------- +// SSE (v0.3 chunk 3) +// --------------------------------------------------------------------------- + +/// Conn::sse(): the response is a chunked text/event-stream; events sent +/// through the EventSender arrive framed; the stream ends (0-chunk) when +/// the producer drops the sender. +#[test] +fn sse_events_round_trip() { + let pipe = Pipeline::new().plug(Router::new().get( + "/events", + |c: Conn, _n: Next| { + let (c, events) = c.sse(); + smarm::spawn(move || { + let _ = events.send("tick", "1"); + let _ = events.data("plain"); + smarm::sleep(Duration::from_millis(10)); + // events drops here: end of stream. + }); + c + }, + )); + let port = spawn_server(pipe); + let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap(); + s.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); + s.write_all(b"GET /events HTTP/1.1\r\nHost: x\r\n\r\n").unwrap(); + + let (head, body) = read_chunked_response(&mut s); + assert!(head.contains("content-type: text/event-stream"), "head: {head}"); + assert!(head.contains("transfer-encoding: chunked"), "head: {head}"); + let body = String::from_utf8(body).unwrap(); + assert!(body.contains("event: tick\ndata: 1\n\n"), "body: {body}"); + assert!(body.contains("data: plain\n\n"), "body: {body}"); +} + +/// With a silent producer, the conn actor emits keep-alive comment chunks +/// at the configured heartbeat interval. +#[test] +fn sse_heartbeats_on_silence() { + let pipe = Pipeline::new().plug(Router::new().get( + "/events", + |c: Conn, _n: Next| { + let (c, events) = c.sse_with_heartbeat(Duration::from_millis(100)); + smarm::spawn(move || { + // Hold the sender open, silently, then end the stream. + smarm::sleep(Duration::from_millis(450)); + drop(events); + }); + c + }, + )); + let port = spawn_server(pipe); + let mut s = TcpStream::connect(("127.0.0.1", port)).unwrap(); + s.set_read_timeout(Some(Duration::from_secs(5))).unwrap(); + s.write_all(b"GET /events HTTP/1.1\r\nHost: x\r\n\r\n").unwrap(); + + let (_head, body) = read_chunked_response(&mut s); + let body = String::from_utf8(body).unwrap(); + let beats = body.matches(": keep-alive").count(); + assert!(beats >= 2, "expected >=2 heartbeats in 450ms at 100ms interval, got {beats}: {body:?}"); +}