Add StopSource and StopToken cancellation types

Author: rusty

Cargo.toml

@@ -38,6 +38,7 @@ async-session = { version = "3.0", optional = true }
 async-sse = "4.0.1"
 async-std = { version = "1.6.5", features = ["unstable"] }
 async-trait = "0.1.41"
+event-listener = "2.5.1"
 femme = { version = "2.1.1", optional = true }
 futures-util = "0.3.6"
 http-client = { version = "6.1.0", default-features = false }

src/cancellation.rs

@@ -0,0 +1,218 @@
+use std::pin::Pin;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use async_std::future::Future;
+use async_std::stream::Stream;
+use async_std::sync::Arc;
+use async_std::task::{Context, Poll};
+
+use event_listener::{Event, EventListener};
+use pin_project_lite::pin_project;
+
+#[derive(Debug)]
+pub struct StopSource {
+    stopped: Arc<AtomicBool>,
+    event: Arc<Event>,
+}
+
+impl StopSource {
+    pub fn new() -> Self {
+        Self {
+            stopped: Arc::new(AtomicBool::new(false)),
+            event: Arc::new(Event::new()),
+        }
+    }
+
+    pub fn token(&self) -> StopToken {
+        StopToken {
+            stopped: self.stopped.clone(),
+            event_listener: self.event.listen(),
+            event: self.event.clone(),
+        }
+    }
+}
+
+impl Drop for StopSource {
+    fn drop(&mut self) {
+        self.stopped.store(true, Ordering::SeqCst);
+        self.event.notify(usize::MAX);
+    }
+}
+
+pin_project! {
+    #[derive(Debug)]
+    pub struct StopToken {
+        #[pin]
+        stopped: Arc<AtomicBool>,
+        #[pin]
+        event_listener: EventListener,
+        event: Arc<Event>,
+    }
+}
+
+impl StopToken {
+    pub fn never() -> Self {
+        let event = Event::new();
+        Self {
+            stopped: Arc::new(AtomicBool::new(false)),
+            event_listener: event.listen(),
+            event: Arc::new(event),
+        }
+    }
+}
+
+impl Clone for StopToken {
+    fn clone(&self) -> Self {
+        Self {
+            stopped: self.stopped.clone(),
+            event_listener: self.event.listen(),
+            event: self.event.clone(),
+        }
+    }
+}
+
+impl Future for StopToken {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let mut this = self.project();
+        let _ = Future::poll(Pin::new(&mut this.event_listener), cx);
+        if this.stopped.load(Ordering::Relaxed) {
+            Poll::Ready(())
+        } else {
+            Poll::Pending
+        }
+    }
+}
+
+pin_project! {
+    #[derive(Debug)]
+    pub struct StopStream<S> {
+        #[pin]
+        stream: S,
+        #[pin]
+        stop_token: StopToken,
+    }
+}
+
+impl<S> StopStream<S> {
+    pub fn new(stream: S, stop_token: StopToken) -> Self {
+        Self { stream, stop_token }
+    }
+}
+
+impl<S> Stream for StopStream<S>
+where
+    S: Stream,
+{
+    type Item = S::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let mut this = self.project();
+        if Future::poll(Pin::new(&mut this.stop_token), cx).is_ready() {
+            Poll::Ready(None)
+        } else {
+            this.stream.poll_next(cx)
+        }
+    }
+}
+
+pub trait StopStreamExt: Sized {
+    fn stop_on(self, stop_token: StopToken) -> StopStream<Self> {
+        StopStream::new(self, stop_token)
+    }
+}
+
+impl<S> StopStreamExt for S where S: Stream {}
+
+pin_project! {
+    #[derive(Debug)]
+    pub struct StopFuture<F> {
+        #[pin]
+        future: F,
+        #[pin]
+        stop_token: StopToken,
+    }
+}
+
+impl<F> StopFuture<F> {
+    pub fn new(future: F, stop_token: StopToken) -> Self {
+        Self { future, stop_token }
+    }
+}
+
+impl<F> Future for StopFuture<F>
+where
+    F: Future,
+{
+    type Output = Option<F::Output>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let mut this = self.project();
+        if Future::poll(Pin::new(&mut this.stop_token), cx).is_ready() {
+            Poll::Ready(None)
+        } else {
+            match this.future.poll(cx) {
+                Poll::Pending => Poll::Pending,
+                Poll::Ready(output) => Poll::Ready(Some(output)),
+            }
+        }
+    }
+}
+
+pub trait StopFutureExt: Sized {
+    fn stop_on(self, stop_token: StopToken) -> StopFuture<Self> {
+        StopFuture::new(self, stop_token)
+    }
+}
+
+impl<F> StopFutureExt for F where F: Future {}
+
+#[cfg(test)]
+mod tests {
+    use std::thread;
+    use std::time::Duration;
+
+    use async_std::prelude::{FutureExt, StreamExt};
+
+    use super::*;
+
+    #[test]
+    fn test_cancellation() {
+        let source = StopSource::new();
+        let stop_token = source.token();
+
+        let pending_stream1 = async_std::stream::pending::<()>();
+        let pending_stream2 = async_std::stream::pending::<()>();
+        let pending_future1 = async_std::future::pending::<()>();
+        let pending_future2 = async_std::future::pending::<()>();
+        let wrapped_stream1 = pending_stream1.stop_on(stop_token.clone());
+        let wrapped_stream2 = pending_stream2.stop_on(stop_token.clone());
+        let wrapped_future1 = pending_future1.stop_on(stop_token.clone());
+        let wrapped_future2 = pending_future2.stop_on(stop_token);
+
+        let join_future = wrapped_stream1
+            .last()
+            .join(wrapped_stream2.last())
+            .join(wrapped_future1)
+            .join(wrapped_future2);
+
+        thread::spawn(move || {
+            let source = source;
+            thread::sleep(Duration::from_secs(1));
+            drop(source);
+        });
+
+        let res = async_std::task::block_on(join_future);
+        assert_eq!(res, (((None, None), None), None));
+    }
+
+    #[test]
+    fn test_never() {
+        let pending_future = async_std::future::pending::<()>();
+        let wrapped_future = pending_future.stop_on(StopToken::never());
+
+        let res = async_std::task::block_on(wrapped_future.timeout(Duration::from_secs(1)));
+        assert!(res.is_err());
+    }
+}

src/lib.rs

@@ -76,6 +76,7 @@ mod route;
 mod router;
 mod server;
 
+pub mod cancellation;
 pub mod convert;
 pub mod listener;
 pub mod log;