Refactor run_async to custom future implementation

Author: benaubin

lucet-runtime/lucet-runtime-internals/src/future.rs

@@ -7,7 +7,6 @@ pub use crate::instance::execution::{KillError, KillState, KillSuccess, KillSwit
 pub use crate::instance::signals::{signal_handler_none, SignalBehavior, SignalHandler};
 pub use crate::instance::state::State;
 
-use crate::alloc::Alloc;
 use crate::context::Context;
 use crate::embed_ctx::CtxMap;
 use crate::error::Error;
@@ -18,13 +17,13 @@ use crate::region::RegionInternal;
 use crate::sysdeps::HOST_PAGE_SIZE_EXPECTED;
 use crate::val::{UntypedRetVal, Val};
 use crate::WASM_PAGE_SIZE;
+use crate::{alloc::Alloc, future::AsyncContext};
 use libc::{c_void, pthread_self, siginfo_t, uintptr_t};
 use lucet_module::InstanceRuntimeData;
 use memoffset::offset_of;
 use std::any::Any;
 use std::cell::{BorrowError, BorrowMutError, Ref, RefCell, RefMut, UnsafeCell};
 use std::convert::TryFrom;
-use std::marker::PhantomData;
 use std::mem;
 use std::ops::{Deref, DerefMut};
 use std::ptr::{self, NonNull};
@@ -228,6 +227,10 @@ pub struct Instance {
     /// Small mutexed state used for remote kill switch functionality
     pub(crate) kill_state: Arc<KillState>,
 
+    /// Indicates whether the instance is running in an async context (`Instance::run_async`)
+    /// or not. Needed by `Vmctx::block_on`.
+    pub(crate) async_ctx: Option<std::sync::Arc<AsyncContext>>,
+
     #[cfg(feature = "concurrent_testpoints")]
     /// Conditionally-present helpers to force permutations of possible races in testing.
     pub lock_testpoints: Arc<LockTestpoints>,
@@ -515,7 +518,7 @@ impl Instance {
     /// in the future.
     pub fn run(&mut self, entrypoint: &str, args: &[Val]) -> Result<RunResult, Error> {
         let func = self.module.get_export_func(entrypoint)?;
-        Ok(self.run_func(func, &args, false, None)?.unwrap())
+        Ok(self.run_func(func, &args, None, None)?.unwrap())
     }
 
     /// Run a function with arguments in the guest context from the [WebAssembly function
@@ -531,7 +534,7 @@ impl Instance {
         args: &[Val],
     ) -> Result<RunResult, Error> {
         let func = self.module.get_func_from_idx(table_idx, func_idx)?;
-        Ok(self.run_func(func, &args, false, None)?.unwrap())
+        Ok(self.run_func(func, &args, None, None)?.unwrap())
     }
 
     /// Resume execution of an instance that has yielded without providing a value to the guest.
@@ -562,19 +565,21 @@ impl Instance {
     /// The foreign code safety caveat of [`Instance::run()`](struct.Instance.html#method.run)
     /// applies.
     pub fn resume_with_val<A: Any + 'static>(&mut self, val: A) -> Result<RunResult, Error> {
-        Ok(self.resume_with_val_impl(val, false, None)?.unwrap())
+        Ok(self
+            .resume_with_val_impl(Box::new(val), None, None)?
+            .unwrap())
     }
 
-    pub(crate) fn resume_with_val_impl<A: Any + 'static>(
+    pub(crate) fn resume_with_val_impl(
         &mut self,
-        val: A,
-        async_context: bool,
+        val: Box<dyn Any + 'static>,
+        async_context: Option<AsyncContext>,
         max_insn_count: Option<u64>,
     ) -> Result<InternalRunResult, Error> {
         match &self.state {
             State::Yielded { expecting, .. } => {
                 // make sure the resumed value is of the right type
-                if !expecting.is::<PhantomData<A>>() {
+                if &(*val).type_id() != expecting {
                     return Err(Error::InvalidArgument(
                         "type mismatch between yielded instance expected value and resumed value",
                     ));
@@ -583,7 +588,7 @@ impl Instance {
             _ => return Err(Error::InvalidArgument("can only resume a yielded instance")),
         }
 
-        self.resumed_val = Some(Box::new(val) as Box<dyn Any + 'static>);
+        self.resumed_val = Some(val);
 
         self.set_instruction_bound_delta(max_insn_count);
         self.swap_and_return(async_context)
@@ -602,6 +607,7 @@ impl Instance {
     /// applies.
     pub(crate) fn resume_bounded(
         &mut self,
+        async_context: AsyncContext,
         max_insn_count: u64,
     ) -> Result<InternalRunResult, Error> {
         if !self.state.is_bound_expired() {
@@ -610,7 +616,7 @@ impl Instance {
             ));
         }
         self.set_instruction_bound_delta(Some(max_insn_count));
-        self.swap_and_return(true)
+        self.swap_and_return(Some(async_context))
     }
 
     /// Run the module's [start function][start], if one exists.
@@ -648,7 +654,7 @@ impl Instance {
             if !self.is_not_started() {
                 return Err(Error::StartAlreadyRun);
             }
-            self.run_func(start, &[], false, None)?;
+            self.run_func(start, &[], None, None)?;
         }
         Ok(())
     }
@@ -1021,6 +1027,7 @@ impl Instance {
             entrypoint: None,
             resumed_val: None,
             terminate_on_heap_oom: false,
+            async_ctx: None,
             _padding: (),
         };
         inst.set_globals_ptr(globals_ptr);
@@ -1090,7 +1097,7 @@ impl Instance {
         &mut self,
         func: FunctionHandle,
         args: &[Val],
-        async_context: bool,
+        async_context: Option<AsyncContext>,
         inst_count_bound: Option<u64>,
     ) -> Result<InternalRunResult, Error> {
         let needs_start = self.state.is_not_started() && !func.is_start_func;
@@ -1191,7 +1198,10 @@ impl Instance {
     /// This must only be called for an instance in a ready, non-fatally faulted, or yielded state,
     /// or in the not-started state on the start function. The public wrappers around this function
     /// should make sure the state is appropriate.
-    fn swap_and_return(&mut self, async_context: bool) -> Result<InternalRunResult, Error> {
+    fn swap_and_return<'a>(
+        &mut self,
+        async_context: Option<AsyncContext>,
+    ) -> Result<InternalRunResult, Error> {
         let is_start_func = self
             .entrypoint
             .expect("we always have an entrypoint by now")
@@ -1203,7 +1213,10 @@ impl Instance {
                 || self.state.is_yielded()
                 || self.state.is_bound_expired()
         );
-        self.state = State::Running { async_context };
+
+        self.async_ctx = async_context.map(|cx| Arc::new(cx));
+
+        self.state = State::Running;
 
         let res = self.with_current_instance(|i| {
             i.with_signals_on(|i| {
@@ -1217,6 +1230,9 @@ impl Instance {
             })
         });
 
+        // remove async ctx
+        self.async_ctx.take();
+
         #[cfg(feature = "concurrent_testpoints")]
         self.lock_testpoints
             .instance_after_clearing_current_instance

lucet-runtime/lucet-runtime-internals/src/instance.rs

@@ -2,7 +2,7 @@ use crate::instance::siginfo_ext::SiginfoExt;
 use crate::instance::{FaultDetails, TerminationDetails, YieldedVal};
 use crate::sysdeps::UContext;
 use libc::{SIGBUS, SIGSEGV};
-use std::any::Any;
+use std::any::TypeId;
 use std::ffi::{CStr, CString};
 
 /// The representation of a Lucet instance's state machine.
@@ -22,11 +22,7 @@ pub enum State {
     /// Transitions to `Ready` when the guest function returns normally, or to `Faulted`,
     /// `Terminating`, or `Yielding` if the instance faults, terminates, or yields, or to
     /// `BoundExpired` if the instance is run with an instruction bound and reaches it.
-    Running {
-        /// Indicates whether the instance is running in an async context (`Instance::run_async`)
-        /// or not. Needed by `Vmctx::block_on`.
-        async_context: bool,
-    },
+    Running,
 
     /// The instance has faulted, potentially fatally.
     ///
@@ -56,11 +52,8 @@ pub enum State {
     /// `RunResult` before anything else happens to the instance.
     Yielding {
         val: YieldedVal,
-        /// A phantom value carrying the type of the expected resumption value.
-        ///
-        /// Concretely, this should only ever be `Box<PhantomData<R>>` where `R` is the type
-        /// the guest expects upon resumption.
-        expecting: Box<dyn Any>,
+        /// The type of the expected resumption value
+        expecting: TypeId,
     },
 
     /// The instance has yielded.
@@ -69,10 +62,7 @@ pub enum State {
     /// instance is reset.
     Yielded {
         /// A phantom value carrying the type of the expected resumption value.
-        ///
-        /// Concretely, this should only ever be `Box<PhantomData<R>>` where `R` is the type
-        /// the guest expects upon resumption.
-        expecting: Box<dyn Any>,
+        expecting: TypeId,
     },
 
     /// The instance has reached an instruction-count bound.
@@ -96,12 +86,7 @@ impl std::fmt::Display for State {
         match self {
             State::NotStarted => write!(f, "not started"),
             State::Ready => write!(f, "ready"),
-            State::Running {
-                async_context: false,
-            } => write!(f, "running"),
-            State::Running {
-                async_context: true,
-            } => write!(f, "running (in async context)"),
+            State::Running => write!(f, "running"),
             State::Faulted {
                 details, siginfo, ..
             } => {
@@ -162,14 +147,6 @@ impl State {
         }
     }
 
-    pub fn is_running_async(&self) -> bool {
-        if let State::Running { async_context } = self {
-            *async_context
-        } else {
-            false
-        }
-    }
-
     pub fn is_faulted(&self) -> bool {
         if let State::Faulted { .. } = self {
             true

lucet-runtime/lucet-runtime-internals/src/instance/state.rs

@@ -13,10 +13,9 @@ use crate::instance::{
     CURRENT_INSTANCE, HOST_CTX,
 };
 use lucet_module::{FunctionHandle, GlobalValue};
-use std::any::Any;
+use std::any::{Any, TypeId};
 use std::borrow::{Borrow, BorrowMut};
 use std::cell::{Ref, RefCell, RefMut};
-use std::marker::PhantomData;
 
 /// An opaque handle to a running instance's context.
 #[derive(Debug)]
@@ -436,10 +435,9 @@ impl Vmctx {
         if is_bound_expiration {
             inst.state = State::BoundExpired;
         } else {
-            let expecting: Box<PhantomData<R>> = Box::new(PhantomData);
             inst.state = State::Yielding {
                 val: YieldedVal::new(val),
-                expecting: expecting as Box<dyn Any>,
+                expecting: TypeId::of::<R>(),
             };
         }
 

lucet-runtime/lucet-runtime-internals/src/vmctx.rs

@@ -101,6 +101,17 @@ pub fn lucet_hostcall(_attr: TokenStream, item: TokenStream) -> TokenStream {
         quote! { lucet_runtime::TerminationDetails }
     };
 
+    let res_ident = quote::format_ident!("res");
+
+    let block_if_async = match raw_sig.asyncness.take() {
+        Some(_) => {
+            quote! { let #res_ident = vmctx.block_on(#res_ident); }
+        }
+        None => {
+            quote! {}
+        }
+    };
+
     let raw_hostcall = quote! {
         #(#attrs)*
         #vis
@@ -111,7 +122,13 @@ pub fn lucet_hostcall(_attr: TokenStream, item: TokenStream) -> TokenStream {
             let vmctx = #vmctx_mod::Vmctx::from_raw(vmctx_raw);
             #vmctx_mod::VmctxInternal::instance_mut(&vmctx).uninterruptable(|| {
                 let res = std::panic::catch_unwind(move || {
-                    #hostcall_ident(&#vmctx_mod::Vmctx::from_raw(vmctx_raw), #(#impl_args),*)
+                    let vmctx = #vmctx_mod::Vmctx::from_raw(vmctx_raw);
+
+                    let #res_ident = #hostcall_ident(&vmctx, #(#impl_args),*);
+
+                    #block_if_async
+
+                    #res_ident
                 });
                 match res {
                     Ok(res) => res,

lucet-runtime/lucet-runtime-macros/src/lib.rs

@@ -1,5 +1,8 @@
 {
     "env": {
-        "hostcall_containing_block_on": "hostcall_containing_block_on"
+        "hostcall_containing_block_on": "hostcall_containing_block_on",
+        "hostcall_containing_yielding_block_on": "hostcall_containing_yielding_block_on",
+        "hostcall_async_containing_yielding_block_on": "hostcall_async_containing_yielding_block_on",
+        "await_manual_future": "await_manual_future"
     }
 }

lucet-runtime/lucet-runtime-tests/guests/async_hostcall/bindings.json

@@ -1,9 +1,30 @@
 #include <stddef.h>
 
 extern void hostcall_containing_block_on(int);
+extern void hostcall_containing_yielding_block_on(int);
+extern int hostcall_async_containing_yielding_block_on(int, int);
 
 int main(void)
 {
     hostcall_containing_block_on(1312);
     return 0;
 }
+
+int yielding()
+{
+    hostcall_containing_yielding_block_on(0);
+    hostcall_containing_yielding_block_on(1);
+    hostcall_containing_yielding_block_on(2);
+    hostcall_containing_yielding_block_on(3);
+
+    int six = hostcall_async_containing_yielding_block_on(3, 6);
+    hostcall_async_containing_yielding_block_on(3, six);
+
+    return 0;
+}
+
+int manual_future()
+{
+    await_manual_future();
+    return 0;
+}