rework encoder

Author: yoshuawuyts

Cargo.toml

@@ -17,15 +17,15 @@ authors = [
 [features]
 
 [dependencies]
-async-std = "1.5.0"
+async-std = { version = "1.5.0", features = ["unstable"] }
 http-types = "1.0.1"
 log = "0.4.8"
 memchr = "2.3.3"
 pin-project-lite = "0.1.4"
 
 [dev-dependencies]
 femme = "1.3.0"
-async-std = { version = "1.5.0", features = ["attributes"] }
+async-std = { version = "1.5.0", features = ["attributes", "unstable"] }
 
 [patch.crates-io]
 async-std = { path = "../async-std" }

src/encoder.rs

@@ -1,71 +1,127 @@
-use async_std::io::prelude::*;
-use async_std::io::Write as AsyncWrite;
+use async_std::sync;
 use std::io;
 use std::time::Duration;
 
-/// An SSE protocol encoder.
-#[derive(Debug)]
-pub struct Encoder<W> {
-    writer: W,
-}
+use async_std::io::Read as AsyncRead;
+use async_std::prelude::*;
+use async_std::task::{ready, Context, Poll};
+use std::pin::Pin;
 
-/// Encode a new SSE connection.
-pub fn encode<W: AsyncWrite + Unpin>(writer: W) -> Encoder<W> {
-    Encoder { writer }
+pin_project_lite::pin_project! {
+    /// An SSE protocol encoder.
+    #[derive(Debug)]
+    pub struct Encoder {
+        buf: Option<Vec<u8>>,
+        #[pin]
+        receiver: sync::Receiver<Vec<u8>>,
+        cursor: usize,
+    }
 }
 
-impl<W> Encoder<W> {
-    /// Access the inner writer from the Encoder.
-    pub fn into_writer(self) -> W {
-        self.writer
+impl AsyncRead for Encoder {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<io::Result<usize>> {
+        // Request a new buffer if we don't have one yet.
+        if let None = self.buf {
+            self.buf = match ready!(Pin::new(&mut self.receiver).poll_next(cx)) {
+                Some(buf) => {
+                    log::trace!("> Received a new buffer with len {}", buf.len());
+                    Some(buf)
+                }
+                None => {
+                    log::trace!("> Encoder done reading");
+                    return Poll::Ready(Ok(0));
+                }
+            };
+        };
+
+        // Write the current buffer to completion.
+        let local_buf = self.buf.as_mut().unwrap();
+        let local_len = local_buf.len();
+        let max = buf.len().min(local_buf.len());
+        buf[..max].clone_from_slice(&local_buf[..max]);
+
+        self.cursor += max;
+
+        // Reset values if we're done reading.
+        if self.cursor == local_len {
+            self.buf = None;
+            self.cursor = 0;
+        };
+
+        // Return bytes read.
+        Poll::Ready(Ok(max))
     }
 }
 
-impl<W: AsyncWrite + Unpin> Encoder<W> {
+// impl AsyncBufRead for Encoder {
+//     fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
+//         match ready!(self.project().receiver.poll_next(cx)) {
+//             Some(buf) => match &self.buf {
+//                 None => self.project().buf = &mut Some(buf),
+//                 Some(local_buf) => local_buf.extend(buf),
+//             },
+//             None => {
+//                 if let None = self.buf {
+//                     self.project().buf = &mut Some(vec![]);
+//                 };
+//             }
+//         };
+//         Poll::Ready(Ok(self.buf.as_ref().unwrap()))
+//     }
+
+//     fn consume(self: Pin<&mut Self>, amt: usize) {
+//         Pin::new(self).cursor += amt;
+//     }
+// }
+
+/// The sending side of the encoder.
+#[derive(Debug)]
+pub struct Sender(sync::Sender<Vec<u8>>);
+
+/// Create a new SSE encoder.
+pub fn encode() -> (Sender, Encoder) {
+    let (sender, receiver) = sync::channel(1);
+    let encoder = Encoder {
+        receiver,
+        buf: None,
+        cursor: 0,
+    };
+    (Sender(sender), encoder)
+}
+
+impl Sender {
     /// Send a new message over SSE.
-    pub async fn send(&mut self, name: &str, data: &[u8], id: Option<&str>) -> io::Result<()> {
+    pub async fn send(&self, name: &str, data: &[u8], id: Option<&str>) {
         // Write the event name
-        self.writer.write_all(b"event:").await?;
-        self.writer.write_all(name.as_bytes()).await?;
-        self.writer.write_all(b"\n").await?;
+        let msg = format!("event:{}\n", name);
+        self.0.send(msg.into_bytes()).await;
 
         // Write the id
         if let Some(id) = id {
-            self.writer.write_all(b"id:").await?;
-            self.writer.write_all(id.as_bytes()).await?;
-            self.writer.write_all(b"\n").await?;
+            self.0.send(format!("id:{}\n", id).into_bytes()).await;
         }
 
-        // Write the section
-        self.writer.write_all(b"data:").await?;
-        self.writer.write_all(data).await?;
-        self.writer.write_all(b"\n").await?;
-
-        // Finalize the message
-        self.writer.write_all(b"\n").await?;
-
-        Ok(())
+        // Write the data section, and end.
+        let mut msg = b"data:".to_vec();
+        msg.extend_from_slice(data);
+        msg.extend_from_slice(b"\n\n");
+        self.0.send(msg).await;
     }
 
     /// Send a new "retry" message over SSE.
-    pub async fn send_retry(&mut self, dur: Duration, id: Option<&str>) -> io::Result<()> {
+    pub async fn send_retry(&self, dur: Duration, id: Option<&str>) {
         // Write the id
         if let Some(id) = id {
-            self.writer.write_all(b"id:").await?;
-            self.writer.write_all(id.as_bytes()).await?;
-            self.writer.write_all(b"\n").await?;
+            self.0.send(format!("id:{}\n", id).into_bytes()).await;
         }
 
-        // Write the section
-        self.writer.write_all(b"retry:").await?;
-        self.writer
-            .write_all(&format!("{}", dur.as_secs_f64() as u64).as_bytes())
-            .await?;
-        self.writer.write_all(b"\n").await?;
-
-        // Finalize the message
-        self.writer.write_all(b"\n").await?;
-
-        Ok(())
+        // Write the retry section, and end.
+        let dur = dur.as_secs_f64() as u64;
+        let msg = format!("retry:{}\n\n", dur);
+        self.0.send(msg.into_bytes()).await;
     }
 }

src/event.rs

@@ -2,7 +2,7 @@ use crate::Message;
 
 use std::time::Duration;
 
-/// The kind of event sent.
+/// The kind of SSE event sent.
 #[derive(Debug, Eq, PartialEq)]
 pub enum Event {
     /// A retry frame, signaling a new retry duration must be used..

src/handshake.rs

@@ -0,0 +1,7 @@
+/// Upgrade an HTTP connection into an SSE session.
+pub fn upgrade(headers: &mut impl AsMut<http_types::Headers>) -> http_types::Result<()> {
+    let headers = headers.as_mut();
+    headers.insert("Cache-Control", "no-cache")?;
+    headers.insert("Content-Type", "text/event-stream")?;
+    Ok(())
+}

src/lib.rs

@@ -4,22 +4,20 @@
 //!
 //! ```
 //! use async_sse::{decode, encode, Event};
-//! use async_std::io::Cursor;
 //! use async_std::prelude::*;
+//! use async_std::io::BufReader;
+//! use async_std::task;
 //!
 //! #[async_std::main]
 //! async fn main() -> http_types::Result<()> {
-//!     let buf = Cursor::new(vec![]);
-//!
-//!     // Encode messages to an AsyncWrite.
-//!     let mut encoder = encode(buf);
-//!     encoder.send("cat", b"chashu", None).await?;
-//!
-//!     let mut buf = encoder.into_writer();
-//!     buf.set_position(0);
-//!
-//!     // Decode messages from an AsyncRead.
-//!     let mut reader = decode(buf);
+//!     // Create an encoder + sender pair and send a message.
+//!     let (sender, encoder) = encode();
+//!     task::spawn(async move {
+//!         sender.send("cat", b"chashu", None).await;
+//!     });
+//!
+//!     // Decode messages using a decoder.
+//!     let mut reader = decode(BufReader::new(encoder));
 //!     let event = reader.next().await.unwrap()?;
 //!     // Match and handle the event
 //!
@@ -40,12 +38,14 @@
 mod decoder;
 mod encoder;
 mod event;
+mod handshake;
 mod lines;
 mod message;
 
 pub use decoder::{decode, Decoder};
-pub use encoder::{encode, Encoder};
+pub use encoder::{encode, Encoder, Sender};
 pub use event::Event;
+pub use handshake::upgrade;
 pub use message::Message;
 
 pub(crate) use lines::Lines;

src/message.rs

@@ -1,4 +1,4 @@
-/// A data event.
+/// An SSE event with a data payload.
 #[derive(Debug, PartialEq, Eq, Hash)]
 pub struct Message {
     /// The ID of this event.

tests/encode.rs

@@ -1,6 +1,7 @@
 use async_sse::{decode, encode, Event};
-use async_std::io::Cursor;
+use async_std::io::BufReader;
 use async_std::prelude::*;
+use async_std::task;
 use std::time::Duration;
 
 /// Assert a Message.
@@ -24,43 +25,42 @@ fn assert_retry(event: &Event, dur: u64) {
         assert_eq!(dur, &expected);
     }
 }
+
 #[async_std::test]
 async fn encode_message() -> http_types::Result<()> {
-    let buf = Cursor::new(vec![]);
-    let mut encoder = encode(buf);
-    encoder.send("cat", b"chashu", None).await?;
-    let mut buf = encoder.into_writer();
-    buf.set_position(0);
+    let (sender, encoder) = encode();
+    task::spawn(async move {
+        sender.send("cat", b"chashu", None).await;
+    });
 
-    let mut reader = decode(buf);
+    let mut reader = decode(BufReader::new(encoder));
     let event = reader.next().await.unwrap()?;
     assert_message(&event, "cat", "chashu", None);
     Ok(())
 }
 
 #[async_std::test]
 async fn encode_message_with_id() -> http_types::Result<()> {
-    let buf = Cursor::new(vec![]);
-    let mut encoder = encode(buf);
-    encoder.send("cat", b"chashu", Some("0")).await?;
-    let mut buf = encoder.into_writer();
-    buf.set_position(0);
+    let (sender, encoder) = encode();
+    task::spawn(async move {
+        sender.send("cat", b"chashu", Some("0")).await;
+    });
 
-    let mut reader = decode(buf);
+    let mut reader = decode(BufReader::new(encoder));
     let event = reader.next().await.unwrap()?;
     assert_message(&event, "cat", "chashu", Some("0"));
     Ok(())
 }
 
 #[async_std::test]
 async fn encode_retry() -> http_types::Result<()> {
-    let buf = Cursor::new(vec![]);
-    let mut encoder = encode(buf);
-    encoder.send_retry(Duration::from_secs(12), None).await?;
-    let mut buf = encoder.into_writer();
-    buf.set_position(0);
+    let (sender, encoder) = encode();
+    task::spawn(async move {
+        let dur = Duration::from_secs(12);
+        sender.send_retry(dur, None).await;
+    });
 
-    let mut reader = decode(buf);
+    let mut reader = decode(BufReader::new(encoder));
     let event = reader.next().await.unwrap()?;
     assert_retry(&event, 12);
     Ok(())