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
// Released under the MIT License.
// Copyright, 2026, by Samuel Williams.
use crate::Cancellation;
use std::future::{Future, poll_fn};
use std::pin::pin;
/// Notify protected work of cancellation and continue awaiting its completion.
///
/// When cancellation is observed, call `on_cancel` exactly once before polling the protected future again. This includes an already-cancelled signal on the first poll. If completion and cancellation are both ready in that poll, the callback runs before returning the future's normal output. A callback panic propagates to the caller.
///
/// The callback is synchronous: request graceful shutdown and wake the work; keep asynchronous draining in the protected future. Its operations must use cancellation inputs which allow draining. This wrapper does not mask signals, prevent its own destruction, or intercept scheduler task cancellation.
///
/// ```
/// use base_scheduler::{Cancellation, defer_cancel};
/// # async fn example() {
///
/// let shutdown = Cancellation::new();
/// let drain = Cancellation::new();
/// shutdown.cancel();
/// let output = defer_cancel(&shutdown, async {
/// drain.cancelled().await;
/// 42
/// }, || { drain.cancel(); }).await;
/// assert_eq!(output, 42);
/// # }
/// ```
pub async fn defer_cancel<FutureType, OnCancel>(
cancellation: &Cancellation,
future: FutureType,
on_cancel: OnCancel,
) -> FutureType::Output
where
FutureType: Future,
OnCancel: FnOnce(),
{
let mut notification = pin!(cancellation.cancelled());
let mut future = pin!(future);
let mut on_cancel = Some(on_cancel);
poll_fn(move |context| {
if on_cancel.is_some() && notification.as_mut().poll(context).is_ready() {
// Remove the callback before invoking user code, including on panic.
on_cancel.take().expect("cancellation callback is present")();
}
future.as_mut().poll(context)
})
.await
}