async-utils/src/lib.rs
86 lines
use std::future::Future;use tokio_util::sync::CancellationToken;#[derive(Debug, PartialEq, Eq)]pub enum CancelErr { Cancelled,}pub trait OrCancelExt: Sized { type Output; fn or_cancel( self, token: &CancellationToken, ) -> impl Future<Output = Result<Self::Output, CancelErr>> + Send;}impl<F> OrCancelExt for Fwhere F: Future + Send, F::Output: Send,{ type Output = F::Output; async fn or_cancel(self, token: &CancellationToken) -> Result<Self::Output, CancelErr> { tokio::select! { _ = token.cancelled() => Err(CancelErr::Cancelled), res = self => Ok(res), } }}#[cfg(test)]mod tests { use super::*; use pretty_assertions::assert_eq; use std::time::Duration; use tokio::task; use tokio::time::sleep; #[tokio::test] async fn returns_ok_when_future_completes_first() { let token = CancellationToken::new(); let value = async { 42 }; let result = value.or_cancel(&token).await; assert_eq!(Ok(42), result); } #[tokio::test] async fn returns_err_when_token_cancelled_first() { let token = CancellationToken::new(); let token_clone = token.clone(); let cancel_handle = task::spawn(async move { sleep(Duration::from_millis(10)).await; token_clone.cancel(); }); let result = async { sleep(Duration::from_millis(100)).await; 7 } .or_cancel(&token) .await; cancel_handle.await.expect("cancel task panicked"); assert_eq!(Err(CancelErr::Cancelled), result); } #[tokio::test] async fn returns_err_when_token_already_cancelled() { let token = CancellationToken::new(); token.cancel(); let result = async { sleep(Duration::from_millis(50)).await; 5 } .or_cancel(&token) .await; assert_eq!(Err(CancelErr::Cancelled), result); }}