wasmer/backend/sys/
async_runtime.rs

1use std::{
2    cell::RefCell,
3    collections::HashMap,
4    future::Future,
5    marker::PhantomData,
6    pin::Pin,
7    ptr,
8    rc::Rc,
9    sync::{LazyLock, atomic::AtomicU64},
10    task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
11};
12
13use corosensei::{Coroutine, CoroutineResult, Yielder};
14use dashmap::DashMap;
15
16use super::entities::function::Function as SysFunction;
17use crate::{
18    AsStoreAsync, AsStoreMut, AsStoreRef, ForcedStoreInstallGuard, LocalRwLockWriteGuard,
19    RuntimeError, Store, StoreAsync, StoreContext, StoreInner, StoreMut, StoreRef, Value,
20};
21use wasmer_types::StoreId;
22
23type HostFuture = Pin<Box<dyn Future<Output = Result<Vec<Value>, RuntimeError>> + 'static>>;
24
25pub(crate) fn call_function_async(
26    function: SysFunction,
27    store: StoreAsync,
28    params: Vec<Value>,
29) -> AsyncCallFuture {
30    AsyncCallFuture::new(function, store, params)
31}
32
33struct AsyncYield(HostFuture);
34
35enum AsyncResume {
36    Start,
37    HostFutureReady(Result<Vec<Value>, RuntimeError>),
38}
39
40type AsyncCallFutureId = u64;
41
42static NEXT_ASYNC_CALL_FUTURE_ID: AtomicU64 = AtomicU64::new(0);
43
44static ASYNC_CALL_FUTURE_WAKERS: LazyLock<DashMap<StoreId, HashMap<AsyncCallFutureId, Waker>>> =
45    LazyLock::new(DashMap::new);
46
47#[allow(clippy::type_complexity)]
48pub(crate) struct AsyncCallFuture {
49    id: AsyncCallFutureId,
50
51    coroutine: Option<Coroutine<AsyncResume, AsyncYield, Result<Box<[Value]>, RuntimeError>>>,
52    pending_store_install: Option<Pin<Box<dyn Future<Output = ForcedStoreInstallGuard>>>>,
53    pending_future: Option<HostFuture>,
54    next_resume: Option<AsyncResume>,
55    result: Option<Result<Box<[Value]>, RuntimeError>>,
56
57    // Store handle we can use to lock the store down
58    store: StoreAsync,
59}
60
61// We can't use any of the existing AsStoreMut types here, since we keep
62// changing the store context underneath us while the coroutine yields.
63// To work around it, we use this dummy struct, which just grabs the store
64// from the store context. Since we always have a store context installed
65// when resuming the coroutine, this is safe in that it can access the store
66// through the store context. HOWEVER, references returned from this struct
67// CAN NOT BE HELD ACROSS A YIELD POINT. We don't do this anywhere in the
68// `Function::call` code.
69struct AsyncCallStoreMut {
70    store_id: StoreId,
71}
72
73impl AsStoreRef for AsyncCallStoreMut {
74    fn as_store_ref(&self) -> StoreRef<'_> {
75        // Safety: This is only used with Function::call, which doesn't store
76        // the returned reference anywhere, including when calling into WASM
77        // code.
78        unsafe {
79            StoreRef {
80                inner: StoreContext::get_current_transient(self.store_id)
81                    .as_ref()
82                    .unwrap(),
83            }
84        }
85    }
86}
87
88impl AsStoreMut for AsyncCallStoreMut {
89    fn as_store_mut(&mut self) -> StoreMut<'_> {
90        // Safety: This is only used with Function::call, which doesn't store
91        // the returned reference anywhere, including when calling into WASM
92        // code.
93        unsafe {
94            StoreMut {
95                inner: StoreContext::get_current_transient(self.store_id)
96                    .as_mut()
97                    .unwrap(),
98            }
99        }
100    }
101
102    fn objects_mut(&mut self) -> &mut crate::StoreObjects {
103        // Safety: This is only used with Function::call, which doesn't store
104        // the returned reference anywhere, including when calling into WASM
105        // code.
106        unsafe {
107            &mut StoreContext::get_current_transient(self.store_id)
108                .as_mut()
109                .unwrap()
110                .objects
111        }
112    }
113}
114
115impl AsyncCallFuture {
116    pub(crate) fn new(function: SysFunction, store: StoreAsync, params: Vec<Value>) -> Self {
117        let store_id = store.id;
118        let coroutine =
119            Coroutine::new(move |yielder: &Yielder<AsyncResume, AsyncYield>, resume| {
120                assert!(matches!(resume, AsyncResume::Start));
121
122                let ctx_state = CoroutineContext::new(yielder);
123                ctx_state.enter();
124                let result = {
125                    let mut store_mut = AsyncCallStoreMut { store_id };
126                    function.call(&mut store_mut, &params)
127                };
128                ctx_state.leave();
129                result
130            });
131
132        Self {
133            id: NEXT_ASYNC_CALL_FUTURE_ID.fetch_add(1, std::sync::atomic::Ordering::SeqCst),
134            coroutine: Some(coroutine),
135            pending_store_install: None,
136            pending_future: None,
137            next_resume: Some(AsyncResume::Start),
138            result: None,
139            store,
140        }
141    }
142
143    fn remove_from_wakers_list(&self) {
144        let mut wakers_entry = match ASYNC_CALL_FUTURE_WAKERS.entry(self.store.store_id()) {
145            dashmap::Entry::Occupied(o) => o,
146            dashmap::Entry::Vacant(v) => return,
147        };
148
149        let mut waker_map_ref = wakers_entry.get_mut();
150        waker_map_ref.remove(&self.id);
151        if waker_map_ref.is_empty() {
152            wakers_entry.remove();
153        }
154    }
155}
156
157impl Future for AsyncCallFuture {
158    type Output = Result<Box<[Value]>, RuntimeError>;
159
160    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
161        loop {
162            let store_id = self.store.store_id();
163
164            #[cfg(feature = "experimental-host-interrupt")]
165            {
166                if super::vm::interrupt_registry::is_interrupted(store_id) {
167                    self.remove_from_wakers_list();
168                    return Poll::Ready(Err(super::vm::Trap::lib(
169                        super::vm::TrapCode::HostInterrupt,
170                    )
171                    .into()));
172                }
173            }
174
175            if let Some(future) = self.pending_future.as_mut() {
176                match future.as_mut().poll(cx) {
177                    Poll::Ready(result) => {
178                        self.pending_future = None;
179                        self.next_resume = Some(AsyncResume::HostFutureReady(result));
180                    }
181                    Poll::Pending => return Poll::Pending,
182                }
183            }
184
185            // If we're ready, return early
186            if self.coroutine.is_none() {
187                self.remove_from_wakers_list();
188                return Poll::Ready(self.result.take().expect("polled after completion"));
189            }
190
191            {
192                let mut wakers_entry = ASYNC_CALL_FUTURE_WAKERS.entry(store_id).or_default();
193                wakers_entry.insert(self.id, cx.waker().clone());
194            }
195
196            // Start a store installation if not in progress already
197            if self.pending_store_install.is_none() {
198                self.pending_store_install = Some(Box::pin(install_store_context(StoreAsync {
199                    id: self.store.id,
200                    inner: self.store.inner.clone(),
201                })));
202            }
203
204            // Acquiring a store lock should be the last step before resuming
205            // the coroutine, to minimize the time we hold the lock.
206            let store_context_guard = match self
207                .pending_store_install
208                .as_mut()
209                .unwrap()
210                .as_mut()
211                .poll(cx)
212            {
213                Poll::Ready(guard) => {
214                    self.pending_store_install = None;
215                    guard
216                }
217                Poll::Pending => return Poll::Pending,
218            };
219
220            let resume_arg = self.next_resume.take().expect("no resume arg available");
221            let coroutine = self.coroutine.as_mut().unwrap();
222            match coroutine.resume(resume_arg) {
223                CoroutineResult::Yield(AsyncYield(fut)) => {
224                    self.pending_future = Some(fut);
225                }
226                CoroutineResult::Return(result) => {
227                    self.coroutine = None;
228                    self.result = Some(result);
229                }
230            }
231
232            // Uninstall the store context to unlock the store after the coroutine
233            // yields or returns.
234            drop(store_context_guard);
235        }
236    }
237}
238
239impl Drop for AsyncCallFuture {
240    fn drop(&mut self) {
241        self.remove_from_wakers_list();
242    }
243}
244
245async fn install_store_context(store: StoreAsync) -> ForcedStoreInstallGuard {
246    match unsafe { crate::StoreContext::try_get_current_async(store.id) } {
247        crate::GetStoreAsyncGuardResult::NotInstalled => {
248            // We always need to acquire a new write lock on the store.
249            let store_guard = store.inner.write().await;
250            unsafe { crate::StoreContext::install_async(store_guard) }
251        }
252        _ => {
253            // If we're already in a store context, it is unsafe to reuse
254            // the existing store ref since it'll also be accessible from
255            // the imported function that tried to poll us, which is a
256            // double mutable borrow.
257            // Note to people who discover this code: this *would* be safe
258            // if we had a separate variation of call_async that just
259            // used the existing coroutine context instead of spawning a
260            // new coroutine. However, the current call_async always spawns
261            // a new coroutine, so we can't allow this; every coroutine
262            // needs to own its write lock on the store to make sure there
263            // are no overlapping mutable borrows. If this is something
264            // you're interested in, feel free to open a GitHub issue outlining
265            // your use-case.
266            panic!(
267                "Function::call_async futures cannot be polled recursively \
268                from within another imported function. If you need to await \
269                a recursive call_async, consider spawning the future into \
270                your async runtime and awaiting the resulting task; \
271                e.g. tokio::task::spawn(func.call_async(...)).await"
272            );
273        }
274    }
275}
276
277pub enum AsyncRuntimeError {
278    YieldOutsideAsyncContext,
279    RuntimeError(RuntimeError),
280}
281
282pub(crate) fn block_on_host_future<Fut>(future: Fut) -> Result<Vec<Value>, AsyncRuntimeError>
283where
284    Fut: Future<Output = Result<Vec<Value>, RuntimeError>> + 'static,
285{
286    match CoroutineContext::get_current() {
287        None => {
288            // If there is no async context or we haven't entered it,
289            // we can still directly run a future that doesn't block
290            // inline.
291            run_immediate(future)
292        }
293        Some(context) => unsafe { context.as_ref().expect("valid context pointer") }
294            .block_on_future(Box::pin(future))
295            .map_err(AsyncRuntimeError::RuntimeError),
296    }
297}
298
299// This function will only *wake* pending futures belonging to the given store.
300// It is expected that the interrupt happens before this function is called, so
301// that the futures can catch and react to `interrupt_registry::is_interrupted`
302// correctly.
303pub(crate) fn notify_pending_futures_of_interrupt(store_id: StoreId) {
304    let dashmap::Entry::Occupied(entry) = ASYNC_CALL_FUTURE_WAKERS.entry(store_id) else {
305        return;
306    };
307
308    for waker in entry.get().values() {
309        waker.wake_by_ref();
310    }
311}
312
313thread_local! {
314    static CURRENT_CONTEXT: RefCell<Vec<*const CoroutineContext>> = const { RefCell::new(Vec::new()) };
315}
316
317struct CoroutineContext {
318    yielder: *const Yielder<AsyncResume, AsyncYield>,
319}
320
321impl CoroutineContext {
322    fn new(yielder: &Yielder<AsyncResume, AsyncYield>) -> Self {
323        Self {
324            yielder: yielder as *const _,
325        }
326    }
327
328    fn enter(&self) {
329        CURRENT_CONTEXT.with(|cell| {
330            let mut borrow = cell.borrow_mut();
331
332            // Push this coroutine on top of the active stack.
333            borrow.push(self as *const _);
334        })
335    }
336
337    // Note: we don't use a drop-style guard here on purpose; if a panic
338    // happens while a coroutine is running, CURRENT_CONTEXT will be in
339    // an inconsistent state. corosensei will unwind all coroutine stacks
340    // anyway, and if we had a guard that would get dropped and try to
341    // leave its context, it'd panic again at the assert_eq! below.
342    fn leave(&self) {
343        CURRENT_CONTEXT.with(|cell| {
344            let mut borrow = cell.borrow_mut();
345
346            // Pop this coroutine from the active stack.
347            assert_eq!(
348                borrow.pop(),
349                Some(self as *const _),
350                "Active coroutine stack corrupted"
351            );
352        });
353    }
354
355    fn get_current() -> Option<*const Self> {
356        CURRENT_CONTEXT.with(|cell| cell.borrow().last().copied())
357    }
358
359    fn block_on_future(&self, future: HostFuture) -> Result<Vec<Value>, RuntimeError> {
360        // Leave the coroutine context since we're yielding back to the
361        // parent stack, and will be inactive until the future is ready.
362        self.leave();
363
364        let yielder = unsafe { self.yielder.as_ref().expect("yielder pointer valid") };
365        let result = match yielder.suspend(AsyncYield(future)) {
366            AsyncResume::HostFutureReady(result) => result,
367            AsyncResume::Start => unreachable!("coroutine resumed without start"),
368        };
369
370        // Once the future is ready, we restore the current coroutine
371        // context.
372        self.enter();
373
374        result
375    }
376}
377
378fn run_immediate(
379    future: impl Future<Output = Result<Vec<Value>, RuntimeError>> + 'static,
380) -> Result<Vec<Value>, AsyncRuntimeError> {
381    let waker = futures::task::noop_waker();
382    let mut cx = Context::from_waker(&waker);
383    let mut future = Box::pin(future);
384    match future.as_mut().poll(&mut cx) {
385        Poll::Ready(result) => result.map_err(AsyncRuntimeError::RuntimeError),
386        Poll::Pending => Err(AsyncRuntimeError::YieldOutsideAsyncContext),
387    }
388}
wasmer/backend/sys/async_runtime.rs

wasmer/backend/sys/
async_runtime.rs
