trillium_sse/
handler.rs

1use crate::{Eventable, encode};
2use futures_lite::{AsyncWriteExt, Stream};
3use std::{
4    fmt::{self, Debug, Formatter},
5    future::{self, Future},
6    io,
7    pin::Pin,
8    task::Poll,
9    time::Duration,
10};
11use sync_wrapper::SyncWrapper;
12use trillium::{Conn, Handler, Info, KnownHeaderName, Status, Upgrade};
13use trillium_server_common::Runtime;
14
15/// The trait that defines an event source for the [`Sse`] handler.
16///
17/// Implement this on a type that holds whatever fanout mechanism your application uses — a
18/// broadcast channel, a subscription registry — and return a per-client [`Stream`] from
19/// [`connect`](SseHandler::connect).
20///
21/// ```
22/// use broadcaster::BroadcastChannel;
23/// use trillium::Conn;
24/// use trillium_sse::{Sse, SseHandler};
25///
26/// struct Notifications {
27///     channel: BroadcastChannel<String>,
28/// }
29///
30/// impl SseHandler for Notifications {
31///     type Event = String;
32///     type EventStream = BroadcastChannel<String>;
33///
34///     async fn connect(&self, _conn: &mut Conn) -> Self::EventStream {
35///         self.channel.clone()
36///     }
37/// }
38///
39/// let handler = Sse::new(Notifications {
40///     channel: BroadcastChannel::new(),
41/// });
42/// ```
43///
44/// This trait is also implemented for any `Fn(&mut Conn) -> Stream`, for the common case
45/// where nothing needs to be awaited in order to build the stream:
46///
47/// ```
48/// use futures_lite::stream;
49/// use trillium::Conn;
50/// use trillium_sse::sse;
51///
52/// let handler = sse(|_: &mut Conn| stream::iter(["one", "two"]));
53/// ```
54pub trait SseHandler: Send + Sync + Sized + 'static {
55    /// The type yielded by this handler's [`EventStream`](SseHandler::EventStream).
56    type Event: Eventable;
57
58    /// A [`Stream`] of events to send to a connected client, built per client in
59    /// [`connect`](SseHandler::connect).
60    type EventStream: Stream<Item = Self::Event> + Unpin + Send + 'static;
61
62    /// Called once per request, to build the stream of events for that client.
63    ///
64    /// The conn is borrowed mutably to allow setting response headers or state, but note that
65    /// [`Sse`] sets the status, headers, and body itself.
66    fn connect(&self, conn: &mut Conn) -> impl Future<Output = Self::EventStream> + Send;
67}
68
69impl<F, S, E> SseHandler for F
70where
71    F: Fn(&mut Conn) -> S + Send + Sync + 'static,
72    S: Stream<Item = E> + Unpin + Send + 'static,
73    E: Eventable,
74{
75    type Event = E;
76    type EventStream = S;
77
78    async fn connect(&self, conn: &mut Conn) -> Self::EventStream {
79        self(conn)
80    }
81}
82
83/// A [`Handler`] that responds to requests with a server-sent event stream.
84///
85/// Build one from any [`SseHandler`] with [`Sse::new`] or [`sse`].
86///
87/// The conn is passed through untouched — continuing on to subsequent handlers — unless the
88/// request's [`Accept`][rfc] header names `text/event-stream`. Wildcard ranges and an absent
89/// `Accept` header do not match, so the same route can also serve other representations.
90///
91/// When the client disconnects, the event stream is dropped: promptly on HTTP/1.x and HTTP/2,
92/// and at the next event or heartbeat on HTTP/3.
93///
94/// [rfc]: https://www.rfc-editor.org/rfc/rfc9110.html#name-accept
95pub struct Sse<H> {
96    handler: H,
97    heartbeat: Option<Duration>,
98    runtime: Option<Runtime>,
99}
100
101/// Builds a new [`Sse`] handler. Alias for [`Sse::new`].
102pub fn sse<H: SseHandler>(sse_handler: H) -> Sse<H> {
103    Sse::new(sse_handler)
104}
105
106impl<H: SseHandler> Sse<H> {
107    /// Builds a new [`Sse`] handler from any [`SseHandler`].
108    pub fn new(handler: H) -> Self {
109        Self {
110            handler,
111            heartbeat: None,
112            runtime: None,
113        }
114    }
115
116    /// Sends an empty comment whenever `heartbeat` elapses without the stream yielding an event.
117    ///
118    /// The interval is measured from the most recent event, not from the start of the response,
119    /// so a busy stream sends no heartbeats at all. Clients discard the comment.
120    ///
121    /// Regular traffic makes an idle stream less likely to be dropped by an intermediary. On
122    /// HTTP/3 it also bounds how long a departed client goes unnoticed, since disconnection is
123    /// detected there by a failed write.
124    pub fn with_heartbeat(mut self, heartbeat: Duration) -> Self {
125        self.heartbeat = Some(heartbeat);
126        self
127    }
128}
129
130/// Whether the client asked for `text/event-stream` by name.
131///
132/// Wildcards (`*/*`, `text/*`) and an absent `Accept` header do not count: they say the client
133/// will take whatever is on offer, not that it speaks the event-stream protocol. Requiring the
134/// media type by name lets an `Sse` handler share a route with handlers serving other
135/// representations. An `Accept` header naming it with `q=0` is a refusal.
136fn accepts_event_stream(conn: &Conn) -> bool {
137    let Some(accept) = conn.request_headers().get_str(KnownHeaderName::Accept) else {
138        return false;
139    };
140
141    accept.split(',').any(|media_range| {
142        let mut parts = media_range.split(';').map(str::trim);
143
144        let matches_range = parts
145            .next()
146            .is_some_and(|range| range.eq_ignore_ascii_case("text/event-stream"));
147
148        matches_range
149            && !parts.any(|parameter| {
150                parameter
151                    .strip_prefix("q=")
152                    .and_then(|q| q.parse::<f32>().ok())
153                    .is_some_and(|q| q <= 0.0)
154            })
155    })
156}
157
158/// Private state key carrying the event stream from [`Handler::run`] to [`Handler::upgrade`].
159/// The [`SyncWrapper`] lets a `!Sync` stream satisfy the state
160/// [`TypeSet`](trillium::TypeSet)'s `Sync` requirement.
161struct SseStream<S>(SyncWrapper<S>);
162
163/// Stray inbound bytes tolerated while probing for disconnection on h1. A conforming client
164/// sends nothing on an event stream, but a pipelining client may have optimistic requests in
165/// flight.
166const READ_ALLOWANCE: usize = 16 * 1024;
167
168enum Tick<E> {
169    Event(E),
170    Heartbeat,
171    StreamEnded,
172    ClientDisconnected,
173}
174
175async fn write_flush(upgrade: &mut Upgrade, bytes: &[u8]) -> io::Result<()> {
176    upgrade.write_all(bytes).await?;
177    upgrade.flush().await
178}
179
180async fn drive_events<S, E>(mut upgrade: Upgrade, stream: S, heartbeat: Option<(Duration, Runtime)>)
181where
182    S: Stream<Item = E> + Unpin + Send + 'static,
183    E: Eventable,
184{
185    let swansong = upgrade.swansong();
186    let mut stream = swansong.interrupt(stream);
187
188    let new_delay = |(duration, runtime): &(Duration, Runtime)| {
189        let duration = *duration;
190        let runtime = runtime.clone();
191        Box::pin(async move { runtime.delay(duration).await })
192            as Pin<Box<dyn Future<Output = ()> + Send>>
193    };
194    let mut delay = heartbeat.as_ref().map(new_delay);
195
196    loop {
197        let tick = future::poll_fn(|cx| {
198            if Pin::new(upgrade.as_mut())
199                .poll_closed(cx, READ_ALLOWANCE)
200                .is_ready()
201            {
202                return Poll::Ready(Tick::ClientDisconnected);
203            }
204
205            match Pin::new(&mut stream).poll_next(cx) {
206                Poll::Ready(Some(event)) => return Poll::Ready(Tick::Event(event)),
207                Poll::Ready(None) => return Poll::Ready(Tick::StreamEnded),
208                Poll::Pending => {}
209            }
210
211            if let Some(delay) = &mut delay
212                && delay.as_mut().poll(cx).is_ready()
213            {
214                return Poll::Ready(Tick::Heartbeat);
215            }
216
217            Poll::Pending
218        })
219        .await;
220
221        match tick {
222            Tick::ClientDisconnected => return,
223            Tick::StreamEnded => break,
224            Tick::Event(event) => {
225                delay = heartbeat.as_ref().map(new_delay);
226                let Some(encoded) = encode(&event) else {
227                    continue;
228                };
229                if write_flush(&mut upgrade, encoded.as_bytes()).await.is_err() {
230                    return;
231                }
232            }
233            Tick::Heartbeat => {
234                delay = heartbeat.as_ref().map(new_delay);
235                if write_flush(&mut upgrade, b":\n\n").await.is_err() {
236                    return;
237                }
238            }
239        }
240    }
241
242    let _ = upgrade.close().await;
243}
244
245impl<H: SseHandler> Handler for Sse<H> {
246    async fn run(&self, mut conn: Conn) -> Conn {
247        if !accepts_event_stream(&conn) {
248            return conn;
249        }
250
251        let stream = self.handler.connect(&mut conn).await;
252
253        conn.with_state(SseStream(SyncWrapper::new(stream)))
254            .with_response_header(KnownHeaderName::ContentType, "text/event-stream")
255            .with_response_header(KnownHeaderName::CacheControl, "no-cache")
256            // Close-delimited framing: the event stream carries neither `Content-Length`
257            // nor `Transfer-Encoding`, running until the connection closes. Chunked
258            // transfer-encoding can disrupt event delivery timing for this protocol.
259            // h2 and h3 strip this h1-only header and frame at the stream layer.
260            .with_response_header(KnownHeaderName::Connection, "close")
261            .with_status(Status::Ok)
262            .halt()
263            .upgrade()
264    }
265
266    async fn init(&mut self, info: &mut Info) {
267        self.runtime = info.shared_state::<Runtime>().cloned();
268
269        if self.heartbeat.is_some() && self.runtime.is_none() {
270            log::warn!(
271                "no runtime in shared state; sse heartbeats are disabled. this handler was \
272                 probably not initialized by a trillium runtime adapter."
273            );
274        }
275    }
276
277    fn has_upgrade(&self, upgrade: &Upgrade) -> bool {
278        upgrade.state().contains::<SseStream<H::EventStream>>()
279    }
280
281    async fn upgrade(&self, mut upgrade: Upgrade) {
282        let Some(SseStream(stream)) = upgrade.state_mut().take::<SseStream<H::EventStream>>()
283        else {
284            return;
285        };
286        let stream = stream.into_inner();
287
288        let heartbeat = self.heartbeat.zip(self.runtime.clone());
289        drive_events(upgrade, stream, heartbeat).await;
290    }
291}
292
293impl<H> Debug for Sse<H>
294where
295    H: Debug,
296{
297    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
298        f.debug_struct("Sse")
299            .field("handler", &self.handler)
300            .field("heartbeat", &self.heartbeat)
301            .field("runtime", &self.runtime)
302            .finish()
303    }
304}
trillium_sse/handler.rs

trillium_sse/
handler.rs
