task_scope/
waker.rs

1use futures::future::OptionFuture;
2use futures_intrusive::channel::{shared::StateReceiveFuture, StateId};
3
4use std::task::{Context, RawWaker, RawWakerVTable, Waker};
5
6use crate::Token;
7
8pub(crate) struct WakerData {
9    pub(crate) token: Token,
10    original: Waker,
11}
12
13// to implement safely, need https://github.com/rust-lang/rfcs/issues/2746
14// hope this works
15pub(crate) unsafe fn retrieve_data<'c, 'w>(cx: &'c mut Context<'w>) -> Option<&'w WakerData> {
16    let waker_ptr: *const Waker = cx.waker();
17    let data_ptr: *const *const WakerData = waker_ptr.cast();
18    let vtable_ptr: *const &'static RawWakerVTable = data_ptr
19        .cast::<u8>()
20        .add(std::mem::size_of::<*const ()>())
21        .cast();
22    let vtable_ref = std::ptr::read(vtable_ptr);
23
24    if vtable_ref != &VTABLE {
25        // polled from outside of a scope
26        // cancellation is not supported
27        None
28    } else {
29        let data_ref = &**data_ptr;
30        Some(data_ref)
31    }
32}
33
34// returns a future that resolves to
35// - Some(Some(state_id, true)) when the current scope requests cancellation
36// - Some(None) when the current scope is dropped
37// - None when the current context doesn't support cancellation
38pub(crate) fn cancellation<'c, 'w>(
39    cx: &'c mut Context<'w>,
40    state_id: StateId,
41) -> OptionFuture<StateReceiveFuture<parking_lot::RawMutex, bool>> {
42    unsafe {
43        if let Some(data) = retrieve_data(cx) {
44            Some(data.token.cancel.receive(state_id)).into()
45        } else {
46            None.into()
47        }
48    }
49}
50
51pub(crate) unsafe fn waker(token: Token, original: Waker) -> Waker {
52    Waker::from_raw(raw_waker(token, original))
53}
54
55static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop);
56
57fn raw_waker(token: Token, original: Waker) -> RawWaker {
58    let data = Box::into_raw(Box::new(WakerData { token, original }));
59
60    RawWaker::new(data.cast(), &VTABLE)
61}
62
63unsafe fn clone(data: *const ()) -> RawWaker {
64    let data_ref = &*data.cast::<WakerData>();
65    raw_waker(data_ref.token.clone(), data_ref.original.clone())
66}
67
68unsafe fn wake(data: *const ()) {
69    let data = Box::from_raw(data.cast::<WakerData>() as *mut WakerData);
70    data.original.wake();
71}
72
73unsafe fn wake_by_ref(data: *const ()) {
74    let data_ref = &*data.cast::<WakerData>();
75    data_ref.original.wake_by_ref();
76}
77
78unsafe fn drop(data: *const ()) {
79    Box::from_raw(data.cast::<WakerData>() as *mut WakerData);
80}
task_scope/waker.rs

task_scope/
waker.rs
