Notify the Sender that the Encoder has been dropped

Author: jbr

Cargo.toml

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

src/encoder.rs

@@ -6,15 +6,26 @@ use async_std::io::Read as AsyncRead;
 use async_std::prelude::*;
 use async_std::task::{ready, Context, Poll};
 use std::pin::Pin;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
 
-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,
+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>,
+}
+
+#[pinned_drop]
+impl PinnedDrop for Encoder {
+    fn drop(self: Pin<&mut Self>) {
+        self.disconnected.store(true, Ordering::Relaxed);
     }
 }
 
@@ -79,53 +90,80 @@ impl AsyncRead for Encoder {
 // }
 
 /// The sending side of the encoder.
-#[derive(Debug)]
-pub struct Sender(sync::Sender<Vec<u8>>);
+#[derive(Debug, Clone)]
+pub struct Sender {
+    sender: sync::Sender<Vec<u8>>,
+    disconnected: Arc<std::sync::atomic::AtomicBool>,
+}
 
 /// Create a new SSE encoder.
 pub fn encode() -> (Sender, Encoder) {
     let (sender, receiver) = sync::channel(1);
+    let disconnected = Arc::new(AtomicBool::new(false));
+
     let encoder = Encoder {
         receiver,
         buf: None,
         cursor: 0,
+        disconnected: disconnected.clone(),
+    };
+
+    let sender = Sender {
+        sender,
+        disconnected,
     };
-    (Sender(sender), encoder)
+
+    (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")
+    }
+}
+
+#[must_use]
 impl Sender {
     /// Send a new message over SSE.
-    pub async fn send(&self, name: &str, data: &str, id: Option<&str>) {
+    pub async fn send(
+        &self,
+        name: &str,
+        data: &str,
+        id: Option<&str>,
+    ) -> Result<(), DisconnectedError> {
+        if self.disconnected.load(Ordering::Relaxed) {
+            return Err(DisconnectedError);
+        }
+
         // Write the event name
         let msg = format!("event:{}\n", name);
-        self.0.send(msg.into_bytes()).await;
+        self.sender.send(msg.into_bytes()).await;
 
         // Write the id
         if let Some(id) = id {
-            self.0.send(format!("id:{}\n", id).into_bytes()).await;
+            self.sender.send(format!("id:{}\n", id).into_bytes()).await;
         }
 
         // Write the data section, and end.
         let msg = format!("data:{}\n\n", data);
-        self.0.send(msg.into_bytes()).await;
+        self.sender.send(msg.into_bytes()).await;
+        Ok(())
     }
 
     /// Send a new "retry" message over SSE.
     pub async fn send_retry(&self, dur: Duration, id: Option<&str>) {
         // Write the id
         if let Some(id) = id {
-            self.0.send(format!("id:{}\n", id).into_bytes()).await;
+            self.sender.send(format!("id:{}\n", id).into_bytes()).await;
         }
 
         // 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;
-    }
-}
-
-impl Clone for Sender {
-    fn clone(&self) -> Self {
-        Self(self.0.clone())
+        self.sender.send(msg.into_bytes()).await;
     }
 }

src/lib.rs

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

src/lines.rs

@@ -2,30 +2,29 @@ use std::mem;
 use std::pin::Pin;
 use std::str;
 
-use pin_project_lite::pin_project;
+use pin_project::pin_project;
 
 use async_std::io::{self, BufRead};
 use async_std::stream::Stream;
 use async_std::task::{ready, Context, Poll};
 
-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,
-    }
+/// 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,
 }
 
 impl<R> Lines<R> {

tests/encode.rs

@@ -29,9 +29,7 @@ fn assert_retry(event: &Event, dur: u64) {
 #[async_std::test]
 async fn encode_message() -> http_types::Result<()> {
     let (sender, encoder) = encode();
-    task::spawn(async move {
-        sender.send("cat", "chashu", None).await;
-    });
+    task::spawn(async move { sender.send("cat", "chashu", None).await });
 
     let mut reader = decode(BufReader::new(encoder));
     let event = reader.next().await.unwrap()?;
@@ -42,9 +40,7 @@ async fn encode_message() -> http_types::Result<()> {
 #[async_std::test]
 async fn encode_message_with_id() -> http_types::Result<()> {
     let (sender, encoder) = encode();
-    task::spawn(async move {
-        sender.send("cat", "chashu", Some("0")).await;
-    });
+    task::spawn(async move { sender.send("cat", "chashu", Some("0")).await });
 
     let mut reader = decode(BufReader::new(encoder));
     let event = reader.next().await.unwrap()?;
@@ -65,3 +61,20 @@ async fn encode_retry() -> http_types::Result<()> {
     assert_retry(&event, 12);
     Ok(())
 }
+
+#[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() });
+
+    //move the encoder into Lines, which gets dropped after this
+    assert_eq!(reader.lines().next().await.unwrap().unwrap(), "event:cat");
+
+    assert_eq!(
+        sender.send("cat", "nori", None).await,
+        Err(async_sse::DisconnectedError)
+    );
+    Ok(())
+}