feat(sse): Conn::sse() sugar — EventSender + pump heartbeats (v0.3 chunk 3)
New src/sse.rs. Conn::sse() (and sse_with_heartbeat) turns the response
into a text/event-stream: status 200 if unset, content-type +
cache-control: no-cache, and a RespBody::Stream whose heartbeat is wired
to the new StreamBody.heartbeat field. Returns (Conn, EventSender);
EventSender is Clone (stream ends when the LAST clone drops) with
send(event, data) / data(data) / comment(text), all returning
Err(SseClosed) once the conn actor has dropped the stream — the producer
exit signal. Multi-line data becomes one data: line per line (pure
format_event, unit-tested).
Heartbeat lives in the pump, as the roadmap leaned: with
StreamBody.heartbeat set, pump_stream waits recv_timeout(interval) and on
expiry writes the payload (": keep-alive" comment chunk) and keeps
waiting. Liveness inversion is deliberate: no request clock governs an
SSE response (request_timeout is read-phase only since chunk 1); a dead
client is detected when an event or heartbeat write stalls past
write_timeout, and a draining shutdown force-stops the recv park like any
stream.
Tests: 4 unit (3 framing, 1 sse() wiring) + 2 integration (event round
trip over decoded chunked stream incl. named + unnamed events and clean
0-chunk EOS on sender drop; >=2 heartbeats observed across 450ms of
producer silence at a 100ms interval).
This commit is contained in:
Claude
+13
-1
@@ -174,11 +174,23 @@ pub enum RespBody {
|
||||
/// body.
|
||||
pub struct StreamBody {
|
||||
pub rx: smarm::Receiver<Vec<u8>>,
|
||||
/// 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<u8>)>,
|
||||
}
|
||||
|
||||
impl StreamBody {
|
||||
pub fn new(rx: smarm::Receiver<Vec<u8>>) -> Self {
|
||||
Self { rx }
|
||||
Self { rx, heartbeat: None }
|
||||
}
|
||||
|
||||
pub fn with_heartbeat(mut self, interval: std::time::Duration, payload: Vec<u8>) -> Self {
|
||||
self.heartbeat = Some((interval, payload));
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+34
-13
@@ -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)?;
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
+178
@@ -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<Vec<u8>>,
|
||||
}
|
||||
|
||||
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<u8> {
|
||||
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::<Vec<u8>>();
|
||||
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:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user