mirror of
https://gitlab.com/fabinfra/fabaccess/bffh.git
synced 2024-10-19 00:02:32 +02:00
155 lines
5.4 KiB
Rust
155 lines
5.4 KiB
Rust
//!
|
|
//! Blocking run of the async processes
|
|
//!
|
|
//!
|
|
use crate::worker;
|
|
use crossbeam_utils::sync::{Parker, Unparker};
|
|
use std::cell::Cell;
|
|
use std::future::Future;
|
|
use std::mem;
|
|
use std::mem::{ManuallyDrop, MaybeUninit};
|
|
use std::pin::Pin;
|
|
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
|
|
use crate::proc_stack::ProcStack;
|
|
|
|
///
|
|
/// This method blocks the current thread until passed future is resolved with an output.
|
|
///
|
|
/// It is called `block_on` or `blocking` in some executors.
|
|
///
|
|
/// # Example
|
|
/// ```rust
|
|
/// use executor::prelude::*;
|
|
/// use lightproc::prelude::*;
|
|
/// let mut sum = 0;
|
|
///
|
|
/// run(
|
|
/// async {
|
|
/// (0..10_000_000).for_each(|_| {
|
|
/// sum += 1;
|
|
/// });
|
|
/// },
|
|
/// ProcStack::default(),
|
|
/// );
|
|
/// ```
|
|
pub fn run<F, T>(future: F, stack: ProcStack) -> T
|
|
where
|
|
F: Future<Output = T>,
|
|
{
|
|
unsafe {
|
|
// An explicitly uninitialized `T`. Until `assume_init` is called this will not call any
|
|
// drop code for T
|
|
let mut out = MaybeUninit::uninit();
|
|
|
|
// Wrap the future into one that stores the result into `out`.
|
|
let future = {
|
|
let out = out.as_mut_ptr();
|
|
|
|
async move {
|
|
*out = future.await;
|
|
}
|
|
};
|
|
|
|
// Pin the future onto the stack.
|
|
pin_utils::pin_mut!(future);
|
|
|
|
// Extend the lifetime of the future to 'static.
|
|
let future = mem::transmute::<
|
|
Pin<&'_ mut dyn Future<Output = ()>>,
|
|
Pin<&'static mut dyn Future<Output = ()>>,
|
|
>(future);
|
|
|
|
// Block on the future and and wait for it to complete.
|
|
worker::set_stack(&stack, || block(future));
|
|
|
|
// Assume that if the future completed and didn't panic it fully initialized its output
|
|
out.assume_init()
|
|
}
|
|
}
|
|
|
|
fn block<F, T>(f: F) -> T
|
|
where
|
|
F: Future<Output = T>,
|
|
{
|
|
thread_local! {
|
|
// May hold a pre-allocated parker that can be reused for efficiency.
|
|
//
|
|
// Note that each invocation of `block` needs its own parker. In particular, if `block`
|
|
// recursively calls itself, we must make sure that each recursive call uses a distinct
|
|
// parker instance.
|
|
static CACHE: Cell<Option<Parker>> = Cell::new(None);
|
|
}
|
|
|
|
pin_utils::pin_mut!(f);
|
|
|
|
CACHE.with(|cache| {
|
|
// Reuse a cached parker or create a new one for this invocation of `block`.
|
|
let parker: Parker = cache.take().unwrap_or_else(|| Parker::new());
|
|
|
|
let ptr = Unparker::into_raw(parker.unparker().clone());
|
|
let vt = vtable();
|
|
|
|
// Waker must not be dropped until it's no longer required. We also happen to know that a
|
|
// Parker contains at least one reference to `Unparker` so the relevant `Unparker` will not
|
|
// be dropped at least until the `Parker` is.
|
|
let waker = unsafe { Waker::from_raw(RawWaker::new(ptr, vt)) };
|
|
let cx = &mut Context::from_waker(&waker);
|
|
|
|
loop {
|
|
if let Poll::Ready(t) = f.as_mut().poll(cx) {
|
|
// Save the parker for the next invocation of `block`.
|
|
cache.set(Some(parker));
|
|
return t;
|
|
}
|
|
parker.park();
|
|
}
|
|
})
|
|
}
|
|
|
|
fn vtable() -> &'static RawWakerVTable {
|
|
/// This function will be called when the RawWaker gets cloned, e.g. when the Waker in which
|
|
/// the RawWaker is stored gets cloned.
|
|
//
|
|
/// The implementation of this function must retain all resources that are required for this
|
|
/// additional instance of a RawWaker and associated task. Calling wake on the resulting
|
|
/// RawWaker should result in a wakeup of the same task that would have been awoken by the
|
|
/// original RawWaker.
|
|
unsafe fn clone_raw(ptr: *const ()) -> RawWaker {
|
|
// [`Unparker`] implements `Clone` and upholds the contract stated above. The current
|
|
// Implementation is simply an Arc over the actual inner values.
|
|
let unparker = Unparker::from_raw(ptr).clone();
|
|
RawWaker::new(Unparker::into_raw(unparker), vtable())
|
|
}
|
|
|
|
/// This function will be called when wake is called on the Waker. It must wake up the task
|
|
/// associated with this RawWaker.
|
|
///
|
|
/// The implementation of this function must make sure to release any resources that are
|
|
/// associated with this instance of a RawWaker and associated task.
|
|
unsafe fn wake_raw(ptr: *const ()) {
|
|
// We reconstruct the Unparker from the pointer here thus ensuring it is dropped at the
|
|
// end of this function call.
|
|
Unparker::from_raw(ptr).unpark();
|
|
}
|
|
|
|
/// This function will be called when wake_by_ref is called on the Waker. It must wake up the
|
|
/// task associated with this RawWaker.
|
|
///
|
|
/// This function is similar to wake, but must not consume the provided data pointer.
|
|
unsafe fn wake_by_ref_raw(ptr: *const ()) {
|
|
// We **must not** drop the resulting Unparker so we wrap it in `ManuallyDrop`.
|
|
let unparker = ManuallyDrop::new(Unparker::from_raw(ptr));
|
|
unparker.unpark();
|
|
}
|
|
|
|
/// This function gets called when a RawWaker gets dropped.
|
|
///
|
|
/// The implementation of this function must make sure to release any resources that are
|
|
/// associated with this instance of a RawWaker and associated task.
|
|
unsafe fn drop_raw(ptr: *const ()) {
|
|
drop(Unparker::from_raw(ptr))
|
|
}
|
|
|
|
&RawWakerVTable::new(clone_raw, wake_raw, wake_by_ref_raw, drop_raw)
|
|
}
|