http_sse/
server.rs

1// Copyright 2019 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use crate::Event;
6use futures::channel::{mpsc, oneshot};
7use futures::future::Future;
8use futures::lock::Mutex;
9use futures::stream::Stream;
10use futures::task::{Context, Poll};
11use http_body_util::StreamBody;
12use hyper::body::Bytes;
13use hyper::{Response, StatusCode};
14use std::mem::replace;
15use std::ops::DerefMut;
16use std::pin::Pin;
17use std::sync::Arc;
18
19pub type Body = StreamBody<BodyAbortStream>;
20
21pub struct SseResponseCreator {
22    buffer_size: usize,
23    clients: Arc<Mutex<Vec<Client>>>,
24}
25
26impl SseResponseCreator {
27    /// hyper `Response` `Body`s created by this `SseResponseCreator` will buffer
28    /// `buffer_size + 1` `Events` before the `Body` stream is closed for falling too far behind.
29    pub fn with_additional_buffer_size(buffer_size: usize) -> (Self, EventSender) {
30        let clients = Arc::new(Mutex::new(vec![]));
31        (Self { buffer_size, clients: Arc::clone(&clients) }, EventSender { clients })
32    }
33
34    /// Creates hyper `Response`s whose `Body`s receive `Event`s from the `EventSender` associated
35    /// with this `SseResponseCreator`.
36    pub async fn create(&self) -> Response<Body> {
37        let (abort_tx, abort_rx) = oneshot::channel();
38        let (chunk_tx, chunk_rx) = mpsc::channel(self.buffer_size);
39        self.clients.lock().await.push(Client { abort_tx, chunk_tx });
40        Response::builder()
41            .status(StatusCode::OK)
42            .header("content-type", "text/event-stream")
43            .body(StreamBody::new(BodyAbortStream { abort_rx, chunk_rx }))
44            .unwrap() // builder arguments are all statically determined, build will not fail
45    }
46}
47
48pub struct EventSender {
49    clients: Arc<Mutex<Vec<Client>>>,
50}
51
52impl EventSender {
53    /// Send an `Event` to each connected client. Clients that have fallen too far behind have
54    /// their connections closed.
55    pub async fn send(&self, event: &Event) {
56        let mut clients_guard = self.clients.lock().await;
57        let clients = replace(DerefMut::deref_mut(&mut clients_guard), vec![]);
58        let clients = clients
59            .into_iter()
60            .filter_map(|mut c| {
61                if c.try_send(event).is_ok() {
62                    Some(c)
63                } else {
64                    let _ = c.abort();
65                    None
66                }
67            })
68            .collect();
69        *clients_guard = clients;
70    }
71
72    /// Number of clients that `send` will attempt to communicate with.
73    pub async fn client_count(&self) -> usize {
74        self.clients.lock().await.len()
75    }
76
77    /// Drops all connected clients. Already existing `Response<Body>`s created by the
78    /// `SseResponseCreator` should return error on subsequent `poll_next`.
79    pub async fn drop_all_clients(&self) {
80        self.clients.lock().await.clear();
81    }
82}
83
84// reimplementation of the body created by hyper::body::body::channel() b/c hyper doesn't allow
85// specifying the buffer size and doesn't provide an abort channel.
86pub struct BodyAbortStream {
87    abort_rx: oneshot::Receiver<()>,
88    chunk_rx: mpsc::Receiver<Bytes>,
89}
90
91impl Stream for BodyAbortStream {
92    type Item = Result<hyper::body::Frame<Bytes>, &'static str>;
93
94    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
95        if let Poll::Ready(_) = Pin::new(&mut self.abort_rx).poll(cx) {
96            return Poll::Ready(Some(Err("client dropped")));
97        }
98        match Pin::new(&mut self.chunk_rx).poll_next(cx) {
99            Poll::Ready(Some(chunk)) => Poll::Ready(Some(Ok(hyper::body::Frame::data(chunk)))),
100            Poll::Ready(None) => Poll::Ready(None),
101            Poll::Pending => Poll::Pending,
102        }
103    }
104}
105
106// Reimplementation of hyper::body::Sender b/c in-tree hyper doesn't allow specifying the buffer
107// size and doesn't provide an abort channel.
108struct Client {
109    abort_tx: oneshot::Sender<()>,
110    chunk_tx: mpsc::Sender<Bytes>,
111}
112
113impl Client {
114    fn try_send(&mut self, event: &Event) -> Result<(), ()> {
115        self.chunk_tx.try_send(event.to_vec().into()).map_err(|_| ())
116    }
117    fn abort(self) {
118        let _ = self.abort_tx.send(());
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125    use assert_matches::assert_matches;
126    use fuchsia_async::{self as fasync};
127    use futures::StreamExt;
128
129    #[fasync::run_singlethreaded(test)]
130    async fn response_headers() {
131        let (sse_response_creator, _) = SseResponseCreator::with_additional_buffer_size(0);
132        let resp = sse_response_creator.create().await;
133
134        assert_eq!(resp.status(), StatusCode::OK);
135        assert_eq!(
136            resp.headers().get("content-type").map(|h| h.as_bytes()),
137            Some(&b"text/event-stream"[..])
138        );
139    }
140
141    #[fasync::run_singlethreaded(test)]
142    async fn response_correct_body_single_event() {
143        let event = Event::from_type_and_data("event_type", "data_contents").unwrap();
144        let (sse_response_creator, event_sender) =
145            SseResponseCreator::with_additional_buffer_size(0);
146        let resp = sse_response_creator.create().await;
147
148        event_sender.send(&event).await;
149        let mut body_stream = resp.into_body();
150        let body_bytes = body_stream.next().await;
151
152        assert_eq!(body_bytes.unwrap().unwrap().data_ref().unwrap().to_vec(), event.to_vec());
153    }
154
155    #[fasync::run_singlethreaded(test)]
156    async fn full_client_dropped_other_clients_continue_to_receive_events() {
157        let event0 = Event::from_type_and_data("event_type0", "data_contents0").unwrap();
158        let (sse_response_creator, event_sender) =
159            SseResponseCreator::with_additional_buffer_size(0);
160        assert_eq!(event_sender.client_count().await, 0);
161
162        let mut body_stream0 = sse_response_creator.create().await.into_body();
163        let mut body_stream1 = sse_response_creator.create().await.into_body();
164        assert_eq!(event_sender.client_count().await, 2);
165
166        event_sender.send(&event0).await;
167
168        let body_bytes1 = body_stream1.next().await;
169
170        assert_matches!(body_bytes1, Some(Ok(chunk)) if chunk.data_ref().unwrap().to_vec() == event0.to_vec());
171
172        let event1 = Event::from_type_and_data("event_type1", "data_contents1").unwrap();
173        event_sender.send(&event1).await;
174        assert_eq!(event_sender.client_count().await, 1);
175
176        let body_bytes0 = body_stream0.next().await;
177        assert_matches!(body_bytes0, Some(Err(_)));
178
179        let body_bytes1 = body_stream1.next().await;
180        assert_eq!(body_bytes1.unwrap().unwrap().data_ref().unwrap().to_vec(), event1.to_vec());
181    }
182
183    #[fasync::run_singlethreaded(test)]
184    async fn drop_all_clients() {
185        let (sse_response_creator, event_sender) =
186            SseResponseCreator::with_additional_buffer_size(0);
187        let mut body_stream0 = sse_response_creator.create().await.into_body();
188        let mut body_stream1 = sse_response_creator.create().await.into_body();
189        assert_eq!(event_sender.client_count().await, 2);
190        event_sender.send(&Event::from_type_and_data("event_type", "data_contents").unwrap()).await;
191
192        event_sender.drop_all_clients().await;
193
194        assert_eq!(event_sender.client_count().await, 0);
195        assert_matches!(body_stream0.next().await, Some(Err(_)));
196        assert_matches!(body_stream1.next().await, Some(Err(_)));
197    }
198}
http_sse/server.rs

http_sse/
server.rs
