.await for async native functions outside of the VM

This avoids manually polling boxed futures within the __call__
implementation for NativeAsyncFunctions.

This instead stashes the boxed future under `Context::async_call` so
that we can `.await` for it's completion within
`Context::run_async_with_budget`.

This approach should integrate much better with async runtimes like
Tokio, since it doesn't involve manually polling with a task::Context +
Waker that aren't managed by the current async runtime.

This also means the thread can block waiting for async IO without
polling for native function completion in a busy loop.

This still needs further iteration, but hopefully serves as a usable
draft / proof of concept.

Author: rib

core/engine/src/context/mod.rs

@@ -14,6 +14,7 @@ use intrinsics::Intrinsics;
 use temporal_rs::tzdb::FsTzdbProvider;
 
 use crate::job::Job;
+use crate::native_function::AsyncCallState;
 use crate::vm::RuntimeLimits;
 use crate::{
     builtins,
@@ -129,6 +130,17 @@ pub struct Context {
     parser_identifier: u32,
 
     data: HostDefined,
+
+    /// State of any boxed future that should be .awaited before continuing
+    /// to execute VM instructions.
+    ///
+    /// XXX: How do we make sure that any
+    /// `AsyncCallState::Finished(Result<JsValue>)` can't be garbage collected?
+    ///
+    /// XXX: there's maybe a better place for this, or better abstraction /
+    /// generalization than this, but this hopefully works for a draft /
+    /// proof-of-concept.
+    pub(crate) async_call: AsyncCallState,
 }
 
 impl std::fmt::Debug for Context {
@@ -1131,6 +1143,7 @@ impl ContextBuilder {
             parser_identifier: 0,
             can_block: self.can_block,
             data: HostDefined::default(),
+            async_call: AsyncCallState::None,
         };
 
         builtins::set_default_global_bindings(&mut context)?;

core/engine/src/native_function/mod.rs

@@ -9,7 +9,6 @@ use std::pin::Pin;
 use std::rc::Rc;
 
 use boa_gc::{custom_trace, Finalize, Gc, Trace};
-use futures_lite::FutureExt as _;
 
 use crate::job::NativeAsyncJob;
 use crate::value::JsVariant;
@@ -134,55 +133,32 @@ enum AsyncCallResult {
     Ready(JsResult<JsValue>),
 }
 
-enum AsyncRunningState {
-    None,
-    Calling,
-    Constructing { new_target: JsValue },
-}
 trait TraceableAsyncFunction: Trace {
-    fn running(&self) -> AsyncRunningState;
     fn call_or_construct(
         &self,
         this: &JsValue,
         args: &[JsValue],
         new_target: &JsValue,
         context: &mut Context,
     ) -> AsyncCallResult;
-    fn poll(&self, context: &mut Context) -> AsyncCallResult {
-        self.call_or_construct(&JsValue::undefined(), &[], &JsValue::undefined(), context)
-    }
 }
 
-#[derive(Finalize)]
-enum AsyncCallState {
+#[derive(Default, Finalize)]
+pub(crate) enum AsyncCallState {
+    #[default]
     None,
     Calling {
-        this: JsValue,
-        args: Vec<JsValue>,
-        f: Pin<Box<dyn Future<Output = JsResult<JsValue>>>>,
-    },
-    Constructing {
-        new_target: JsValue,
-        args: Vec<JsValue>,
         f: Pin<Box<dyn Future<Output = JsResult<JsValue>>>>,
     },
+    Finished(JsResult<JsValue>),
 }
 unsafe impl Trace for AsyncCallState {
     custom_trace!(this, mark, {
         match this {
-            AsyncCallState::None => {}
-            AsyncCallState::Calling { this, args, f: _ } => {
-                mark(this);
-                mark(args);
-            }
-            AsyncCallState::Constructing {
-                new_target,
-                args,
-                f: _,
-            } => {
-                mark(new_target);
-                mark(args);
-            }
+            AsyncCallState::None
+            | AsyncCallState::Calling { .. }
+            | AsyncCallState::Finished(Err(_)) => {}
+            AsyncCallState::Finished(Ok(value)) => mark(value),
         }
     });
 }
@@ -191,83 +167,36 @@ unsafe impl Trace for AsyncCallState {
 struct NativeAsyncFunction<T: Trace> {
     start: Rc<SpawnAsyncFunctionFn>,
     captures: T,
-    state: RefCell<AsyncCallState>,
 }
 unsafe impl<T: Trace> Trace for NativeAsyncFunction<T> {
     custom_trace!(this, mark, {
         mark(&this.captures);
-        mark(&*this.state.borrow());
     });
 }
 
 impl<T> TraceableAsyncFunction for NativeAsyncFunction<T>
 where
     T: Trace,
 {
-    fn running(&self) -> AsyncRunningState {
-        match &*self.state.borrow() {
-            AsyncCallState::None => AsyncRunningState::None,
-            AsyncCallState::Calling { .. } => AsyncRunningState::Calling,
-            AsyncCallState::Constructing { new_target, .. } => AsyncRunningState::Constructing {
-                new_target: new_target.clone(),
-            },
-        }
-    }
-
     fn call_or_construct(
         &self,
         this: &JsValue,
         args: &[JsValue],
         new_target: &JsValue,
         context: &mut Context,
     ) -> AsyncCallResult {
-        let mut state = self.state.borrow_mut();
-        match *state {
-            AsyncCallState::None => {
-                if new_target.is_undefined() {
-                    let f = (*self.start)(this, args, context);
-                    let f = match f {
-                        Ok(f) => f,
-                        Err(e) => return AsyncCallResult::Ready(Err(e)),
-                    };
-                    *state = AsyncCallState::Calling {
-                        this: this.clone(),
-                        args: args.to_vec(),
-                        f,
-                    };
-                } else {
-                    let f = (*self.start)(new_target, args, context);
-                    let f = match f {
-                        Ok(f) => f,
-                        Err(e) => return AsyncCallResult::Ready(Err(e)),
-                    };
-                    *state = AsyncCallState::Constructing {
-                        new_target: new_target.clone(),
-                        args: args.to_vec(),
-                        f,
-                    };
-                }
+        let f = if new_target.is_undefined() {
+            (*self.start)(this, args, context)
+        } else {
+            (*self.start)(new_target, args, context)
+        };
+        let f = match f {
+            Ok(f) => f,
+            Err(e) => return AsyncCallResult::Ready(Err(e)),
+        };
+        context.async_call = AsyncCallState::Calling { f };
 
-                AsyncCallResult::Pending
-            }
-            AsyncCallState::Calling { ref mut f, .. }
-            | AsyncCallState::Constructing { ref mut f, .. } => {
-                // FIXME: figure out how to work with an async Context / Waker from the application (e.g. from Tokio)
-                let waker = std::task::Waker::noop();
-                let mut context = std::task::Context::from_waker(waker);
-                let result = f.poll(&mut context);
-                match result {
-                    std::task::Poll::Pending => {
-                        //println!("Pending");
-                        AsyncCallResult::Pending
-                    }
-                    std::task::Poll::Ready(result) => {
-                        *state = AsyncCallState::None;
-                        AsyncCallResult::Ready(result)
-                    }
-                }
-            }
-        }
+        AsyncCallResult::Pending
     }
 }
 
@@ -495,7 +424,6 @@ impl NativeFunction {
         let ptr = Gc::into_raw(Gc::new(NativeAsyncFunction {
             start: Rc::new(f),
             captures,
-            state: RefCell::new(AsyncCallState::None),
         }));
 
         // SAFETY: The pointer returned by `into_raw` is only used to coerce to a trait object,
@@ -608,46 +536,39 @@ impl NativeFunction {
         context: &mut Context,
     ) -> AsyncCallResult {
         //println!("[NativeFunction] call, arg_count: {}", argument_count);
+        let args = context
+            .vm
+            .stack
+            .calling_convention_pop_arguments(argument_count);
+        let func = context.vm.stack.pop();
+        let this = context.vm.stack.pop();
+        let this_ref = if is_constructor {
+            &JsValue::undefined()
+        } else {
+            &this
+        };
+
         if let Inner::AsyncFn(ref f) = self.inner {
-            match f.running() {
-                AsyncRunningState::None => {
-                    let args = context
-                        .vm
-                        .stack
-                        .calling_convention_pop_arguments(argument_count);
-                    let func = context.vm.stack.pop();
-                    let this = context.vm.stack.pop();
-                    let this_ref = if is_constructor {
-                        &JsValue::undefined()
-                    } else {
-                        &this
-                    };
-                    let result =
-                        f.call_or_construct(this_ref, &args, &JsValue::undefined(), context);
-                    if matches!(result, AsyncCallResult::Pending) {
-                        context.vm.stack.push(this);
-                        context.vm.stack.push(func);
-                        context.vm.stack.calling_convention_push_arguments(&args);
-                    }
-                    result
+            let result = match &context.async_call {
+                AsyncCallState::None => {
+                    f.call_or_construct(this_ref, &args, &JsValue::undefined(), context)
                 }
-                AsyncRunningState::Calling => f.poll(context),
-                AsyncRunningState::Constructing { new_target: _ } => {
-                    unreachable!()
+                AsyncCallState::Calling { .. } => AsyncCallResult::Pending,
+                AsyncCallState::Finished(_) => {
+                    let AsyncCallState::Finished(result) = std::mem::take(&mut context.async_call)
+                    else {
+                        unreachable!()
+                    };
+                    AsyncCallResult::Ready(result)
                 }
+            };
+            if matches!(result, AsyncCallResult::Pending) {
+                context.vm.stack.push(this);
+                context.vm.stack.push(func);
+                context.vm.stack.calling_convention_push_arguments(&args);
             }
+            result
         } else {
-            let args = context
-                .vm
-                .stack
-                .calling_convention_pop_arguments(argument_count);
-            let _func = context.vm.stack.pop();
-            let this = context.vm.stack.pop();
-            let this = if is_constructor {
-                JsValue::undefined()
-            } else {
-                this
-            };
             AsyncCallResult::Ready(match self.inner {
                 Inner::PointerFn(f) => f(&this, &args, context),
                 Inner::Closure(ref c) => c.call(&this, &args, context),
@@ -661,43 +582,42 @@ impl NativeFunction {
         argument_count: usize,
         context: &mut Context,
     ) -> (JsValue, AsyncCallResult) {
-        if let Inner::AsyncFn(ref f) = self.inner {
-            match f.running() {
-                AsyncRunningState::None => {
-                    let new_target = context.vm.stack.pop();
-                    let args = context
-                        .vm
-                        .stack
-                        .calling_convention_pop_arguments(argument_count);
-                    let _func = context.vm.stack.pop();
-                    let _this = context.vm.stack.pop();
-                    (
-                        new_target.clone(),
-                        f.call_or_construct(&JsValue::undefined(), &args, &new_target, context),
-                    )
-                }
-                AsyncRunningState::Constructing { new_target } => {
-                    (new_target.clone(), f.poll(context))
+        let new_target = context.vm.stack.pop();
+        let args = context
+            .vm
+            .stack
+            .calling_convention_pop_arguments(argument_count);
+        let func = context.vm.stack.pop();
+        let this = context.vm.stack.pop();
+        let result = if let Inner::AsyncFn(ref f) = self.inner {
+            let result = match &context.async_call {
+                AsyncCallState::None => {
+                    f.call_or_construct(&JsValue::undefined(), &args, &new_target, context)
                 }
-                AsyncRunningState::Calling => {
-                    unreachable!()
+                AsyncCallState::Calling { .. } => AsyncCallResult::Pending,
+                AsyncCallState::Finished(_) => {
+                    let AsyncCallState::Finished(result) = std::mem::take(&mut context.async_call)
+                    else {
+                        unreachable!()
+                    };
+                    AsyncCallResult::Ready(result)
                 }
+            };
+            if matches!(result, AsyncCallResult::Pending) {
+                context.vm.stack.push(this);
+                context.vm.stack.push(func);
+                context.vm.stack.calling_convention_push_arguments(&args);
+                context.vm.stack.push(new_target.clone());
             }
+            result
         } else {
-            let new_target = context.vm.stack.pop();
-            let args = context
-                .vm
-                .stack
-                .calling_convention_pop_arguments(argument_count);
-            let _func = context.vm.stack.pop();
-            let _this = context.vm.stack.pop();
-            let result = AsyncCallResult::Ready(match self.inner {
+            AsyncCallResult::Ready(match self.inner {
                 Inner::PointerFn(f) => f(&new_target, &args, context),
                 Inner::Closure(ref c) => c.call(&new_target, &args, context),
                 Inner::AsyncFn(_) => unreachable!(),
-            });
-            (new_target.clone(), result)
-        }
+            })
+        };
+        (new_target.clone(), result)
     }
 
     /// Converts this `NativeFunction` into a `JsFunction` without setting its name or length.

core/engine/src/vm/mod.rs

@@ -7,6 +7,7 @@
 use crate::{
     builtins::promise::{PromiseCapability, ResolvingFunctions},
     environments::EnvironmentStack,
+    native_function::AsyncCallState,
     object::JsFunction,
     realm::Realm,
     script::Script,
@@ -819,7 +820,15 @@ impl Context {
 
             if runtime_budget == 0 {
                 runtime_budget = budget;
-                yield_now().await;
+                match &mut self.async_call {
+                    AsyncCallState::None | AsyncCallState::Finished(_) => {
+                        yield_now().await;
+                    }
+                    AsyncCallState::Calling { f } => {
+                        let result = f.await;
+                        self.async_call = AsyncCallState::Finished(result);
+                    }
+                }
             }
         }
 
@@ -846,6 +855,7 @@ impl Context {
                 ControlFlow::Continue(_) => {}
                 ControlFlow::Break(value) => return value,
             }
+            debug_assert!(matches!(self.async_call, AsyncCallState::None));
         }
 
         CompletionRecord::Throw(JsError::from_native(JsNativeError::error()))

core/engine/src/vm/opcode/call/mod.rs

@@ -197,8 +197,7 @@ impl Call {
             .__call__(argument_count.into())
             .async_resolve(context)?
         {
-            ResolvedCallValue::Ready => Ok(OpStatus::Finished),
-            ResolvedCallValue::Complete => Ok(OpStatus::Finished),
+            ResolvedCallValue::Ready | ResolvedCallValue::Complete => Ok(OpStatus::Finished),
             ResolvedCallValue::Pending => {
                 //println!("Pending call");
                 Ok(OpStatus::Pending)