timvw_rust_runtime/
sse.rs

1//! Server-sent events (`text/event-stream`): a parser, the events of a response, and their `data:` decoded as JSON.
2//!
3//! [`Response::events`] gives the events of a response as they arrive; [`Events::json`] decodes each `data:` payload
4//! to a type of your own. A `data: [DONE]` event or the end of the body ends the stream. [`Decoder`] is the parser
5//! underneath, for a body read some other way.
6//!
7//! ```no_run
8//! # async fn demo(client: timvw_rust_runtime::Client) -> Result<(), timvw_rust_runtime::Error> {
9//! use futures_util::StreamExt;
10//!
11//! #[derive(serde::Deserialize)]
12//! struct Chunk {
13//!     id: String,
14//! }
15//!
16//! let response = client
17//!     .request(timvw_rust_runtime::Method::POST, "https://api.example.com/v1/stream")
18//!     .json(&serde_json::json!({"stream": true}))
19//!     .send()
20//!     .await?;
21//! let mut chunks = response.events().json::<Chunk>();
22//! while let Some(chunk) = chunks.next().await {
23//!     println!("{}", chunk?.id);
24//! }
25//! # Ok(())
26//! # }
27//! ```
28
29use std::collections::VecDeque;
30use std::fmt;
31#[cfg(feature = "time")]
32use std::future::Future;
33use std::marker::PhantomData;
34use std::pin::Pin;
35use std::task::{Context, Poll};
36
37use futures_core::Stream;
38use serde::de::DeserializeOwned;
39
40use crate::{Error, MaybeSend, Response};
41
42/// The `data:` value that ends a stream.
43pub const DONE: &str = "[DONE]";
44
45/// One server-sent event, before its data is decoded.
46#[derive(Clone, Debug, Default, PartialEq, Eq)]
47pub struct Event {
48    /// The `event:` field, when the event has one.
49    pub event: Option<String>,
50    /// The `data:` lines of the event, joined with `\n`.
51    pub data: String,
52    /// The last `id:` field seen on the stream, when there was one.
53    pub id: Option<String>,
54}
55
56/// An incremental `text/event-stream` parser, following the HTML standard: lines end with CR, LF or CRLF, a line
57/// starting with `:` is a comment, an empty line dispatches the event, and an event without `data:` is not
58/// dispatched. Feed it the body as it arrives.
59#[derive(Clone, Debug, Default)]
60pub struct Decoder {
61    buffer: Vec<u8>,
62    after_cr: bool,
63    started: bool,
64    event: Option<String>,
65    data: String,
66    has_data: bool,
67    id: Option<String>,
68}
69
70impl Decoder {
71    /// A parser at the start of a stream.
72    pub fn new() -> Self {
73        Self::default()
74    }
75
76    /// Takes the next chunk of the body and appends the events it completes to `events`.
77    pub fn feed(&mut self, chunk: &[u8], events: &mut VecDeque<Event>) {
78        let mut chunk = chunk;
79        if !self.started {
80            // A UTF-8 byte order mark at the start of the stream is skipped.
81            let pending = self.buffer.len() + chunk.len();
82            let mut head = self.buffer.clone();
83            head.extend_from_slice(&chunk[..chunk.len().min(3)]);
84            if pending < 3 && b"\xEF\xBB\xBF".starts_with(&head) {
85                self.buffer.extend_from_slice(chunk);
86                return;
87            }
88            self.started = true;
89            let mut all = std::mem::take(&mut self.buffer);
90            all.extend_from_slice(chunk);
91            let rest = all.strip_prefix(b"\xEF\xBB\xBF").unwrap_or(&all).to_vec();
92            self.feed_lines(&rest, events);
93            return;
94        }
95        if self.after_cr {
96            if let Some(rest) = chunk.strip_prefix(b"\n") {
97                chunk = rest;
98            }
99            self.after_cr = false;
100        }
101        self.feed_lines(chunk, events);
102    }
103
104    fn feed_lines(&mut self, chunk: &[u8], events: &mut VecDeque<Event>) {
105        let mut start = 0;
106        let mut i = 0;
107        while i < chunk.len() {
108            match chunk[i] {
109                b'\n' | b'\r' => {
110                    let mut line = std::mem::take(&mut self.buffer);
111                    line.extend_from_slice(&chunk[start..i]);
112                    self.line(&line, events);
113                    if chunk[i] == b'\r' {
114                        if i + 1 < chunk.len() {
115                            if chunk[i + 1] == b'\n' {
116                                i += 1;
117                            }
118                        } else {
119                            self.after_cr = true;
120                        }
121                    }
122                    start = i + 1;
123                }
124                _ => {}
125            }
126            i += 1;
127        }
128        self.buffer.extend_from_slice(&chunk[start..]);
129    }
130
131    fn line(&mut self, line: &[u8], events: &mut VecDeque<Event>) {
132        if line.is_empty() {
133            if self.has_data {
134                events.push_back(Event {
135                    event: self.event.take(),
136                    data: std::mem::take(&mut self.data),
137                    id: self.id.clone(),
138                });
139            }
140            self.event = None;
141            self.data.clear();
142            self.has_data = false;
143            return;
144        }
145        if line[0] == b':' {
146            return;
147        }
148        let line = String::from_utf8_lossy(line);
149        let (field, value) = match line.find(':') {
150            Some(colon) => {
151                let value = &line[colon + 1..];
152                (&line[..colon], value.strip_prefix(' ').unwrap_or(value))
153            }
154            None => (&line[..], ""),
155        };
156        match field {
157            "data" => {
158                if self.has_data {
159                    self.data.push('\n');
160                }
161                self.data.push_str(value);
162                self.has_data = true;
163            }
164            "event" => self.event = Some(value.to_owned()),
165            "id" if !value.contains('\0') => self.id = Some(value.to_owned()),
166            _ => {}
167        }
168    }
169
170    /// The body ended: an event not completed by an empty line is discarded, as the standard requires.
171    pub fn finish(&mut self) {
172        self.buffer.clear();
173        self.event = None;
174        self.data.clear();
175        self.has_data = false;
176    }
177}
178
179/// The events of a `text/event-stream` body, as they come. A `data: [DONE]` event ends it, as does the end of the body
180/// or a failure reading it (yielded as `Err` first); the body is dropped then, which frees the connection.
181pub struct Events {
182    bytes: Option<Body>,
183    decoder: Decoder,
184    ready: VecDeque<Event>,
185    #[cfg(feature = "time")]
186    idle: Option<Idle>,
187}
188
189/// The idle timeout of [`Events`]: its limit, and the timer started at the first read (so that `Events` can be made
190/// outside a tokio runtime).
191#[cfg(feature = "time")]
192struct Idle {
193    limit: std::time::Duration,
194    timer: Option<crate::timer::Sleep>,
195}
196
197#[cfg(not(target_family = "wasm"))]
198type Chunk = Box<dyn AsRef<[u8]> + Send>;
199#[cfg(target_family = "wasm")]
200type Chunk = Box<dyn AsRef<[u8]>>;
201type Body = crate::maybe::BoxStream<'static, Result<Chunk, Error>>;
202
203impl Events {
204    /// The events of a response body. The status and the content type are not checked.
205    pub fn from_response(response: Response) -> Self {
206        Self::from_bytes(response.bytes_stream())
207    }
208
209    /// The events of a body read as a stream of chunks. The stream and its chunks must be `Send`, except on
210    /// WebAssembly ([`MaybeSend`]).
211    pub fn from_bytes<S, B>(bytes: S) -> Self
212    where
213        S: Stream<Item = Result<B, Error>> + MaybeSend + 'static,
214        B: AsRef<[u8]> + MaybeSend + 'static,
215    {
216        Self {
217            bytes: Some(Box::pin(Chunks(Box::pin(bytes)))),
218            decoder: Decoder::new(),
219            ready: VecDeque::new(),
220            #[cfg(feature = "time")]
221            idle: None,
222        }
223    }
224
225    /// Ends the stream with an error when no bytes arrive for `limit`: a server or a proxy that keeps the connection
226    /// open but stopped sending. The error ([`Error::is_timeout`] and [`Error::is_body`]) is the last item, and the
227    /// body is dropped. Every chunk restarts the clock, comments (`: keep-alive`) too; there is no limit on the whole
228    /// stream. Feature `time`; the timer is tokio's, or the JavaScript host's `setTimeout` on WebAssembly.
229    ///
230    /// ```no_run
231    /// # async fn demo(response: timvw_rust_runtime::Response) {
232    /// use std::time::Duration;
233    ///
234    /// let events = response.events().idle_timeout(Duration::from_secs(300));
235    /// # }
236    /// ```
237    #[cfg(feature = "time")]
238    pub fn idle_timeout(mut self, limit: std::time::Duration) -> Self {
239        self.idle = Some(Idle { limit, timer: None });
240        self
241    }
242
243    /// Whether the idle timeout ran out while waiting for bytes; restarts it when `reset`.
244    #[cfg(feature = "time")]
245    fn idle(&mut self, cx: &mut Context<'_>, reset: bool) -> Option<Error> {
246        let idle = self.idle.as_mut()?;
247        let deadline = crate::timer::Instant::now() + idle.limit;
248        match &mut idle.timer {
249            Some(timer) if reset => timer.reset(deadline),
250            Some(_) => {}
251            None => idle.timer = Some(crate::timer::Sleep::until(deadline)),
252        }
253        if reset {
254            return None;
255        }
256        let timer = idle.timer.as_mut()?;
257        match Pin::new(timer).poll(cx) {
258            Poll::Ready(()) => Some(Error::body(crate::error::TimedOut(format!(
259                "no server-sent event data for {:?}",
260                idle.limit
261            )))),
262            Poll::Pending => None,
263        }
264    }
265}
266
267impl Stream for Events {
268    type Item = Result<Event, Error>;
269
270    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
271        let this = &mut *self;
272        loop {
273            if let Some(event) = this.ready.pop_front() {
274                if event.data == DONE {
275                    this.ready.clear();
276                    this.bytes = None;
277                    return Poll::Ready(None);
278                }
279                return Poll::Ready(Some(Ok(event)));
280            }
281            let Some(bytes) = this.bytes.as_mut() else {
282                return Poll::Ready(None);
283            };
284            match bytes.as_mut().poll_next(cx) {
285                Poll::Pending => {
286                    #[cfg(feature = "time")]
287                    if let Some(error) = this.idle(cx, false) {
288                        this.bytes = None;
289                        this.decoder.finish();
290                        this.ready.clear();
291                        return Poll::Ready(Some(Err(error)));
292                    }
293                    return Poll::Pending;
294                }
295                Poll::Ready(Some(Ok(chunk))) => {
296                    #[cfg(feature = "time")]
297                    this.idle(cx, true);
298                    this.decoder.feed((*chunk).as_ref(), &mut this.ready)
299                }
300                Poll::Ready(Some(Err(error))) => {
301                    this.bytes = None;
302                    this.decoder.finish();
303                    return Poll::Ready(Some(Err(error)));
304                }
305                Poll::Ready(None) => {
306                    this.bytes = None;
307                    this.decoder.finish();
308                }
309            }
310        }
311    }
312}
313
314impl fmt::Debug for Events {
315    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
316        let mut debug = f.debug_struct("Events");
317        #[cfg(feature = "time")]
318        debug.field("idle_timeout", &self.idle.as_ref().map(|idle| idle.limit));
319        debug
320            .field("open", &self.bytes.is_some())
321            .field("ready", &self.ready)
322            .finish_non_exhaustive()
323    }
324}
325
326/// A byte stream whose chunks are boxed, so that [`Events`] holds any body.
327struct Chunks<S>(Pin<Box<S>>);
328
329impl<S, B> Stream for Chunks<S>
330where
331    S: Stream<Item = Result<B, Error>>,
332    B: AsRef<[u8]> + MaybeSend + 'static,
333{
334    type Item = Result<Chunk, Error>;
335
336    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
337        self.0
338            .as_mut()
339            .poll_next(cx)
340            .map(|item| item.map(|chunk| chunk.map(|bytes| Box::new(bytes) as Chunk)))
341    }
342}
343
344impl Events {
345    /// Decodes the `data:` payload of every event as JSON, to `T`.
346    ///
347    /// A payload that does not decode yields an error for which [`Error::is_decode`] is true, and the stream goes on
348    /// with the next event; an event with an empty payload yields nothing. A failure reading the body yields an error
349    /// and ends the stream.
350    pub fn json<T: DeserializeOwned>(self) -> JsonEvents<T> {
351        JsonEvents {
352            events: self,
353            types: PhantomData,
354        }
355    }
356}
357
358/// The events of a stream with their `data:` decoded as JSON: [`Events::json`].
359pub struct JsonEvents<T> {
360    events: Events,
361    types: PhantomData<fn() -> T>,
362}
363
364impl<T> JsonEvents<T> {
365    /// The events not yet read, undecoded.
366    pub fn into_events(self) -> Events {
367        self.events
368    }
369}
370
371impl<T: DeserializeOwned> Stream for JsonEvents<T> {
372    type Item = Result<T, Error>;
373
374    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
375        loop {
376            match Pin::new(&mut self.events).poll_next(cx) {
377                Poll::Ready(Some(Ok(event))) if event.data.trim().is_empty() => {}
378                Poll::Ready(Some(Ok(event))) => {
379                    return Poll::Ready(Some(
380                        serde_json::from_str(&event.data).map_err(Error::decode),
381                    ));
382                }
383                Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
384                Poll::Ready(None) => return Poll::Ready(None),
385                Poll::Pending => return Poll::Pending,
386            }
387        }
388    }
389}
390
391impl<T> fmt::Debug for JsonEvents<T> {
392    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
393        f.debug_struct("JsonEvents")
394            .field("events", &self.events)
395            .finish()
396    }
397}
398
399#[cfg(test)]
400mod tests {
401    use super::*;
402    use futures_util::StreamExt;
403
404    fn events(chunks: &[&'static str]) -> Events {
405        let chunks: Vec<Result<&'static [u8], Error>> =
406            chunks.iter().map(|chunk| Ok(chunk.as_bytes())).collect();
407        Events::from_bytes(futures_util::stream::iter(chunks))
408    }
409
410    #[test]
411    fn parses_across_chunks_and_line_endings() {
412        let mut decoder = Decoder::new();
413        let mut out = VecDeque::new();
414        decoder.feed(b"\xEF\xBB\xBFevent: a\r", &mut out);
415        decoder.feed(b"\ndata: 1\r\ndata:", &mut out);
416        decoder.feed(b"2\n: comment\nid: 7\n\ndata: x\n", &mut out);
417        decoder.finish();
418        assert_eq!(
419            out.into_iter().collect::<Vec<_>>(),
420            vec![Event {
421                event: Some("a".into()),
422                data: "1\n2".into(),
423                id: Some("7".into()),
424            }]
425        );
426    }
427
428    #[tokio::test]
429    async fn done_ends_the_stream() {
430        let data: Vec<_> = events(&["data: 1\n\n", "data: [DONE]\n\ndata: 2\n\n"])
431            .map(|event| event.unwrap().data)
432            .collect()
433            .await;
434        assert_eq!(data, vec!["1"]);
435    }
436
437    #[cfg(feature = "time")]
438    #[tokio::test(start_paused = true)]
439    async fn an_idle_stream_times_out() {
440        use std::time::Duration;
441        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<Result<&'static [u8], Error>>();
442        let rx = futures_util::stream::poll_fn(move |cx| rx.poll_recv(cx));
443        let mut events = Events::from_bytes(rx).idle_timeout(Duration::from_secs(10));
444        tx.send(Ok(b"data: 1\n\n".as_slice())).unwrap();
445        assert_eq!(events.next().await.unwrap().unwrap().data, "1");
446        tokio::time::sleep(Duration::from_secs(8)).await;
447        // a keep-alive comment restarts the clock
448        tx.send(Ok(b": ping\n".as_slice())).unwrap();
449        tokio::time::sleep(Duration::from_secs(8)).await;
450        tx.send(Ok(b"data: 2\n\n".as_slice())).unwrap();
451        assert_eq!(events.next().await.unwrap().unwrap().data, "2");
452        let started = tokio::time::Instant::now();
453        let error = events.next().await.unwrap().unwrap_err();
454        assert!(error.is_timeout() && error.is_body(), "{error:?}");
455        assert_eq!(started.elapsed(), Duration::from_secs(10));
456        assert!(events.next().await.is_none());
457        drop(tx);
458    }
459
460    #[tokio::test]
461    async fn decodes_json_and_goes_on_after_a_bad_payload() {
462        #[derive(serde::Deserialize, Debug, PartialEq)]
463        struct Chunk {
464            n: u32,
465        }
466        let items: Vec<_> =
467            events(&["data: {\"n\":1}\n\ndata: oops\n\ndata:\n\ndata: {\"n\":2}\n\n"])
468                .json::<Chunk>()
469                .collect()
470                .await;
471        assert_eq!(items.len(), 3);
472        assert_eq!(items[0].as_ref().unwrap(), &Chunk { n: 1 });
473        assert!(items[1].as_ref().unwrap_err().is_decode());
474        assert_eq!(items[2].as_ref().unwrap(), &Chunk { n: 2 });
475    }
476}
timvw_rust_runtime/sse.rs

timvw_rust_runtime/
sse.rs
