nordisk/knut

crates/knut-wasm-host/src/sse.rs

//! Server-sent events: a stream of already-encoded frame strings → a
//! `text/event-stream` body.
//!
//! What goes in is a `String` the caller made. This crate does not encode it, does
//! not parse it, and does not know that `{"done":true}` means the run finished —
//! only that the caller sent it last. The *terminal* frames are the caller's too:
//! [`FrameSender::finish`] exists to say "this one is the last", not to invent one.
//!
//! The framing itself is the only thing owned here, and it is the part that is easy
//! to get subtly wrong: SSE ends an event at a blank line, so a payload containing
//! a newline must carry a `data:` prefix on EVERY line or one frame silently
//! becomes two.
//!
//! The body ends when every [`FrameSender`] is dropped. No runtime is started here;
//! the caller's does the polling.

use std::convert::Infallible;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;

use bytes::{BufMut, Bytes, BytesMut};
use hyper::body::{Body, Frame, SizeHint};
use hyper::http::{Response, header};
use tokio::sync::mpsc;

/// The content type an SSE body must be served with.
pub const SSE_CONTENT_TYPE: &str = "text/event-stream";

/// A comment line — bytes that keep an idle connection (and any proxy in front of
/// it) alive without being an event. See [`SseBody::with_keep_alive`].
const KEEP_ALIVE_COMMENT: &[u8] = b":\n\n";

/// Frame `payload` as ONE server-sent event: a `data:` line per line of the
/// payload, then the blank line that ends the event.
#[must_use]
pub fn event_bytes(payload: &str) -> Bytes {
    // `data: ` + payload + a newline per line + the terminating blank line.
    let mut out = BytesMut::with_capacity(payload.len() + 16);
    for line in payload.split('\n') {
        out.put_slice(b"data: ");
        // A trailing `\r` would ride into the payload the browser decodes.
        out.put_slice(line.strip_suffix('\r').unwrap_or(line).as_bytes());
        out.put_u8(b'\n');
    }
    out.put_u8(b'\n');
    out.freeze()
}

/// Open a frame channel: a sender for the caller's producer task, and the body to
/// hand back as the response.
#[must_use]
pub fn sse_channel() -> (FrameSender, SseBody) {
    let (tx, rx) = mpsc::unbounded_channel::<String>();
    (
        FrameSender { tx },
        SseBody {
            rx,
            keep_alive: None,
        },
    )
}

/// The producing half of [`sse_channel`]. Cloneable, and the body lives until the
/// last clone is dropped.
#[derive(Clone, Debug)]
pub struct FrameSender {
    tx: mpsc::UnboundedSender<String>,
}

impl FrameSender {
    /// Queue one already-encoded frame. `false` means the client is gone — a
    /// producer that cares can stop; one that does not can ignore it, which is why
    /// this returns a bool instead of an error nobody handles.
    pub fn frame(&self, encoded: impl Into<String>) -> bool {
        self.tx.send(encoded.into()).is_ok()
    }

    /// Queue the caller's TERMINAL frame — whatever string the caller's own
    /// protocol ends a stream with. Identical machinery to [`FrameSender::frame`];
    /// separate so the intent is in the call site rather than in a comment.
    pub fn finish(&self, terminal: impl Into<String>) -> bool {
        self.frame(terminal)
    }

    /// Whether the receiving body has been dropped (the browser went away).
    #[must_use]
    pub fn is_closed(&self) -> bool {
        self.tx.is_closed()
    }
}

/// The consuming half: an [`http_body::Body`] of `text/event-stream` bytes.
///
/// [`http_body::Body`]: hyper::body::Body
#[derive(Debug)]
pub struct SseBody {
    rx: mpsc::UnboundedReceiver<String>,
    keep_alive: Option<tokio::time::Interval>,
}

impl SseBody {
    /// Emit a comment line every `every` while no frame is ready, so an idle stream
    /// (and any proxy between) does not time the connection out. Needs a tokio
    /// timer on the polling runtime; without this the body is pure `sync`.
    #[must_use]
    pub fn with_keep_alive(mut self, every: Duration) -> Self {
        let mut interval = tokio::time::interval(every);
        // The first tick of a tokio interval fires immediately; a keep-alive
        // comment before the first frame would be noise.
        interval.reset();
        self.keep_alive = Some(interval);
        self
    }

    /// The response this body belongs in: `200` + [`SSE_CONTENT_TYPE`] + no caching.
    #[must_use]
    pub fn into_response(self) -> Response<Self> {
        Response::builder()
            .header(header::CONTENT_TYPE, SSE_CONTENT_TYPE)
            .header(header::CACHE_CONTROL, "no-store")
            .body(self)
            .expect("a content type always builds a response")
    }
}

impl Body for SseBody {
    type Data = Bytes;
    type Error = Infallible;

    fn poll_frame(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
        let this = self.get_mut();
        match this.rx.poll_recv(cx) {
            Poll::Ready(Some(payload)) => Poll::Ready(Some(Ok(Frame::data(event_bytes(&payload))))),
            // Every sender dropped: the stream is over.
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Pending => {
                let ticked = this
                    .keep_alive
                    .as_mut()
                    .is_some_and(|interval| interval.poll_tick(cx).is_ready());
                if ticked {
                    Poll::Ready(Some(Ok(Frame::data(Bytes::from_static(KEEP_ALIVE_COMMENT)))))
                } else {
                    Poll::Pending
                }
            }
        }
    }

    fn is_end_stream(&self) -> bool {
        false
    }

    fn size_hint(&self) -> SizeHint {
        SizeHint::default()
    }
}