halo-agent 0.1.1

Halo Agent and Workflow execution kernel
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

use std::{future::Future, task::Poll};

use futures_util::future::poll_fn;

use crate::{Cancellation, Error};

pub(super) fn check(cancellation: &Cancellation) -> Result<(), Error> {
    if cancellation.is_cancelled() {
        Err(Error::new("CANCELLED", "operation cancelled"))
    } else {
        Ok(())
    }
}

pub(crate) async fn cancellable<T, E: From<Error>>(
    future: impl Future<Output = Result<T, E>>,
    cancellation: &Cancellation,
) -> Result<T, E> {
    let mut future = std::pin::pin!(future);
    poll_fn(|cx| {
        cancellation.register(cx.waker());
        if let Err(error) = check(cancellation) {
            return Poll::Ready(Err(error.into()));
        }
        let result = future.as_mut().poll(cx);
        if let Err(error) = check(cancellation) {
            return Poll::Ready(Err(error.into()));
        }
        result
    })
    .await
}