Implement SSE using responses (#98)

src/response/mod.rs

@@ -0,0 +1,466 @@
+//! Types and traits for generating responses.
+
+use crate::{
+    body::{box_body, BoxBody},
+    Error,
+};
+use bytes::Bytes;
+use http::{header, HeaderMap, HeaderValue, Response, StatusCode};
+use http_body::{
+    combinators::{MapData, MapErr},
+    Empty, Full,
+};
+use std::{borrow::Cow, convert::Infallible};
+use tower::{util::Either, BoxError};
+
+#[doc(no_inline)]
+pub use crate::Json;
+
+pub mod sse;
+
+pub use sse::{sse, Sse};
+
+/// Trait for generating responses.
+///
+/// Types that implement `IntoResponse` can be returned from handlers.
+///
+/// # Implementing `IntoResponse`
+///
+/// You generally shouldn't have to implement `IntoResponse` manually, as axum
+/// provides implementations for many common types.
+///
+/// A manual implementation should only be necessary if you have a custom
+/// response body type:
+///
+/// ```rust
+/// use axum::{prelude::*, response::IntoResponse};
+/// use http_body::Body;
+/// use http::{Response, HeaderMap};
+/// use bytes::Bytes;
+/// use std::{
+///     convert::Infallible,
+///     task::{Poll, Context},
+///     pin::Pin,
+/// };
+///
+/// struct MyBody;
+///
+/// // First implement `Body` for `MyBody`. This could for example use
+/// // some custom streaming protocol.
+/// impl Body for MyBody {
+///     type Data = Bytes;
+///     type Error = Infallible;
+///
+///     fn poll_data(
+///         self: Pin<&mut Self>,
+///         cx: &mut Context<'_>
+///     ) -> Poll<Option<Result<Self::Data, Self::Error>>> {
+///         # unimplemented!()
+///         // ...
+///     }
+///
+///     fn poll_trailers(
+///         self: Pin<&mut Self>,
+///         cx: &mut Context<'_>
+///     ) -> Poll<Result<Option<HeaderMap>, Self::Error>> {
+///         # unimplemented!()
+///         // ...
+///     }
+/// }
+///
+/// // Now we can implement `IntoResponse` directly for `MyBody`
+/// impl IntoResponse for MyBody {
+///     type Body = Self;
+///     type BodyError = <Self as Body>::Error;
+///
+///     fn into_response(self) -> Response<Self::Body> {
+///         Response::new(self)
+///     }
+/// }
+///
+/// // We don't need to implement `IntoResponse for Response<MyBody>` as that is
+/// // covered by a blanket implementation in axum.
+///
+/// // `MyBody` can now be returned from handlers.
+/// let app = route("/", get(|| async { MyBody }));
+/// # async {
+/// # hyper::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap();
+/// # };
+/// ```
+pub trait IntoResponse {
+    /// The body type of the response.
+    ///
+    /// Unless you're implementing this trait for a custom body type, these are
+    /// some common types you can use:
+    ///
+    /// - [`axum::body::Body`]: A good default that supports most use cases.
+    /// - [`axum::body::Empty<Bytes>`]: When you know your response is always
+    /// empty.
+    /// - [`axum::body::Full<Bytes>`]: When you know your response always
+    /// contains exactly one chunk.
+    /// - [`axum::body::BoxBody`]: If you need to unify multiple body types into
+    /// one, or return a body type that cannot be named. Can be created with
+    /// [`box_body`].
+    ///
+    /// [`axum::body::Body`]: crate::body::Body
+    /// [`axum::body::Empty<Bytes>`]: crate::body::Empty
+    /// [`axum::body::Full<Bytes>`]: crate::body::Full
+    /// [`axum::body::BoxBody`]: crate::body::BoxBody
+    type Body: http_body::Body<Data = Bytes, Error = Self::BodyError> + Send + Sync + 'static;
+
+    /// The error type `Self::Body` might generate.
+    ///
+    /// Generally it should be possible to set this to:
+    ///
+    /// ```rust,ignore
+    /// type BodyError = <Self::Body as axum::body::HttpBody>::Error;
+    /// ```
+    ///
+    /// This associated type exists mainly to make returning `impl IntoResponse`
+    /// possible and to simplify trait bounds internally in axum.
+    type BodyError: Into<BoxError>;
+
+    /// Create a response.
+    fn into_response(self) -> Response<Self::Body>;
+}
+
+impl IntoResponse for () {
+    type Body = Empty<Bytes>;
+    type BodyError = Infallible;
+
+    fn into_response(self) -> Response<Self::Body> {
+        Response::new(Empty::new())
+    }
+}
+
+impl IntoResponse for Infallible {
+    type Body = Empty<Bytes>;
+    type BodyError = Infallible;
+
+    fn into_response(self) -> Response<Self::Body> {
+        match self {}
+    }
+}
+
+impl<T, K> IntoResponse for Either<T, K>
+where
+    T: IntoResponse,
+    K: IntoResponse,
+{
+    type Body = BoxBody;
+    type BodyError = Error;
+
+    fn into_response(self) -> Response<Self::Body> {
+        match self {
+            Either::A(inner) => inner.into_response().map(box_body),
+            Either::B(inner) => inner.into_response().map(box_body),
+        }
+    }
+}
+
+impl<T, E> IntoResponse for Result<T, E>
+where
+    T: IntoResponse,
+    E: IntoResponse,
+{
+    type Body = BoxBody;
+    type BodyError = Error;
+
+    fn into_response(self) -> Response<Self::Body> {
+        match self {
+            Ok(value) => value.into_response().map(box_body),
+            Err(err) => err.into_response().map(box_body),
+        }
+    }
+}
+
+impl<B> IntoResponse for Response<B>
+where
+    B: http_body::Body<Data = Bytes> + Send + Sync + 'static,
+    B::Error: Into<BoxError>,
+{
+    type Body = B;
+    type BodyError = <B as http_body::Body>::Error;
+
+    fn into_response(self) -> Self {
+        self
+    }
+}
+
+macro_rules! impl_into_response_for_body {
+    ($body:ty) => {
+        impl IntoResponse for $body {
+            type Body = $body;
+            type BodyError = <$body as http_body::Body>::Error;
+
+            fn into_response(self) -> Response<Self> {
+                Response::new(self)
+            }
+        }
+    };
+}
+
+impl_into_response_for_body!(hyper::Body);
+impl_into_response_for_body!(Full<Bytes>);
+impl_into_response_for_body!(Empty<Bytes>);
+
+impl<E> IntoResponse for http_body::combinators::BoxBody<Bytes, E>
+where
+    E: Into<BoxError> + 'static,
+{
+    type Body = Self;
+    type BodyError = E;
+
+    fn into_response(self) -> Response<Self> {
+        Response::new(self)
+    }
+}
+
+impl<B, F> IntoResponse for MapData<B, F>
+where
+    B: http_body::Body + Send + Sync + 'static,
+    F: FnMut(B::Data) -> Bytes + Send + Sync + 'static,
+    B::Error: Into<BoxError>,
+{
+    type Body = Self;
+    type BodyError = <B as http_body::Body>::Error;
+
+    fn into_response(self) -> Response<Self::Body> {
+        Response::new(self)
+    }
+}
+
+impl<B, F, E> IntoResponse for MapErr<B, F>
+where
+    B: http_body::Body<Data = Bytes> + Send + Sync + 'static,
+    F: FnMut(B::Error) -> E + Send + Sync + 'static,
+    E: Into<BoxError>,
+{
+    type Body = Self;
+    type BodyError = E;
+
+    fn into_response(self) -> Response<Self::Body> {
+        Response::new(self)
+    }
+}
+
+impl IntoResponse for &'static str {
+    type Body = Full<Bytes>;
+    type BodyError = Infallible;
+
+    #[inline]
+    fn into_response(self) -> Response<Self::Body> {
+        Cow::Borrowed(self).into_response()
+    }
+}
+
+impl IntoResponse for String {
+    type Body = Full<Bytes>;
+    type BodyError = Infallible;
+
+    #[inline]
+    fn into_response(self) -> Response<Self::Body> {
+        Cow::<'static, str>::Owned(self).into_response()
+    }
+}
+
+impl IntoResponse for std::borrow::Cow<'static, str> {
+    type Body = Full<Bytes>;
+    type BodyError = Infallible;
+
+    fn into_response(self) -> Response<Self::Body> {
+        let mut res = Response::new(Full::from(self));
+        res.headers_mut()
+            .insert(header::CONTENT_TYPE, HeaderValue::from_static("text/plain"));
+        res
+    }
+}
+
+impl IntoResponse for Bytes {
+    type Body = Full<Bytes>;
+    type BodyError = Infallible;
+
+    fn into_response(self) -> Response<Self::Body> {
+        let mut res = Response::new(Full::from(self));
+        res.headers_mut().insert(
+            header::CONTENT_TYPE,
+            HeaderValue::from_static("application/octet-stream"),
+        );
+        res
+    }
+}
+
+impl IntoResponse for &'static [u8] {
+    type Body = Full<Bytes>;
+    type BodyError = Infallible;
+
+    fn into_response(self) -> Response<Self::Body> {
+        let mut res = Response::new(Full::from(self));
+        res.headers_mut().insert(
+            header::CONTENT_TYPE,
+            HeaderValue::from_static("application/octet-stream"),
+        );
+        res
+    }
+}
+
+impl IntoResponse for Vec<u8> {
+    type Body = Full<Bytes>;
+    type BodyError = Infallible;
+
+    fn into_response(self) -> Response<Self::Body> {
+        let mut res = Response::new(Full::from(self));
+        res.headers_mut().insert(
+            header::CONTENT_TYPE,
+            HeaderValue::from_static("application/octet-stream"),
+        );
+        res
+    }
+}
+
+impl IntoResponse for std::borrow::Cow<'static, [u8]> {
+    type Body = Full<Bytes>;
+    type BodyError = Infallible;
+
+    fn into_response(self) -> Response<Self::Body> {
+        let mut res = Response::new(Full::from(self));
+        res.headers_mut().insert(
+            header::CONTENT_TYPE,
+            HeaderValue::from_static("application/octet-stream"),
+        );
+        res
+    }
+}
+
+impl IntoResponse for StatusCode {
+    type Body = Empty<Bytes>;
+    type BodyError = Infallible;
+
+    fn into_response(self) -> Response<Self::Body> {
+        Response::builder().status(self).body(Empty::new()).unwrap()
+    }
+}
+
+impl<T> IntoResponse for (StatusCode, T)
+where
+    T: IntoResponse,
+{
+    type Body = T::Body;
+    type BodyError = T::BodyError;
+
+    fn into_response(self) -> Response<T::Body> {
+        let mut res = self.1.into_response();
+        *res.status_mut() = self.0;
+        res
+    }
+}
+
+impl<T> IntoResponse for (HeaderMap, T)
+where
+    T: IntoResponse,
+{
+    type Body = T::Body;
+    type BodyError = T::BodyError;
+
+    fn into_response(self) -> Response<T::Body> {
+        let mut res = self.1.into_response();
+        res.headers_mut().extend(self.0);
+        res
+    }
+}
+
+impl<T> IntoResponse for (StatusCode, HeaderMap, T)
+where
+    T: IntoResponse,
+{
+    type Body = T::Body;
+    type BodyError = T::BodyError;
+
+    fn into_response(self) -> Response<T::Body> {
+        let mut res = self.2.into_response();
+        *res.status_mut() = self.0;
+        res.headers_mut().extend(self.1);
+        res
+    }
+}
+
+impl IntoResponse for HeaderMap {
+    type Body = Empty<Bytes>;
+    type BodyError = Infallible;
+
+    fn into_response(self) -> Response<Self::Body> {
+        let mut res = Response::new(Empty::new());
+        *res.headers_mut() = self;
+        res
+    }
+}
+
+/// An HTML response.
+///
+/// Will automatically get `Content-Type: text/html`.
+#[derive(Clone, Copy, Debug)]
+pub struct Html<T>(pub T);
+
+impl<T> IntoResponse for Html<T>
+where
+    T: Into<Full<Bytes>>,
+{
+    type Body = Full<Bytes>;
+    type BodyError = Infallible;
+
+    fn into_response(self) -> Response<Self::Body> {
+        let mut res = Response::new(self.0.into());
+        res.headers_mut()
+            .insert(header::CONTENT_TYPE, HeaderValue::from_static("text/html"));
+        res
+    }
+}
+
+impl<T> From<T> for Html<T> {
+    fn from(inner: T) -> Self {
+        Self(inner)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::body::Body;
+    use http::header::{HeaderMap, HeaderName};
+
+    #[test]
+    fn test_merge_headers() {
+        struct MyResponse;
+
+        impl IntoResponse for MyResponse {
+            type Body = Body;
+            type BodyError = <Self::Body as http_body::Body>::Error;
+
+            fn into_response(self) -> Response<Body> {
+                let mut resp = Response::new(String::new().into());
+                resp.headers_mut()
+                    .insert(HeaderName::from_static("a"), HeaderValue::from_static("1"));
+                resp
+            }
+        }
+
+        fn check(resp: impl IntoResponse) {
+            let resp = resp.into_response();
+            assert_eq!(
+                resp.headers().get(HeaderName::from_static("a")).unwrap(),
+                &HeaderValue::from_static("1")
+            );
+            assert_eq!(
+                resp.headers().get(HeaderName::from_static("b")).unwrap(),
+                &HeaderValue::from_static("2")
+            );
+        }
+
+        let headers: HeaderMap =
+            std::iter::once((HeaderName::from_static("b"), HeaderValue::from_static("2")))
+                .collect();
+
+        check((headers.clone(), MyResponse));
+        check((StatusCode::OK, headers, MyResponse));
+    }
+}

src/response/sse.rs

@@ -0,0 +1,387 @@
+//! Server-Sent Events (SSE) responses.
+//!
+//! # Example
+//!
+//! ```
+//! use axum::prelude::*;
+//! use axum::response::sse::{sse, Event, KeepAlive, Sse};
+//! use std::{time::Duration, convert::Infallible};
+//! use tokio_stream::StreamExt as _ ;
+//! use futures::stream::{self, Stream};
+//!
+//! let app = route("/sse", get(sse_handler));
+//!
+//! async fn sse_handler() -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
+//!     // A `Stream` that repeats an event every second
+//!     let stream = stream::repeat_with(|| Event::default().data("hi!"))
+//!         .map(Ok)
+//!         .throttle(Duration::from_secs(1));
+//!
+//!     sse(stream).keep_alive(KeepAlive::default())
+//! }
+//! # async {
+//! # hyper::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap();
+//! # };
+//! ```
+
+use crate::response::IntoResponse;
+use bytes::Bytes;
+use futures_util::{
+    ready,
+    stream::{Stream, TryStream},
+};
+use http::Response;
+use http_body::Body as HttpBody;
+use pin_project_lite::pin_project;
+use serde::Serialize;
+use std::{
+    borrow::Cow,
+    fmt,
+    fmt::Write,
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+    time::Duration,
+};
+use sync_wrapper::SyncWrapper;
+use tokio::time::Sleep;
+use tower::BoxError;
+
+/// Create a new [`Sse`] response that will respond with the given stream of
+/// [`Event`]s.
+///
+/// See the [module docs](self) for more details.
+pub fn sse<S>(stream: S) -> Sse<S>
+where
+    S: TryStream<Ok = Event> + Send + 'static,
+    S::Error: Into<BoxError>,
+{
+    Sse {
+        stream,
+        keep_alive: None,
+    }
+}
+
+/// An SSE response, created by [`sse`].
+#[derive(Clone)]
+pub struct Sse<S> {
+    stream: S,
+    keep_alive: Option<KeepAlive>,
+}
+
+impl<S> Sse<S> {
+    /// Configure the interval between keep-alive messages.
+    ///
+    /// Defaults to no keep-alive messages.
+    pub fn keep_alive(mut self, keep_alive: KeepAlive) -> Self {
+        self.keep_alive = Some(keep_alive);
+        self
+    }
+}
+
+impl<S> fmt::Debug for Sse<S> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Sse")
+            .field("stream", &format_args!("{}", std::any::type_name::<S>()))
+            .field("keep_alive", &self.keep_alive)
+            .finish()
+    }
+}
+
+impl<S, E> IntoResponse for Sse<S>
+where
+    S: Stream<Item = Result<Event, E>> + Send + 'static,
+    E: Into<BoxError>,
+{
+    type Body = Body<S>;
+    type BodyError = E;
+
+    fn into_response(self) -> Response<Self::Body> {
+        let body = Body {
+            event_stream: SyncWrapper::new(self.stream),
+            keep_alive: self.keep_alive.map(KeepAliveStream::new),
+        };
+
+        Response::builder()
+            .header(http::header::CONTENT_TYPE, "text/event-stream")
+            .header(http::header::CACHE_CONTROL, "no-cache")
+            .body(body)
+            .unwrap()
+    }
+}
+
+pin_project! {
+    /// The body of an SSE response.
+    #[derive(Debug)]
+    pub struct Body<S> {
+        #[pin]
+        event_stream: SyncWrapper<S>,
+        #[pin]
+        keep_alive: Option<KeepAliveStream>,
+    }
+}
+
+impl<S, E> HttpBody for Body<S>
+where
+    S: Stream<Item = Result<Event, E>>,
+{
+    type Data = Bytes;
+    type Error = E;
+
+    fn poll_data(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Self::Data, Self::Error>>> {
+        let this = self.project();
+
+        match this.event_stream.get_pin_mut().poll_next(cx) {
+            Poll::Pending => {
+                if let Some(keep_alive) = this.keep_alive.as_pin_mut() {
+                    keep_alive
+                        .poll_event(cx)
+                        .map(|e| Some(Ok(Bytes::from(e.to_string()))))
+                } else {
+                    Poll::Pending
+                }
+            }
+            Poll::Ready(Some(Ok(event))) => {
+                if let Some(keep_alive) = this.keep_alive.as_pin_mut() {
+                    keep_alive.reset();
+                }
+                Poll::Ready(Some(Ok(Bytes::from(event.to_string()))))
+            }
+            Poll::Ready(Some(Err(error))) => Poll::Ready(Some(Err(error))),
+            Poll::Ready(None) => Poll::Ready(None),
+        }
+    }
+
+    fn poll_trailers(
+        self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+    ) -> Poll<Result<Option<http::HeaderMap>, Self::Error>> {
+        Poll::Ready(Ok(None))
+    }
+}
+
+/// Server-sent event
+#[derive(Default, Debug)]
+pub struct Event {
+    name: Option<String>,
+    id: Option<String>,
+    data: Option<DataType>,
+    event: Option<String>,
+    comment: Option<String>,
+    retry: Option<Duration>,
+}
+
+// Server-sent event data type
+#[derive(Debug)]
+enum DataType {
+    Text(String),
+    Json(String),
+}
+
+impl Event {
+    /// Set Server-sent event data
+    /// data field(s) ("data:<content>")
+    pub fn data<T>(mut self, data: T) -> Event
+    where
+        T: Into<String>,
+    {
+        self.data = Some(DataType::Text(data.into()));
+        self
+    }
+
+    /// Set Server-sent event data
+    /// data field(s) ("data:<content>")
+    pub fn json_data<T>(mut self, data: T) -> Result<Event, serde_json::Error>
+    where
+        T: Serialize,
+    {
+        self.data = Some(DataType::Json(serde_json::to_string(&data)?));
+        Ok(self)
+    }
+
+    /// Set Server-sent event comment
+    /// Comment field (":<comment-text>")
+    pub fn comment<T>(mut self, comment: T) -> Event
+    where
+        T: Into<String>,
+    {
+        self.comment = Some(comment.into());
+        self
+    }
+
+    /// Set Server-sent event event
+    /// Event name field ("event:<event-name>")
+    pub fn event<T>(mut self, event: T) -> Event
+    where
+        T: Into<String>,
+    {
+        self.event = Some(event.into());
+        self
+    }
+
+    /// Set Server-sent event retry
+    /// Retry timeout field ("retry:<timeout>")
+    pub fn retry(mut self, duration: Duration) -> Event {
+        self.retry = Some(duration);
+        self
+    }
+
+    /// Set Server-sent event id
+    /// Identifier field ("id:<identifier>")
+    pub fn id<T>(mut self, id: T) -> Event
+    where
+        T: Into<String>,
+    {
+        self.id = Some(id.into());
+        self
+    }
+}
+
+impl fmt::Display for Event {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        if let Some(comment) = &self.comment {
+            ":".fmt(f)?;
+            comment.fmt(f)?;
+            f.write_char('\n')?;
+        }
+
+        if let Some(event) = &self.event {
+            "event:".fmt(f)?;
+            event.fmt(f)?;
+            f.write_char('\n')?;
+        }
+
+        match &self.data {
+            Some(DataType::Text(data)) => {
+                for line in data.split('\n') {
+                    "data:".fmt(f)?;
+                    line.fmt(f)?;
+                    f.write_char('\n')?;
+                }
+            }
+            Some(DataType::Json(data)) => {
+                "data:".fmt(f)?;
+                data.fmt(f)?;
+                f.write_char('\n')?;
+            }
+            None => {}
+        }
+
+        if let Some(id) = &self.id {
+            "id:".fmt(f)?;
+            id.fmt(f)?;
+            f.write_char('\n')?;
+        }
+
+        if let Some(duration) = &self.retry {
+            "retry:".fmt(f)?;
+
+            let secs = duration.as_secs();
+            let millis = duration.subsec_millis();
+
+            if secs > 0 {
+                // format seconds
+                secs.fmt(f)?;
+
+                // pad milliseconds
+                if millis < 10 {
+                    f.write_str("00")?;
+                } else if millis < 100 {
+                    f.write_char('0')?;
+                }
+            }
+
+            // format milliseconds
+            millis.fmt(f)?;
+
+            f.write_char('\n')?;
+        }
+
+        f.write_char('\n')?;
+
+        Ok(())
+    }
+}
+
+/// Configure the interval between keep-alive messages, the content
+/// of each message, and the associated stream.
+#[derive(Debug, Clone)]
+pub struct KeepAlive {
+    comment_text: Cow<'static, str>,
+    max_interval: Duration,
+}
+
+impl KeepAlive {
+    /// Create a new `KeepAlive`.
+    pub fn new() -> Self {
+        Self {
+            comment_text: Cow::Borrowed(""),
+            max_interval: Duration::from_secs(15),
+        }
+    }
+
+    /// Customize the interval between keep-alive messages.
+    ///
+    /// Default is 15 seconds.
+    pub fn interval(mut self, time: Duration) -> Self {
+        self.max_interval = time;
+        self
+    }
+
+    /// Customize the text of the keep-alive message.
+    ///
+    /// Default is an empty comment.
+    pub fn text<I>(mut self, text: I) -> Self
+    where
+        I: Into<Cow<'static, str>>,
+    {
+        self.comment_text = text.into();
+        self
+    }
+}
+
+impl Default for KeepAlive {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+pin_project! {
+    #[derive(Debug)]
+    struct KeepAliveStream {
+        keep_alive: KeepAlive,
+        #[pin]
+        alive_timer: Sleep,
+    }
+}
+
+impl KeepAliveStream {
+    fn new(keep_alive: KeepAlive) -> Self {
+        Self {
+            alive_timer: tokio::time::sleep(keep_alive.max_interval),
+            keep_alive,
+        }
+    }
+
+    fn reset(self: Pin<&mut Self>) {
+        let this = self.project();
+        this.alive_timer
+            .reset(tokio::time::Instant::now() + this.keep_alive.max_interval);
+    }
+
+    fn poll_event(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Event> {
+        let this = self.as_mut().project();
+
+        ready!(this.alive_timer.poll(cx));
+
+        let comment_str = this.keep_alive.comment_text.clone();
+        let event = Event::default().comment(comment_str);
+
+        self.reset();
+
+        Poll::Ready(event)
+    }
+}