File size: 2,149 Bytes
9ebf6d4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | use std::future::Future;
use tokio_util::sync::CancellationToken;
/// Stack budget for threads that poll Codex async work.
pub const THREAD_STACK_SIZE_BYTES: usize = 16 * 1024 * 1024;
#[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 F
where
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);
}
}
|