bastion_executor/
run.rs

1//!
2//! Blocking run of the async processes
3//!
4//!
5use crate::worker;
6use crossbeam_utils::sync::Parker;
7use lightproc::proc_stack::ProcStack;
8use std::cell::{Cell, UnsafeCell};
9use std::future::Future;
10use std::mem;
11use std::mem::ManuallyDrop;
12use std::pin::Pin;
13use std::sync::Arc;
14use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
15
16///
17/// This method blocks the current thread until passed future is resolved with an output (including the panic).
18///
19/// It is called `block_on` or `blocking` in some executors.
20///
21/// # Example
22/// ```rust
23/// use bastion_executor::prelude::*;
24/// use lightproc::prelude::*;
25/// let mut sum = 0;
26///
27/// run(
28///     async {
29///         (0..10_000_000).for_each(|_| {
30///             sum += 1;
31///         });
32///     },
33///     ProcStack::default(),
34/// );
35/// ```
36pub fn run<F, T>(future: F, stack: ProcStack) -> T
37where
38    F: Future<Output = T>,
39{
40    unsafe {
41        // A place on the stack where the result will be stored.
42        let out = &mut UnsafeCell::new(None);
43
44        // Wrap the future into one that stores the result into `out`.
45        let future = {
46            let out = out.get();
47
48            async move {
49                *out = Some(future.await);
50            }
51        };
52
53        // Pin the future onto the stack.
54        pin_utils::pin_mut!(future);
55
56        // Transmute the future into one that is futurestatic.
57        let future = mem::transmute::<
58            Pin<&'_ mut dyn Future<Output = ()>>,
59            Pin<&'static mut dyn Future<Output = ()>>,
60        >(future);
61
62        // Block on the future and and wait for it to complete.
63        worker::set_stack(&stack, || block(future));
64
65        // Take out the result.
66        match (*out.get()).take() {
67            Some(v) => v,
68            _ => unimplemented!(),
69        }
70    }
71}
72
73fn block<F, T>(f: F) -> T
74where
75    F: Future<Output = T>,
76{
77    thread_local! {
78        // May hold a pre-allocated parker that can be reused for efficiency.
79        //
80        // Note that each invocation of `block` needs its own parker. In particular, if `block`
81        // recursively calls itself, we must make sure that each recursive call uses a distinct
82        // parker instance.
83        static CACHE: Cell<Option<Arc<Parker>>> = Cell::new(None);
84    }
85
86    pin_utils::pin_mut!(f);
87
88    CACHE.with(|cache| {
89        // Reuse a cached parker or create a new one for this invocation of `block`.
90        let arc_parker: Arc<Parker> = cache.take().unwrap_or_else(|| Arc::new(Parker::new()));
91
92        let ptr = (&*arc_parker as *const Parker) as *const ();
93        let vt = vtable();
94
95        let waker = unsafe { ManuallyDrop::new(Waker::from_raw(RawWaker::new(ptr, vt))) };
96        let cx = &mut Context::from_waker(&waker);
97
98        loop {
99            if let Poll::Ready(t) = f.as_mut().poll(cx) {
100                // Save the parker for the next invocation of `block`.
101                cache.set(Some(arc_parker));
102                return t;
103            }
104            arc_parker.park();
105        }
106    })
107}
108
109fn vtable() -> &'static RawWakerVTable {
110    unsafe fn clone_raw(ptr: *const ()) -> RawWaker {
111        #![allow(clippy::redundant_clone)]
112        let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const Parker));
113        mem::forget(arc.clone());
114        RawWaker::new(ptr, vtable())
115    }
116
117    unsafe fn wake_raw(ptr: *const ()) {
118        let arc = Arc::from_raw(ptr as *const Parker);
119        arc.unparker().unpark();
120    }
121
122    unsafe fn wake_by_ref_raw(ptr: *const ()) {
123        let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const Parker));
124        arc.unparker().unpark();
125    }
126
127    unsafe fn drop_raw(ptr: *const ()) {
128        drop(Arc::from_raw(ptr as *const Parker))
129    }
130
131    &RawWakerVTable::new(clone_raw, wake_raw, wake_by_ref_raw, drop_raw)
132}
bastion_executor/run.rs

bastion_executor/
run.rs
