use std::cell::{Cell, RefCell};
use std::future::Future;
use std::pin::Pin;
use std::rc::Rc;
use std::sync::atomic::AtomicBool;
use std::task::Context;
use self::waker::WakerState;
/// Must be called on every top-level future corresponding to a method call of a
/// canister by the IC, other than async functions marked `#[update]` or similar.
#[cfg_attr(not(target_arch = "wasm32"), allow(unused_variables, unreachable_code))]
pub fn spawn<F: 'static + Future<Output = ()>>(future: F) {
#[cfg(not(target_arch = "wasm32"))]
panic!("Cannot be run outside of wasm!"); // really, just cannot be run in a multi-threaded environment
let pinned_future = Box::pin(future);
let waker_state = Rc::new(WakerState {
future: RefCell::new(pinned_future),
previous_trap: Cell::new(false),
});
let waker = waker::waker(Rc::clone(&waker_state));
let _ = waker_state
.future
.borrow_mut()
.as_mut()
.poll(&mut Context::from_waker(&waker));
}
pub(crate) static CLEANUP: AtomicBool = AtomicBool::new(false);
// This module contains the implementation of a waker we're using for waking
// top-level futures (the ones returned by canister methods). Rc handles the
// heap management for us. Hence, it will be deallocated once we exit the scope and
// we're not interested in the result, as it can only be a unit `()` if the
// waker was used as intended.
// Sizable unsafe code is mandatory here; Future::poll cannot be executed without implementing
// RawWaker in terms of raw pointers.
mod waker {
use super::*;
use std::{
rc::Rc,
sync::atomic::Ordering,
task::{RawWaker, RawWakerVTable, Waker},
};
// The fields have separate RefCells in order to be modified separately.
pub(crate) struct WakerState {
pub future: RefCell<Pin<Box<dyn Future<Output = ()>>>>,
pub previous_trap: Cell<bool>,
}
static MY_VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop);
/// # Safety
///
/// The pointer must be an owning (i.e. represented in the refcount), Rc-allocated pointer to a `WakerState`.
unsafe fn raw_waker(ptr: *const ()) -> RawWaker {
// SAFETY: All the function pointers in MY_VTABLE correctly operate on the pointer in question.
RawWaker::new(ptr, &MY_VTABLE)
}
/// # Safety
///
/// This function should only be called by a [Waker] created by [`waker`].
unsafe fn clone(ptr: *const ()) -> RawWaker {
// SAFETY: The function's contract guarantees that this pointer is an Rc to a WakerState, and borrows the data from ptr.
unsafe {
Rc::increment_strong_count(ptr);
raw_waker(ptr)
}
}
// Our waker will be called if one of the response callbacks is triggered.
// Then, the waker will restore the future from the pointer we passed into the
// waker inside the `spawn` function and poll the future again. Rc takes care of
// the heap management for us. If CLEANUP is set, then we're recovering from
// a callback trap, and want to drop the future without executing any more of it;
// if previous_trap is set, then we already recovered from a callback trap in a
// different callback, and should immediately trap again in this one.
//
/// # Safety
///
/// This function should only be called by a [Waker] created by [`waker`].
unsafe fn wake(ptr: *const ()) {
// SAFETY: The function's contract guarantees that the pointer is an Rc to a WakerState, and that this call takes ownership of the data.
let state = unsafe { Rc::from_raw(ptr as *const WakerState) };
// Must check CLEANUP *before* previous_trap, as we may be recovering from the following immediate trap.
if super::CLEANUP.load(Ordering::Relaxed) {
state.previous_trap.set(true);
} else if state.previous_trap.get() {
crate::trap("Call already trapped");
} else {
let waker = waker(Rc::clone(&state));
let Ok(mut borrow) = state.future.try_borrow_mut() else {
// If this is already borrowed, then wake was called from inside poll. There's not a lot we can do about this - we are not
// a true scheduler and so cannot immediately schedule another poll, nor can we reentrantly lock the future. So we ignore it.
// This will be disappointing to types like FuturesUnordered that expected this to work, but since the only source of asynchrony
// and thus a guaranteed source of wakeup notifications is the ic0.call_new callback, this shouldn't cause any actual problems.
return;
};
let pinned_future = borrow.as_mut();
let _ = pinned_future.poll(&mut Context::from_waker(&waker));
}
}
/// # Safety
///
/// This function should only be called by a [Waker] created by [waker].
unsafe fn wake_by_ref(ptr: *const ()) {
// SAFETY:
// The function's contract guarantees that the pointer is an Rc to a WakerState, and that this call borrows the data.
// wake has the same contract, except it takes ownership instead of borrowing. Which just requires incrementing the refcount.
unsafe {
Rc::increment_strong_count(ptr as *const WakerState);
wake(ptr);
}
}
/// # Safety
///
/// This function should only be called by a [Waker] created by [waker].
unsafe fn drop(ptr: *const ()) {
// SAFETY: The function contract guarantees that the pointer is an Rc to a WakerState, and that this call takes ownership of the data.
unsafe {
Rc::from_raw(ptr as *const WakerState);
}
}
/// Creates a new Waker.
pub(crate) fn waker(state: Rc<WakerState>) -> Waker {
let ptr = Rc::into_raw(state);
// SAFETY:
// The pointer is an owning, Rc-allocated pointer to a WakerState, and therefore can be passed to raw_waker
// The functions in the vtable are passed said ptr
// The functions in the vtable uphold RawWaker's contract
unsafe { Waker::from_raw(raw_waker(ptr as *const ())) }
}
}