abortable.rs

use crate::task::AtomicWaker;
use alloc::sync::Arc;
use core::fmt;
use core::pin::Pin;
use core::sync::atomic::{AtomicBool, Ordering};
use futures_core::future::Future;
use futures_core::task::{Context, Poll};
use futures_core::Stream;
use pin_project_lite::pin_project;

pin_project! {
    /// A future/stream which can be remotely short-circuited using an `AbortHandle`.
    #[derive(Debug, Clone)]
    #[must_use = "futures/streams do nothing unless you poll them"]
    pub struct Abortable<T> {
        #[pin]
        task: T,
        inner: Arc<AbortInner>,
    }
}

impl<T> Abortable<T> {
    /// Creates a new `Abortable` future/stream using an existing `AbortRegistration`.
    /// `AbortRegistration`s can be acquired through `AbortHandle::new`.
    ///
    /// When `abort` is called on the handle tied to `reg` or if `abort` has
    /// already been called, the future/stream will complete immediately without making
    /// any further progress.
    ///
    /// # Examples:
    ///
    /// Usage with futures:
    ///
    /// ```
    /// # futures::executor::block_on(async {
    /// use futures::future::{Abortable, AbortHandle, Aborted};
    ///
    /// let (abort_handle, abort_registration) = AbortHandle::new_pair();
    /// let future = Abortable::new(async { 2 }, abort_registration);
    /// abort_handle.abort();
    /// assert_eq!(future.await, Err(Aborted));
    /// # });
    /// ```
    ///
    /// Usage with streams:
    ///
    /// ```
    /// # futures::executor::block_on(async {
    /// # use futures::future::{Abortable, AbortHandle};
    /// # use futures::stream::{self, StreamExt};
    ///
    /// let (abort_handle, abort_registration) = AbortHandle::new_pair();
    /// let mut stream = Abortable::stream(stream::iter(vec![1, 2, 3]), abort_registration);
    /// abort_handle.abort();
    /// assert_eq!(stream.next().await, None);
    /// # });
    /// ```
    pub fn new(future: T, reg: AbortRegistration) -> Self {
        Self { task: future, inner: reg.inner }
    }
}

/// A registration handle for an `Abortable` task.
/// Values of this type can be acquired from `AbortHandle::new` and are used
/// in calls to `Abortable::new`.
#[derive(Debug)]
pub struct AbortRegistration {
    inner: Arc<AbortInner>,
}

/// A handle to an `Abortable` task.
#[derive(Debug, Clone)]
pub struct AbortHandle {
    inner: Arc<AbortInner>,
}

impl AbortHandle {
    /// Creates an (`AbortHandle`, `AbortRegistration`) pair which can be used
    /// to abort a running future or stream.
    ///
    /// This function is usually paired with a call to [`Abortable::new`].
    pub fn new_pair() -> (Self, AbortRegistration) {
        let inner =
            Arc::new(AbortInner { waker: AtomicWaker::new(), cancel: AtomicBool::new(false) });

        (Self { inner: inner.clone() }, AbortRegistration { inner })
    }
}

// Inner type storing the waker to awaken and a bool indicating that it
// should be cancelled.
#[derive(Debug)]
struct AbortInner {
    waker: AtomicWaker,
    cancel: AtomicBool,
}

/// Indicator that the `Abortable` task was aborted.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct Aborted;

impl fmt::Display for Aborted {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "`Abortable` future has been aborted")
    }
}

#[cfg(feature = "std")]
impl std::error::Error for Aborted {}

impl<T> Abortable<T> {
    fn try_poll<I>(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        poll: impl Fn(Pin<&mut T>, &mut Context<'_>) -> Poll<I>,
    ) -> Poll<Result<I, Aborted>> {
        // Check if the task has been aborted
        if self.inner.cancel.load(Ordering::Relaxed) {
            return Poll::Ready(Err(Aborted));
        }

        // attempt to complete the task
        if let Poll::Ready(x) = poll(self.as_mut().project().task, cx) {
            return Poll::Ready(Ok(x));
        }

        // Register to receive a wakeup if the task is aborted in the future
        self.inner.waker.register(cx.waker());

        // Check to see if the task was aborted between the first check and
        // registration.
        // Checking with `Relaxed` is sufficient because `register` introduces an
        // `AcqRel` barrier.
        if self.inner.cancel.load(Ordering::Relaxed) {
            return Poll::Ready(Err(Aborted));
        }

        Poll::Pending
    }
}

impl<Fut> Future for Abortable<Fut>
where
    Fut: Future,
{
    type Output = Result<Fut::Output, Aborted>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.try_poll(cx, |fut, cx| fut.poll(cx))
    }
}

impl<St> Stream for Abortable<St>
where
    St: Stream,
{
    type Item = St::Item;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.try_poll(cx, |stream, cx| stream.poll_next(cx)).map(Result::ok).map(Option::flatten)
    }
}

impl AbortHandle {
    /// Abort the `Abortable` stream/future associated with this handle.
    ///
    /// Notifies the Abortable task associated with this handle that it
    /// should abort. Note that if the task is currently being polled on
    /// another thread, it will not immediately stop running. Instead, it will
    /// continue to run until its poll method returns.
    pub fn abort(&self) {
        self.inner.cancel.store(true, Ordering::Relaxed);
        self.inner.waker.wake();
    }
}