trillium_sse/
lib.rs

1//! # Trillium tools for server sent events
2//!
3//! There are two ways to use this crate.
4//!
5//! ## The [`Sse`] handler
6//!
7//! [`sse`] builds a [`Handler`](trillium::Handler) from any [`SseHandler`], which produces a
8//! [`Stream`] of [`Eventable`] items per connected client. This is the fuller-featured of the
9//! two, because a handler is initialized by the server and so can send heartbeats.
10//!
11//! Requests whose `Accept` header excludes `text/event-stream` are passed through to subsequent
12//! handlers untouched.
13//!
14//! ```
15//! use broadcaster::BroadcastChannel;
16//! use std::time::Duration;
17//! use trillium::Conn;
18//! use trillium_sse::sse;
19//!
20//! let channel = BroadcastChannel::<String>::new();
21//!
22//! let handler = sse(move |_: &mut Conn| channel.clone()).with_heartbeat(Duration::from_secs(15));
23//! ```
24//!
25//! ## [`SseConnExt`]
26//!
27//! [`SseConnExt`] is an extension trait for [`trillium::Conn`] whose
28//! [`with_sse_stream`](crate::SseConnExt::with_sse_stream) chainable method takes a [`Stream`]
29//! where the `Item` implements [`Eventable`]. Use it when you already have a conn in hand and
30//! want to respond with an event stream.
31//!
32//! ```
33//! use broadcaster::BroadcastChannel;
34//! use trillium::{Conn, conn_unwrap};
35//! use trillium_sse::SseConnExt;
36//!
37//! type Channel = BroadcastChannel<String>;
38//!
39//! fn get_sse(mut conn: Conn) -> Conn {
40//!     let broadcaster = conn_unwrap!(conn.take_state::<Channel>(), conn);
41//!     conn.with_sse_stream(broadcaster)
42//! }
43//! ```
44//!
45//! Often, you will want this stream to be something like a channel, but
46//! the specifics of that are dependent on the event fanout
47//! characteristics of your application.
48//!
49//! ## Events
50//!
51//! This crate implements [`Eventable`] for an [`Event`] type that you can
52//! use in your application, for `String`, and for `&'static str`. You can
53//! also implement [`Eventable`] for any type in your application.
54//!
55//! In addition to data events, the stream can carry comments — messages that clients ignore,
56//! sent periodically so that an idle stream is still producing traffic. See
57//! [`Event::new_comment`], or
58//! [`Sse::with_heartbeat`] /
59//! [`with_sse_stream_and_heartbeat`](crate::SseConnExt::with_sse_stream_and_heartbeat) to have them
60//! sent automatically whenever the stream goes quiet.
61#![forbid(unsafe_code)]
62#![deny(
63    missing_copy_implementations,
64    rustdoc::missing_crate_level_docs,
65    missing_debug_implementations,
66    nonstandard_style,
67    unused_qualifications
68)]
69#![warn(missing_docs)]
70
71#[cfg(test)]
72#[doc = include_str!("../README.md")]
73mod readme {}
74
75mod handler;
76mod heartbeat;
77
78use futures_lite::{AsyncRead, stream::Stream};
79pub use handler::{Sse, SseHandler, sse};
80use heartbeat::WithHeartbeat;
81use std::{
82    borrow::Cow,
83    fmt::Write,
84    io,
85    marker::PhantomData,
86    pin::Pin,
87    task::{Context, Poll},
88    time::Duration,
89};
90use trillium::{Body, Conn, KnownHeaderName, Status};
91use trillium_server_common::Runtime;
92
93struct SseBody<S, E> {
94    stream: S,
95    buffer: Vec<u8>,
96    event: PhantomData<E>,
97}
98
99impl<S, E> SseBody<S, E>
100where
101    S: Stream<Item = E> + Unpin + Send + 'static,
102    E: Eventable,
103{
104    pub fn new(stream: S) -> Self {
105        Self {
106            stream,
107            buffer: Vec::new(),
108            event: PhantomData,
109        }
110    }
111}
112
113fn write_multiline_field(output: &mut String, prefix: &str, value: &str) {
114    for line in value.split('\n') {
115        let line = line.strip_suffix('\r').unwrap_or(line);
116        if line.is_empty() {
117            writeln!(output, "{prefix}").unwrap();
118        } else {
119            writeln!(output, "{prefix} {line}").unwrap();
120        }
121    }
122}
123
124/// Returns `None` for an event with no fields set at all, which would otherwise be written as a
125/// bare message terminator.
126fn encode(event: &impl Eventable) -> Option<String> {
127    let mut output = String::new();
128
129    if let Some(comment) = event.comment() {
130        write_multiline_field(&mut output, ":", comment);
131    }
132
133    if let Some(event_type) = event.event_type() {
134        writeln!(&mut output, "event: {event_type}").ok()?;
135    }
136
137    if let Some(id) = event.id() {
138        writeln!(&mut output, "id: {id}").ok()?;
139    }
140
141    if let Some(retry) = event.retry() {
142        writeln!(&mut output, "retry: {}", retry.as_millis()).ok()?;
143    }
144
145    if let Some(data) = event.data() {
146        write_multiline_field(&mut output, "data:", data);
147    }
148
149    if output.is_empty() {
150        None
151    } else {
152        writeln!(&mut output).ok()?;
153        Some(output)
154    }
155}
156
157impl<S, E> AsyncRead for SseBody<S, E>
158where
159    S: Stream<Item = E> + Unpin + Send + 'static,
160    E: Eventable,
161{
162    fn poll_read(
163        self: Pin<&mut Self>,
164        cx: &mut Context<'_>,
165        buf: &mut [u8],
166    ) -> Poll<io::Result<usize>> {
167        let Self { buffer, stream, .. } = self.get_mut();
168
169        let buffer_read = buffer.len().min(buf.len());
170        if buffer_read > 0 {
171            buf[0..buffer_read].copy_from_slice(&buffer[0..buffer_read]);
172            buffer.drain(0..buffer_read);
173            return Poll::Ready(Ok(buffer_read));
174        }
175
176        loop {
177            break match Pin::new(&mut *stream).poll_next(cx) {
178                Poll::Pending => Poll::Pending,
179                Poll::Ready(Some(item)) => {
180                    let Some(data) = encode(&item) else { continue };
181                    let data = data.into_bytes();
182                    let writable_len = data.len().min(buf.len());
183                    buf[0..writable_len].copy_from_slice(&data[0..writable_len]);
184                    if writable_len < data.len() {
185                        buffer.extend_from_slice(&data[writable_len..]);
186                    }
187                    Poll::Ready(Ok(writable_len))
188                }
189
190                Poll::Ready(None) => Poll::Ready(Ok(0)),
191            };
192        }
193    }
194}
195
196impl<S, E> From<SseBody<S, E>> for Body
197where
198    S: Stream<Item = E> + Unpin + Send + 'static,
199    E: Eventable,
200{
201    fn from(sse_body: SseBody<S, E>) -> Self {
202        Body::new_streaming(sse_body, None)
203    }
204}
205
206/// Extension trait for server sent events
207pub trait SseConnExt {
208    /// builds and sets a streaming response body that conforms to the
209    /// [server-sent-events
210    /// spec](https://html.spec.whatwg.org/multipage/server-sent-events.html#server-sent-events)
211    /// from a Stream of any [`Eventable`] type (such as
212    /// [`Event`], as well as setting appropiate headers for
213    /// this response.
214    fn with_sse_stream<S, E>(self, sse_stream: S) -> Self
215    where
216        S: Stream<Item = E> + Unpin + Send + 'static,
217        E: Eventable;
218
219    /// as [`with_sse_stream`](SseConnExt::with_sse_stream), but sends an empty comment whenever
220    /// `heartbeat` elapses without the stream yielding an event.
221    ///
222    /// The interval is measured from the most recent event, not from the start of the response,
223    /// so a busy stream sends no heartbeats at all. Clients discard the comment.
224    fn with_sse_stream_and_heartbeat<S, E>(self, sse_stream: S, heartbeat: Duration) -> Self
225    where
226        S: Stream<Item = E> + Unpin + Send + 'static,
227        E: Eventable;
228}
229
230impl SseConnExt for Conn {
231    fn with_sse_stream<S, E>(self, sse_stream: S) -> Self
232    where
233        S: Stream<Item = E> + Unpin + Send + 'static,
234        E: Eventable,
235    {
236        let body = SseBody::new(self.swansong().interrupt(sse_stream));
237        self.set_sse_headers().with_body(body)
238    }
239
240    fn with_sse_stream_and_heartbeat<S, E>(self, sse_stream: S, heartbeat: Duration) -> Self
241    where
242        S: Stream<Item = E> + Unpin + Send + 'static,
243        E: Eventable,
244    {
245        let Some(runtime) = self.shared_state::<Runtime>().cloned() else {
246            log::warn!(
247                "no runtime in shared state; sending sse stream without a heartbeat. this conn \
248                 was probably not served by a trillium server."
249            );
250            return self.with_sse_stream(sse_stream);
251        };
252
253        let stream = WithHeartbeat::new(sse_stream, runtime, heartbeat);
254        let body = SseBody::new(self.swansong().interrupt(stream));
255        self.set_sse_headers().with_body(body)
256    }
257}
258
259trait SseHeaders {
260    fn set_sse_headers(self) -> Self;
261}
262
263impl SseHeaders for Conn {
264    fn set_sse_headers(self) -> Self {
265        self.with_response_header(KnownHeaderName::ContentType, "text/event-stream")
266            .with_response_header(KnownHeaderName::CacheControl, "no-cache")
267            // Close-delimited framing: the event stream carries neither `Content-Length`
268            // nor `Transfer-Encoding`, running until the connection closes. Chunked
269            // transfer-encoding can disrupt event delivery timing for this protocol.
270            .with_response_header(KnownHeaderName::Connection, "close")
271            .with_status(Status::Ok)
272            .halt()
273    }
274}
275
276/// A trait that allows any Unpin + Send + Sync type to act as an event.
277///
278/// For a concrete implementation of this trait, you can use [`Event`],
279/// but it is also implemented for [`String`] and [`&'static str`].
280pub trait Eventable: Unpin + Send + Sync + 'static {
281    /// return the data for this event, if any
282    ///
283    /// Returning `None` yields a message with no `data:` field. Clients dispatch no event for
284    /// such a message, so it is only useful in combination with [`comment`](Eventable::comment).
285    fn data(&self) -> Option<&str>;
286
287    /// return a comment to send alongside this event, optionally
288    ///
289    /// Comments are ignored by clients. They are chiefly used as a heartbeat, sent periodically
290    /// so that an idle event stream is still producing traffic.
291    fn comment(&self) -> Option<&str> {
292        None
293    }
294
295    /// return the event type, optionally
296    fn event_type(&self) -> Option<&str> {
297        None
298    }
299
300    /// return a unique event id, optionally
301    fn id(&self) -> Option<&str> {
302        None
303    }
304
305    /// return a reconnection time to request of the client, optionally
306    ///
307    /// Sent as a `retry:` field in milliseconds, truncated. Clients that reconnect on their own
308    /// — such as the browser `EventSource` — wait this long before doing so. Whether and how it
309    /// is honored is entirely up to the client.
310    fn retry(&self) -> Option<Duration> {
311        None
312    }
313}
314
315impl Eventable for Event {
316    fn data(&self) -> Option<&str> {
317        Event::data(self)
318    }
319
320    fn comment(&self) -> Option<&str> {
321        Event::comment(self)
322    }
323
324    fn event_type(&self) -> Option<&str> {
325        Event::event_type(self)
326    }
327
328    fn id(&self) -> Option<&str> {
329        Event::id(self)
330    }
331
332    fn retry(&self) -> Option<Duration> {
333        Event::retry(self)
334    }
335}
336
337impl Eventable for &'static str {
338    fn data(&self) -> Option<&str> {
339        Some(self)
340    }
341}
342
343impl Eventable for String {
344    fn data(&self) -> Option<&str> {
345        Some(self)
346    }
347}
348
349/// Events are a concrete implementation of the [`Eventable`] trait.
350#[derive(Debug, Clone, Eq, PartialEq, Default, fieldwork::Fieldwork)]
351#[fieldwork(get, set, get_mut, with, option_set_some, into)]
352pub struct Event {
353    /// the data for this event
354    data: Option<Cow<'static, str>>,
355    /// a comment for this event
356    comment: Option<Cow<'static, str>>,
357    /// the type for this event
358    #[field(with = with_type, set = set_type, get_mut = type_mut)]
359    event_type: Option<Cow<'static, str>>,
360    /// the id for this event
361    id: Option<Cow<'static, str>>,
362    /// reconnection time for this stream
363    #[field(copy, into = false)]
364    retry: Option<Duration>,
365}
366
367impl From<&'static str> for Event {
368    fn from(s: &'static str) -> Self {
369        Self::from(Cow::Borrowed(s))
370    }
371}
372
373impl From<String> for Event {
374    fn from(s: String) -> Self {
375        Self::from(Cow::Owned(s))
376    }
377}
378
379impl From<Cow<'static, str>> for Event {
380    fn from(data: Cow<'static, str>) -> Self {
381        Event {
382            data: Some(data),
383            ..Self::default()
384        }
385    }
386}
387
388impl Event {
389    /// builds a new [`Event`]
390    ///
391    /// by default, this event has no event type. to set an event type,
392    /// use [`Event::with_type`] or [`Event::set_type`]
393    pub fn new(data: impl Into<Cow<'static, str>>) -> Self {
394        Self::from(data.into())
395    }
396
397    /// builds a new comment-only [`Event`], with no data
398    ///
399    /// Clients ignore comments and dispatch no event for this message. Sending one periodically
400    /// keeps an otherwise idle event stream from being closed by intermediaries.
401    ///
402    /// ```
403    /// let event = trillium_sse::Event::new_comment("heartbeat");
404    /// assert_eq!(event.comment(), Some("heartbeat"));
405    /// assert_eq!(event.data(), None);
406    /// ```
407    pub fn new_comment(comment: impl Into<Cow<'static, str>>) -> Self {
408        Self {
409            comment: Some(comment.into()),
410            ..Self::default()
411        }
412    }
413}
trillium_sse/lib.rs

trillium_sse/
lib.rs
