weaveffi/abi/
cancel.rs

1//! Reference-counted cancel tokens and the future combinator that makes
2//! cancellation of an async call observable.
3//!
4//! A consumer creates a token with `{prefix}_cancel_token_create` and passes
5//! it to a cancellable async launcher. The launcher takes its own reference
6//! (a [`CancelToken`]), so the consumer may cancel and destroy its reference
7//! at any time afterwards, even before the call completes. Cancelling wakes
8//! every future registered on the token; the launcher's [`Cancellable`]
9//! wrapper then drops the producer's future and the call completes with
10//! [`CANCELLED_ERROR_CODE`](crate::abi::CANCELLED_ERROR_CODE).
11//!
12//! Each in-flight call registers once, when it's first polled, and keeps its
13//! current waker in an [`AtomicWaker`] of its own, so a poll refreshes the
14//! waker without taking a lock.
15
16use std::cell::UnsafeCell;
17use std::future::Future;
18use std::pin::Pin;
19use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
20use std::sync::{Arc, Mutex, PoisonError};
21use std::task::{Context, Poll, Waker};
22
23/// A waker slot that one task registers into and any thread wakes, without a
24/// lock (the algorithm of `futures`' `AtomicWaker`).
25#[derive(Default)]
26pub struct AtomicWaker {
27    state: AtomicUsize,
28    waker: UnsafeCell<Option<Waker>>,
29}
30
31const WAITING: usize = 0;
32const REGISTERING: usize = 0b01;
33const WAKING: usize = 0b10;
34
35// SAFETY: the `state` protocol gives exactly one thread at a time access to
36// `waker` (the registering thread holds REGISTERING, the waking thread holds
37// WAKING), and `Waker` itself is `Send + Sync`.
38unsafe impl Send for AtomicWaker {}
39// SAFETY: see the `Send` impl.
40unsafe impl Sync for AtomicWaker {}
41
42impl std::fmt::Debug for AtomicWaker {
43    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        f.debug_struct("AtomicWaker").finish_non_exhaustive()
45    }
46}
47
48impl AtomicWaker {
49    /// Store `waker` as the one to wake, replacing any earlier one. A wake
50    /// that races the registration wakes `waker` immediately.
51    pub fn register(&self, waker: &Waker) {
52        match self
53            .state
54            .compare_exchange(WAITING, REGISTERING, Ordering::Acquire, Ordering::Acquire)
55            .unwrap_or_else(|actual| actual)
56        {
57            WAITING => {
58                // SAFETY: holding REGISTERING gives exclusive access.
59                let slot = unsafe { &mut *self.waker.get() };
60                match slot {
61                    Some(old) if old.will_wake(waker) => {}
62                    _ => *slot = Some(waker.clone()),
63                }
64                if self
65                    .state
66                    .compare_exchange(REGISTERING, WAITING, Ordering::AcqRel, Ordering::Acquire)
67                    .is_err()
68                {
69                    // A wake arrived while registering (state is now
70                    // REGISTERING | WAKING): deliver it here.
71                    // SAFETY: the waker defers to us while REGISTERING is set.
72                    let woken = unsafe { (*self.waker.get()).take() };
73                    self.state.swap(WAITING, Ordering::AcqRel);
74                    if let Some(w) = woken {
75                        w.wake();
76                    }
77                }
78            }
79            // A concurrent wake is in progress: make sure this task polls
80            // again.
81            WAKING => waker.wake_by_ref(),
82            // Registered concurrently from another poll of the same task;
83            // nothing to do.
84            _ => {}
85        }
86    }
87
88    /// Wake the registered waker, if any.
89    pub fn wake(&self) {
90        if let WAITING = self.state.fetch_or(WAKING, Ordering::AcqRel) {
91            // SAFETY: holding WAKING (without REGISTERING) gives exclusive
92            // access.
93            let woken = unsafe { (*self.waker.get()).take() };
94            self.state.fetch_and(!WAKING, Ordering::Release);
95            if let Some(w) = woken {
96                w.wake();
97            }
98        }
99    }
100}
101
102/// The shared state behind a cancel token (`{prefix}_cancel_token` in C, an
103/// opaque type).
104///
105/// A `*mut FfiCancelToken` handed to C is an [`Arc`] turned into a raw
106/// pointer, so the consumer's reference and every [`CancelToken`] the
107/// producer holds keep the state alive independently.
108#[derive(Debug, Default)]
109pub struct FfiCancelToken {
110    cancelled: AtomicBool,
111    /// One slot per in-flight call; touched once when a call starts and once
112    /// when it ends, never per poll.
113    registrations: Mutex<Vec<Arc<AtomicWaker>>>,
114}
115
116impl FfiCancelToken {
117    fn cancel(&self) {
118        self.cancelled.store(true, Ordering::SeqCst);
119        let registrations = std::mem::take(&mut *self.lock());
120        for r in registrations {
121            r.wake();
122        }
123    }
124
125    fn is_cancelled(&self) -> bool {
126        self.cancelled.load(Ordering::SeqCst)
127    }
128
129    fn lock(&self) -> std::sync::MutexGuard<'_, Vec<Arc<AtomicWaker>>> {
130        self.registrations
131            .lock()
132            .unwrap_or_else(PoisonError::into_inner)
133    }
134
135    fn register(&self, slot: &Arc<AtomicWaker>) {
136        self.lock().push(Arc::clone(slot));
137    }
138
139    fn deregister(&self, slot: &Arc<AtomicWaker>) {
140        self.lock().retain(|r| !Arc::ptr_eq(r, slot));
141    }
142}
143
144/// The body of `{prefix}_cancel_token_create`: a fresh, uncancelled token
145/// holding one reference, which the consumer releases with
146/// [`cancel_token_destroy`].
147#[must_use]
148pub fn cancel_token_create() -> *mut FfiCancelToken {
149    crate::abi::leak::track(crate::abi::leak::TOKENS, 1);
150    Arc::into_raw(Arc::new(FfiCancelToken::default())).cast_mut()
151}
152
153/// The body of `{prefix}_cancel_token_cancel`: request cancellation and wake
154/// every future waiting on the token. Idempotent; a null token is a no-op.
155///
156/// # Safety
157///
158/// `token` must be null or a token from [`cancel_token_create`] whose
159/// consumer reference hasn't been destroyed.
160pub unsafe fn cancel_token_cancel(token: *mut FfiCancelToken) {
161    // SAFETY: the caller guarantees a live token or null.
162    if let Some(t) = unsafe { token.as_ref() } {
163        t.cancel();
164    }
165}
166
167/// The body of `{prefix}_cancel_token_is_cancelled`. A null token reads as
168/// never cancelled.
169///
170/// # Safety
171///
172/// Same contract as [`cancel_token_cancel`].
173#[must_use]
174pub unsafe fn cancel_token_is_cancelled(token: *const FfiCancelToken) -> bool {
175    // SAFETY: the caller guarantees a live token or null.
176    unsafe { token.as_ref() }.is_some_and(FfiCancelToken::is_cancelled)
177}
178
179/// The body of `{prefix}_cancel_token_destroy`: release the consumer's
180/// reference. In-flight calls keep their own references, so this is safe to
181/// call while a call that received the token is still running. A null token
182/// is a no-op.
183///
184/// # Safety
185///
186/// `token` must be null or a token from [`cancel_token_create`], and the
187/// consumer must not use it again.
188pub unsafe fn cancel_token_destroy(token: *mut FfiCancelToken) {
189    if token.is_null() {
190        return;
191    }
192    crate::abi::leak::track(crate::abi::leak::TOKENS, -1);
193    // SAFETY: `token` came from `Arc::into_raw` and its reference is
194    // released exactly once.
195    unsafe { Arc::decrement_strong_count(token.cast_const()) };
196}
197
198/// The producer's handle on a consumer's cancel token, accepted as the final
199/// parameter of a `#[weaveffi::cancellable]` `async fn`.
200///
201/// The handle holds its own reference to the token, so it stays valid after
202/// the consumer destroys its reference. Cancellation is enforced by the
203/// runtime (the producer's future is dropped as soon as the token fires); a
204/// producer polls [`is_cancelled`](Self::is_cancelled) only when it wants to
205/// clean up cooperatively, for example between batches of work on another
206/// thread.
207#[derive(Debug, Default)]
208pub struct CancelToken {
209    state: Option<Arc<FfiCancelToken>>,
210}
211
212impl CancelToken {
213    /// Take a producer reference on the token a launcher received. A null
214    /// token yields a handle that's never cancelled.
215    ///
216    /// This is the entry point the `#[weaveffi::module]` expansion calls;
217    /// producers receive an already-built token.
218    ///
219    /// # Safety
220    ///
221    /// `raw` must be null or a live token from [`cancel_token_create`].
222    #[doc(hidden)]
223    #[must_use]
224    pub unsafe fn from_raw(raw: *const FfiCancelToken) -> Self {
225        if raw.is_null() {
226            return Self::default();
227        }
228        // SAFETY: `raw` came from `Arc::into_raw` and is live, so adding a
229        // reference and adopting it is sound.
230        let state = unsafe {
231            Arc::increment_strong_count(raw);
232            Arc::from_raw(raw)
233        };
234        crate::abi::leak::track(crate::abi::leak::TOKENS, 1);
235        Self { state: Some(state) }
236    }
237
238    /// Whether the consumer has requested cancellation.
239    #[must_use]
240    pub fn is_cancelled(&self) -> bool {
241        self.state
242            .as_deref()
243            .is_some_and(FfiCancelToken::is_cancelled)
244    }
245}
246
247impl Drop for CancelToken {
248    fn drop(&mut self) {
249        if self.state.is_some() {
250            crate::abi::leak::track(crate::abi::leak::TOKENS, -1);
251        }
252    }
253}
254
255impl Clone for CancelToken {
256    fn clone(&self) -> Self {
257        if self.state.is_some() {
258            crate::abi::leak::track(crate::abi::leak::TOKENS, 1);
259        }
260        Self {
261            state: self.state.clone(),
262        }
263    }
264}
265
266/// A future that resolves to `Some(output)` when `inner` completes, or to
267/// `None` as soon as the token is cancelled, dropping `inner` first.
268///
269/// Async launchers for `cancellable` functions wrap the producer's future in
270/// one of these. The waker is registered on the token before the flag is
271/// checked, so a cancellation racing a poll is never missed.
272#[derive(Debug)]
273pub struct Cancellable<F> {
274    inner: Option<Pin<Box<F>>>,
275    token: CancelToken,
276    slot: Option<Arc<AtomicWaker>>,
277}
278
279impl<F: Future> Cancellable<F> {
280    /// Race `fut` against `token`. A handle on a null token never cancels.
281    pub fn new(fut: F, token: CancelToken) -> Self {
282        Self {
283            inner: Some(Box::pin(fut)),
284            token,
285            slot: None,
286        }
287    }
288}
289
290impl<F: Future> Future for Cancellable<F> {
291    type Output = Option<F::Output>;
292
293    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
294        let this = &mut *self;
295        if let Some(state) = this.token.state.as_deref() {
296            let slot = this.slot.get_or_insert_with(|| {
297                let slot = Arc::new(AtomicWaker::default());
298                state.register(&slot);
299                slot
300            });
301            slot.register(cx.waker());
302            if state.is_cancelled() {
303                this.inner = None;
304                state.deregister(slot);
305                return Poll::Ready(None);
306            }
307        }
308        let Some(inner) = this.inner.as_mut() else {
309            // Already resolved; a well-behaved executor never polls again.
310            return Poll::Pending;
311        };
312        match inner.as_mut().poll(cx) {
313            Poll::Pending => Poll::Pending,
314            Poll::Ready(out) => {
315                this.inner = None;
316                if let (Some(state), Some(slot)) = (this.token.state.as_deref(), &this.slot) {
317                    state.deregister(slot);
318                }
319                Poll::Ready(Some(out))
320            }
321        }
322    }
323}
324
325impl<F> Drop for Cancellable<F> {
326    fn drop(&mut self) {
327        if let (Some(state), Some(slot)) = (self.token.state.as_deref(), &self.slot) {
328            state.deregister(slot);
329        }
330    }
331}
332
333#[cfg(test)]
334mod tests {
335    use super::*;
336
337    #[test]
338    fn lifecycle_and_null_safety() {
339        let token = cancel_token_create();
340        assert!(!unsafe { cancel_token_is_cancelled(token) });
341        unsafe { cancel_token_cancel(token) };
342        unsafe { cancel_token_cancel(token) };
343        assert!(unsafe { cancel_token_is_cancelled(token) });
344        unsafe { cancel_token_destroy(token) };
345
346        unsafe {
347            cancel_token_cancel(std::ptr::null_mut());
348            assert!(!cancel_token_is_cancelled(std::ptr::null()));
349            cancel_token_destroy(std::ptr::null_mut());
350        }
351        assert!(!unsafe { CancelToken::from_raw(std::ptr::null()) }.is_cancelled());
352    }
353
354    #[test]
355    fn producer_reference_outlives_the_consumer() {
356        let token = cancel_token_create();
357        let handle = unsafe { CancelToken::from_raw(token) };
358        unsafe { cancel_token_cancel(token) };
359        unsafe { cancel_token_destroy(token) };
360        assert!(handle.is_cancelled());
361    }
362
363    #[test]
364    fn cancelling_wakes_and_drops_the_inner_future() {
365        struct Flag(Arc<AtomicBool>);
366        impl Drop for Flag {
367            fn drop(&mut self) {
368                self.0.store(true, Ordering::SeqCst);
369            }
370        }
371
372        let token = cancel_token_create();
373        let handle = unsafe { CancelToken::from_raw(token) };
374        let dropped = Arc::new(AtomicBool::new(false));
375        let guard = Flag(Arc::clone(&dropped));
376        let fut = Cancellable::new(
377            async move {
378                let _guard = guard;
379                std::future::pending::<()>().await;
380            },
381            handle,
382        );
383        let raw = token as usize;
384        let canceller = std::thread::spawn(move || {
385            std::thread::sleep(std::time::Duration::from_millis(20));
386            unsafe { cancel_token_cancel(raw as *mut FfiCancelToken) };
387            unsafe { cancel_token_destroy(raw as *mut FfiCancelToken) };
388        });
389        assert!(crate::abi::block_on(fut).is_none());
390        assert!(dropped.load(Ordering::SeqCst));
391        canceller.join().unwrap();
392    }
393
394    #[test]
395    fn atomic_waker_wakes_the_latest_registration() {
396        struct Count(AtomicUsize);
397        impl std::task::Wake for Count {
398            fn wake(self: Arc<Self>) {
399                self.0.fetch_add(1, Ordering::SeqCst);
400            }
401        }
402        let a = Arc::new(Count(AtomicUsize::new(0)));
403        let b = Arc::new(Count(AtomicUsize::new(0)));
404        let slot = AtomicWaker::default();
405        slot.wake();
406        slot.register(&Waker::from(Arc::clone(&a)));
407        slot.register(&Waker::from(Arc::clone(&b)));
408        slot.wake();
409        slot.wake();
410        assert_eq!(a.0.load(Ordering::SeqCst), 0);
411        assert_eq!(b.0.load(Ordering::SeqCst), 1);
412    }
413
414    #[test]
415    fn one_token_cancels_every_call_it_was_passed_to() {
416        let token = cancel_token_create();
417        let calls: Vec<_> = (0..3)
418            .map(|_| {
419                Cancellable::new(std::future::pending::<()>(), unsafe {
420                    CancelToken::from_raw(token)
421                })
422            })
423            .collect();
424        let raw = token as usize;
425        let canceller = std::thread::spawn(move || {
426            std::thread::sleep(std::time::Duration::from_millis(20));
427            unsafe { cancel_token_cancel(raw as *mut FfiCancelToken) };
428        });
429        for call in calls {
430            assert!(crate::abi::block_on(call).is_none());
431        }
432        canceller.join().unwrap();
433        unsafe { cancel_token_destroy(token) };
434    }
435
436    #[test]
437    fn an_uncancelled_future_completes() {
438        let handle = CancelToken::default();
439        assert_eq!(
440            crate::abi::block_on(Cancellable::new(async { 3 }, handle)),
441            Some(3)
442        );
443    }
444}
weaveffi/abi/cancel.rs

weaveffi/abi/
cancel.rs
