codex_async_utils/
lib.rs

1use std::future::Future;
2use tokio_util::sync::CancellationToken;
3
4#[derive(Debug, PartialEq, Eq)]
5pub enum CancelErr {
6    Cancelled,
7}
8
9pub trait OrCancelExt: Sized {
10    type Output;
11
12    fn or_cancel(
13        self,
14        token: &CancellationToken,
15    ) -> impl Future<Output = Result<Self::Output, CancelErr>> + Send;
16}
17
18impl<F> OrCancelExt for F
19where
20    F: Future + Send,
21    F::Output: Send,
22{
23    type Output = F::Output;
24
25    async fn or_cancel(self, token: &CancellationToken) -> Result<Self::Output, CancelErr> {
26        tokio::select! {
27            _ = token.cancelled() => Err(CancelErr::Cancelled),
28            res = self => Ok(res),
29        }
30    }
31}
32
33#[cfg(test)]
34mod tests {
35    use super::*;
36    use pretty_assertions::assert_eq;
37    use std::time::Duration;
38    use tokio::task;
39    use tokio::time::sleep;
40
41    #[tokio::test]
42    async fn returns_ok_when_future_completes_first() {
43        let token = CancellationToken::new();
44        let value = async { 42 };
45
46        let result = value.or_cancel(&token).await;
47
48        assert_eq!(Ok(42), result);
49    }
50
51    #[tokio::test]
52    async fn returns_err_when_token_cancelled_first() {
53        let token = CancellationToken::new();
54        let token_clone = token.clone();
55
56        let cancel_handle = task::spawn(async move {
57            sleep(Duration::from_millis(10)).await;
58            token_clone.cancel();
59        });
60
61        let result = async {
62            sleep(Duration::from_millis(100)).await;
63            7
64        }
65        .or_cancel(&token)
66        .await;
67
68        cancel_handle.await.expect("cancel task panicked");
69        assert_eq!(Err(CancelErr::Cancelled), result);
70    }
71
72    #[tokio::test]
73    async fn returns_err_when_token_already_cancelled() {
74        let token = CancellationToken::new();
75        token.cancel();
76
77        let result = async {
78            sleep(Duration::from_millis(50)).await;
79            5
80        }
81        .or_cancel(&token)
82        .await;
83
84        assert_eq!(Err(CancelErr::Cancelled), result);
85    }
86}
codex_async_utils/lib.rs

codex_async_utils/
lib.rs
