Merge pull request #6 from jbr/use-async-channel

use async-channel to notify senders of disconnection

Author: yoshuawuyts

Cargo.toml

@@ -21,7 +21,8 @@ async-std = { version = "1.6.0", features = ["unstable"] }
 http-types = "2.0.1"
 log = "0.4.8"
 memchr = "2.3.3"
-pin-project = "0.4.22"
+pin-project-lite = "0.1.4"
+async-channel = "1.1.1"
 
 [dev-dependencies]
 femme = "2.0.0"

src/encoder.rs

@@ -1,31 +1,19 @@
-use async_std::sync;
-use std::io;
-use std::time::Duration;
-
 use async_std::io::Read as AsyncRead;
 use async_std::prelude::*;
 use async_std::task::{ready, Context, Poll};
+
+use std::io;
 use std::pin::Pin;
-use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::Arc;
-
-use pin_project::{pin_project, pinned_drop};
-
-#[pin_project(PinnedDrop)]
-/// An SSE protocol encoder.
-#[derive(Debug)]
-pub struct Encoder {
-    buf: Option<Vec<u8>>,
-    #[pin]
-    receiver: sync::Receiver<Vec<u8>>,
-    cursor: usize,
-    disconnected: Arc<AtomicBool>,
-}
+use std::time::Duration;
 
-#[pinned_drop]
-impl PinnedDrop for Encoder {
-    fn drop(self: Pin<&mut Self>) {
-        self.disconnected.store(true, Ordering::Relaxed);
+pin_project_lite::pin_project! {
+    /// An SSE protocol encoder.
+    #[derive(Debug)]
+    pub struct Encoder {
+        buf: Option<Vec<u8>>,
+        #[pin]
+        receiver: async_channel::Receiver<Vec<u8>>,
+        cursor: usize,
     }
 }
 
@@ -91,79 +79,56 @@ impl AsyncRead for Encoder {
 
 /// The sending side of the encoder.
 #[derive(Debug, Clone)]
-pub struct Sender {
-    sender: sync::Sender<Vec<u8>>,
-    disconnected: Arc<std::sync::atomic::AtomicBool>,
-}
+pub struct Sender(async_channel::Sender<Vec<u8>>);
 
 /// Create a new SSE encoder.
 pub fn encode() -> (Sender, Encoder) {
-    let (sender, receiver) = sync::channel(1);
-    let disconnected = Arc::new(AtomicBool::new(false));
-
+    let (sender, receiver) = async_channel::bounded(1);
     let encoder = Encoder {
         receiver,
         buf: None,
         cursor: 0,
-        disconnected: disconnected.clone(),
     };
-
-    let sender = Sender {
-        sender,
-        disconnected,
-    };
-
-    (sender, encoder)
+    (Sender(sender), encoder)
 }
 
-/// An error that represents that the [Encoder] has been dropped.
-#[derive(Debug, Eq, PartialEq)]
-pub struct DisconnectedError;
-impl std::error::Error for DisconnectedError {}
-impl std::fmt::Display for DisconnectedError {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "Disconnected")
+impl Sender {
+    async fn inner_send(&self, bytes: impl Into<Vec<u8>>) -> io::Result<()> {
+        self.0
+            .send(bytes.into())
+            .await
+            .map_err(|_| io::Error::new(io::ErrorKind::ConnectionAborted, "sse disconnected"))
     }
-}
 
-#[must_use]
-impl Sender {
     /// Send a new message over SSE.
-    pub async fn send(
-        &self,
-        name: &str,
-        data: &str,
-        id: Option<&str>,
-    ) -> Result<(), DisconnectedError> {
-        if self.disconnected.load(Ordering::Relaxed) {
-            return Err(DisconnectedError);
-        }
-
+    pub async fn send(&self, name: &str, data: &str, id: Option<&str>) -> io::Result<()> {
         // Write the event name
         let msg = format!("event:{}\n", name);
-        self.sender.send(msg.into_bytes()).await;
+        self.inner_send(msg).await?;
 
         // Write the id
         if let Some(id) = id {
-            self.sender.send(format!("id:{}\n", id).into_bytes()).await;
+            self.inner_send(format!("id:{}\n", id)).await?;
         }
 
         // Write the data section, and end.
         let msg = format!("data:{}\n\n", data);
-        self.sender.send(msg.into_bytes()).await;
+        self.inner_send(msg).await?;
+
         Ok(())
     }
 
     /// Send a new "retry" message over SSE.
-    pub async fn send_retry(&self, dur: Duration, id: Option<&str>) {
+    pub async fn send_retry(&self, dur: Duration, id: Option<&str>) -> io::Result<()> {
         // Write the id
         if let Some(id) = id {
-            self.sender.send(format!("id:{}\n", id).into_bytes()).await;
+            self.inner_send(format!("id:{}\n", id)).await?;
         }
 
         // Write the retry section, and end.
         let dur = dur.as_secs_f64() as u64;
         let msg = format!("retry:{}\n\n", dur);
-        self.sender.send(msg.into_bytes()).await;
+        self.inner_send(msg).await?;
+        Ok(())
     }
 }

src/lib.rs

@@ -43,7 +43,7 @@ mod lines;
 mod message;
 
 pub use decoder::{decode, Decoder};
-pub use encoder::{encode, DisconnectedError, Encoder, Sender};
+pub use encoder::{encode, Encoder, Sender};
 pub use event::Event;
 pub use handshake::upgrade;
 pub use message::Message;

src/lines.rs

@@ -2,29 +2,30 @@ use std::mem;
 use std::pin::Pin;
 use std::str;
 
-use pin_project::pin_project;
+use pin_project_lite::pin_project;
 
 use async_std::io::{self, BufRead};
 use async_std::stream::Stream;
 use async_std::task::{ready, Context, Poll};
 
-/// A stream of lines in a byte stream.
-///
-/// This stream is created by the [`lines`] method on types that implement [`BufRead`].
-///
-/// This type is an async version of [`std::io::Lines`].
-///
-/// [`lines`]: trait.BufRead.html#method.lines
-/// [`BufRead`]: trait.BufRead.html
-/// [`std::io::Lines`]: https://doc.rust-lang.org/std/io/struct.Lines.html
-#[pin_project]
-#[derive(Debug)]
-pub(crate) struct Lines<R> {
-    #[pin]
-    pub(crate) reader: R,
-    pub(crate) buf: String,
-    pub(crate) bytes: Vec<u8>,
-    pub(crate) read: usize,
+pin_project! {
+    /// A stream of lines in a byte stream.
+    ///
+    /// This stream is created by the [`lines`] method on types that implement [`BufRead`].
+    ///
+    /// This type is an async version of [`std::io::Lines`].
+    ///
+    /// [`lines`]: trait.BufRead.html#method.lines
+    /// [`BufRead`]: trait.BufRead.html
+    /// [`std::io::Lines`]: https://doc.rust-lang.org/std/io/struct.Lines.html
+    #[derive(Debug)]
+    pub(crate) struct Lines<R> {
+        #[pin]
+        pub(crate) reader: R,
+        pub(crate) buf: String,
+        pub(crate) bytes: Vec<u8>,
+        pub(crate) read: usize,
+    }
 }
 
 impl<R> Lines<R> {

tests/encode.rs

@@ -53,7 +53,7 @@ async fn encode_retry() -> http_types::Result<()> {
     let (sender, encoder) = encode();
     task::spawn(async move {
         let dur = Duration::from_secs(12);
-        sender.send_retry(dur, None).await;
+        sender.send_retry(dur, None).await.unwrap();
     });
 
     let mut reader = decode(BufReader::new(encoder));
@@ -65,16 +65,20 @@ async fn encode_retry() -> http_types::Result<()> {
 #[async_std::test]
 async fn dropping_encoder() -> http_types::Result<()> {
     let (sender, encoder) = encode();
-    let reader = BufReader::new(encoder);
     let sender_clone = sender.clone();
-    task::spawn(async move { sender_clone.send("cat", "chashu", Some("0")).await.unwrap() });
+    task::spawn(async move { sender_clone.send("cat", "chashu", None).await });
 
-    //move the encoder into Lines, which gets dropped after this
-    assert_eq!(reader.lines().next().await.unwrap().unwrap(), "event:cat");
+    let mut reader = decode(BufReader::new(encoder));
+    let event = reader.next().await.unwrap()?;
+    assert_message(&event, "cat", "chashu", None);
+
+    std::mem::drop(reader);
 
+    let response = sender.send("cat", "chashu", None).await;
+    assert!(response.is_err());
     assert_eq!(
-        sender.send("cat", "nori", None).await,
-        Err(async_sse::DisconnectedError)
+        response.unwrap_err().kind(),
+        async_std::io::ErrorKind::ConnectionAborted
     );
     Ok(())
 }