wait_token/

lib.rs

use futures::{FutureExt, StreamExt};
use std::{
    future::Future,
    ops::Deref,
    pin::Pin,
    sync::{
        Arc,
        atomic::{self, AtomicBool},
    },
    time::Duration,
};

#[derive(Debug, Clone)]
pub struct WaitTokenRepeat<Tk> {
    tk: Tk,
    frq: Duration,
}
impl WaitTokenRepeat<WaitToken> {
    pub async fn next(&mut self) -> bool {
        let result = tokio::select! {
            _ = tokio::time::sleep(self.frq) =>  {true}
            _ = self.tk.cancelled() =>  {false}
        };
        result
    }
}

#[derive(Debug, Default)]
pub struct WaitTokenGuard {
    tk: WaitToken,
}
impl Drop for WaitTokenGuard {
    fn drop(&mut self) {
        self.tk.cancel();
    }
}
impl Deref for WaitTokenGuard {
    type Target = WaitToken;

    fn deref(&self) -> &Self::Target {
        &self.tk
    }
}
impl WaitTokenGuard {
    pub fn new() -> Self {
        Self {
            tk: WaitToken::new(),
        }
    }
    pub fn clone_tk(&self) -> WaitToken {
        self.deref().clone()
    }
}

/// WaitToken (aka CancellationToken) is an object used for 'cancellation' events.
///
/// Example:
/// ```rust
/// // Create a new token
/// let cancellation = WaitToken::new();
///
/// // Clone token to use it in another thread. You can call cancel on this object too!
/// let cancellation_clone = cancellation.clone();
/// tokio::spawn(async move {
///     // wait for token to cancel
///     cancellation_clone.await;
///     println!("Finished");
/// })
///
/// // Child tokens will cancel when parent is cancelled.
/// // 'cancel()' on child tokens will not cancel parent
/// let child = cancellation.make_child_token();
/// tokio::spawn(async move {
///     // Wait for token to cancel
///     child.await;
///     println!("Child is finished");
/// })
///
/// println!("Cancelled");
/// cancellation.cancel();
///
/// // This await is just for demonstration. No need to await after cancel
/// // Note that 'await' takes by-value so you may need to clone
/// cancellation.clone().await;
///
/// // Undo cancellation of token
/// cancellation.reset();
///
/// // This await will not return because token is not cancelled anymore
/// cancellation.clone().await;
///
/// // This is unreachable
/// println!("Something went wrong!");
/// ```
#[derive(Debug, Clone)]
pub struct WaitToken {
    inner: Arc<WaitTokenRaw>,
}

pub type CancellationToken = WaitToken;

#[derive(Debug)]
struct WaitTokenRaw {
    state: AtomicBool,
    parent: Option<WaitToken>,
    future: Option<tokio::sync::Notify>,
}

impl WaitTokenRaw {
    pub fn is_cancelled(&self) -> bool {
        if let Some(parent) = &self.parent {
            if parent.is_cancelled() {
                return true;
            }
        }
        self.state.load(atomic::Ordering::Acquire)
    }

    fn notified<'a: 'b, 'b>(&'a self) -> tokio::sync::futures::Notified<'b> {
        if let Some(parent) = &self.parent {
            return parent.inner.notified();
        }
        self.future.as_ref().unwrap().notified()
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CancelState {
    AlreadyCancelled,
    Cancelled,
}
impl CancelState {
    pub fn already_cancelled(self) -> bool {
        self == CancelState::AlreadyCancelled
    }
    pub fn just_cancelled(self) -> bool {
        self == CancelState::Cancelled
    }
}

impl WaitToken {
    /// Create a new token
    pub fn new() -> Self {
        Self {
            inner: Arc::new(WaitTokenRaw {
                state: false.into(),
                parent: None,
                future: Some(tokio::sync::Notify::new()),
            }),
        }
    }

    /// Spawns background fetch to listem SIGTERM & SIGINT events and calls cancel
    pub fn spawn_terminator(&self) -> std::thread::JoinHandle<Result<(), std::io::Error>> {
        let this = self.clone();
        std::thread::spawn(move || {
            let mut signals = signal_hook::iterator::Signals::new([
                signal_hook::consts::SIGTERM,
                signal_hook::consts::SIGINT,
            ])
            .inspect_err(|_| {
                this.cancel();
            })?;

            for _ in &mut signals {
                this.cancel();
            }

            Ok(())
        })
    }

    /// Create an already cancelled token
    pub fn ready() -> Self {
        Self {
            inner: Arc::new(WaitTokenRaw {
                state: true.into(),
                parent: None,
                future: Some(tokio::sync::Notify::new()),
            }),
        }
    }

    fn wake(&self) {
        if let Some(parent) = &self.inner.parent {
            parent.wake();
            return;
        }
        if let Some(future) = &self.inner.future {
            future.notify_waiters();
        }
    }

    fn notified(&self) -> tokio::sync::futures::Notified<'_> {
        self.inner.notified()
    }

    /// Set state to 'cancelled'. Futures waiting for token will complete and streams terminated.
    pub fn cancel(&self) -> CancelState {
        if self.inner.state.swap(true, atomic::Ordering::Release) {
            return CancelState::AlreadyCancelled;
        }
        self.wake();
        CancelState::Cancelled
    }

    /// Set state to 'not cancelled'
    pub fn reset(&self) {
        self.inner.state.store(false, atomic::Ordering::Release);
    }

    pub fn is_cancelled(&self) -> bool {
        self.inner.is_cancelled()
    }

    pub fn cancelled(&self) -> WaitTokenCancelled {
        WaitTokenCancelled {
            token: self.clone(),
            future: Box::pin(unsafe {
                std::mem::transmute::<
                    tokio::sync::futures::Notified<'_>,
                    tokio::sync::futures::Notified<'_>,
                >(self.notified())
            }),
        }
    }

    /// Make a stream that will emit items until token is cancelled
    pub fn until_cancel(
        &self,
    ) -> futures::stream::TakeUntil<futures::stream::Repeat<()>, WaitTokenCancelled> {
        futures::stream::repeat(()).take_until(self.cancelled())
    }

    pub fn on_cancel(&self) -> futures::stream::Once<WaitTokenCancelled> {
        futures::stream::once(self.cancelled())
    }
    pub async fn on_cancel_then<Fut: Future + Send>(self, fut: Fut) -> Fut::Output {
        self.cancelled().then(|_| fut).await
    }
    /// Make an object that will wait for delay every iteration and return false if token is cancelled.
    /// Common usage is with while loop.
    ///
    /// Example:
    /// ```rust
    /// let mut repeater = cancellation.repeat_until_cancel(Duration::from_secs(1));
    /// while repeater.next().await {
    ///     do_something();
    /// }
    /// ```
    pub fn repeat_until_cancel(&self, frq: Duration) -> WaitTokenRepeat<Self> {
        WaitTokenRepeat {
            tk: self.clone(),
            frq,
        }
    }

    /// Run future. Future will be aborted when token is cancelled
    pub async fn run_fn<Fut: Future>(&self, fut: Fut) -> Option<Fut::Output> {
        tokio::select! {
            a = fut => Some(a),
            _ = self.cancelled() => None,
        }
    }

    /// Create a sub-token that will cancel if current token cancelled. <br/>
    /// Cancel on child token will not cancel parent
    pub fn make_child_token(&self) -> Self {
        Self {
            inner: Arc::new(WaitTokenRaw {
                state: false.into(),
                parent: Some(self.clone()),
                future: None,
            }),
        }
    }

    /// Create a guard that will cancel wait token on drop
    pub fn guard(&self) -> WaitTokenGuard {
        WaitTokenGuard { tk: self.clone() }
    }
}

impl Default for WaitToken {
    fn default() -> Self {
        Self::new()
    }
}

pub struct WaitTokenCancelled {
    token: WaitToken,
    future: Pin<Box<tokio::sync::futures::Notified<'static>>>,
}

impl Future for WaitTokenCancelled {
    type Output = ();

    fn poll(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Self::Output> {
        loop {
            if self.token.is_cancelled() {
                return std::task::Poll::Ready(());
            }

            if self.future.poll_unpin(cx).is_pending() {
                return std::task::Poll::Pending;
            }

            if self.token.is_cancelled() {
                return std::task::Poll::Ready(());
            }

            let future = unsafe {
                std::mem::transmute::<
                    tokio::sync::futures::Notified<'_>,
                    tokio::sync::futures::Notified<'_>,
                >(self.token.notified())
            };
            self.future.set(future);
        }
    }
}