Add support for core re-entrancy + test

Author: arjunr2

crates/wasmtime/src/runtime/func.rs

@@ -1388,7 +1388,7 @@ impl Func {
                 flat_params,
                 &mut caller.store.0,
             )?;
-            rr::core_hooks::replay_host_func_return(values_vec, &mut caller.store.0)?;
+            rr::core_hooks::replay_host_func_return(values_vec, &mut caller)?;
         }
 
         // Restore our `val_vec` back into the store so it's usable for the next
@@ -2544,7 +2544,7 @@ impl HostContext {
                 // Replay the return values
                 rr::core_hooks::replay_host_func_return(
                     unsafe { &mut args.as_mut()[..num_results] },
-                    caller.store.0,
+                    &mut caller,
                 )?;
             }
 

crates/wasmtime/src/runtime/rr/core.rs

@@ -346,22 +346,6 @@ pub trait Replayer: Iterator<Item = Result<RREvent, ReplayError>> {
         T::try_from(self.next_event()?).map_err(|e| e.into())
     }
 
-    /// Pop the next replay event and calls `f` with a desired type conversion
-    ///
-    /// ## Errors
-    ///
-    /// See [`next_event_typed`](Replayer::next_event_typed)
-    #[inline]
-    fn next_event_and<T, F>(&mut self, f: F) -> Result<(), ReplayError>
-    where
-        T: TryFrom<RREvent>,
-        ReplayError: From<<T as TryFrom<RREvent>>::Error>,
-        F: FnOnce(T) -> Result<(), ReplayError>,
-    {
-        let call_event = self.next_event_typed()?;
-        Ok(f(call_event)?)
-    }
-
     /// Conditionally process the next validation recorded event and if
     /// replay validation is enabled, run the validation check
     ///
@@ -591,6 +575,18 @@ mod tests {
     use tempfile::{NamedTempFile, TempPath};
     use wasmtime_environ::FuncIndex;
 
+    impl ReplayBuffer {
+        fn next_event_and<T, F>(&mut self, f: F) -> Result<(), ReplayError>
+        where
+            T: TryFrom<RREvent>,
+            ReplayError: From<<T as TryFrom<RREvent>>::Error>,
+            F: FnOnce(T) -> Result<(), ReplayError>,
+        {
+            let event = self.next_event_typed::<T>()?;
+            f(event)
+        }
+    }
+
     fn rr_harness<S, T>(record_fn: S, replay_fn: T) -> Result<()>
     where
         S: FnOnce(&mut RecordBuffer) -> Result<()>,

crates/wasmtime/src/runtime/rr/hooks/core_hooks.rs

@@ -1,12 +1,14 @@
 use crate::rr::FlatBytes;
 #[cfg(feature = "rr")]
 use crate::rr::{
-    RRFuncArgVals, ResultEvent, common_events::HostFuncEntryEvent,
-    common_events::HostFuncReturnEvent, common_events::WasmFuncReturnEvent,
-    core_events::WasmFuncEntryEvent,
+    RREvent, RRFuncArgVals, ReplayError, ReplayHostContext, Replayer, ResultEvent,
+    common_events::HostFuncEntryEvent, common_events::HostFuncReturnEvent,
+    common_events::WasmFuncReturnEvent, core_events::WasmFuncEntryEvent,
 };
 use crate::store::StoreOpaque;
-use crate::{FuncType, StoreContextMut, ValRaw, WasmFuncOrigin, prelude::*};
+use crate::{Caller, FuncType, StoreContextMut, ValRaw, WasmFuncOrigin, prelude::*};
+#[cfg(feature = "rr")]
+use wasmtime_environ::EntityIndex;
 
 /// Record and replay hook operation for core wasm function entry events
 ///
@@ -121,15 +123,74 @@ where
 
 /// Replay hook operation for host function return events
 #[inline]
-pub fn replay_host_func_return<T>(args: &mut [T], store: &mut StoreOpaque) -> Result<()>
+pub fn replay_host_func_return<T, U: 'static>(
+    args: &mut [T],
+    caller: &mut Caller<'_, U>,
+) -> Result<()>
 where
     T: FlatBytes,
 {
     #[cfg(feature = "rr")]
-    store.next_replay_event_and(|event: HostFuncReturnEvent| {
-        event.args.into_raw_slice(args);
-        Ok(())
-    })?;
-    let _ = (args, store);
+    {
+        // Core wasm can be re-entrant, so we need to check for this
+        let mut complete = false;
+        while !complete {
+            let buf = caller.store.0.replay_buffer_mut().unwrap();
+            let event = buf.next_event()?;
+            match event {
+                RREvent::HostFuncReturn(event) => {
+                    event.args.into_raw_slice(args);
+                    complete = true;
+                }
+                // Re-entrant call into wasm function: this resembles the implementation in [`ReplayInstance`]
+                RREvent::CoreWasmFuncEntry(event) => {
+                    let entity = EntityIndex::from(event.origin.index);
+
+                    // SAFETY: The store's data is always of type `ReplayHostContext<T>` when created by
+                    // the replay driver. As an additional guarantee, we assert that replay is indeed
+                    // truly enabled.
+                    assert!(caller.store.0.replay_enabled());
+                    let replay_data = unsafe {
+                        let raw_ptr: *const U = caller.store.data();
+                        &*(raw_ptr as *const ReplayHostContext)
+                    };
+
+                    // Grab the correct module instance
+                    let instance = replay_data.get_module_instance(event.origin.instance)?;
+
+                    let mut store = &mut caller.store;
+                    let func = instance
+                        ._get_export(store.0, entity)
+                        .into_func()
+                        .ok_or(ReplayError::InvalidCoreFuncIndex(entity))?;
+
+                    let params_ty = func.ty(&store).params().collect::<Vec<_>>();
+
+                    // Obtain the argument values for function call
+                    let mut results = vec![crate::Val::I64(0); func.ty(&store).results().len()];
+                    let params = event.args.to_val_vec(&mut store, params_ty);
+
+                    // Call the function
+                    //
+                    // This is almost a mirror of the usage in [`crate::Func::call_impl`]
+                    func.call_impl_check_args(&mut store, &params, &mut results)?;
+                    unsafe {
+                        func.call_impl_do_call(
+                            &mut store,
+                            params.as_slice(),
+                            results.as_mut_slice(),
+                        )?;
+                    }
+                }
+                _ => {
+                    bail!(
+                        "Unexpected event during core wasm host function replay: {:?}",
+                        event
+                    );
+                }
+            }
+        }
+    }
+    let _ = (args, caller);
     Ok(())
 }

crates/wasmtime/src/runtime/rr/replay_driver.rs

@@ -43,25 +43,73 @@ impl ReplayEnvironment {
     }
 
     /// Instantiate a new [`ReplayInstance`] using a [`ReplayReader`] in context of this environment
-    pub fn instantiate(&self, reader: impl ReplayReader + 'static) -> Result<ReplayInstance<()>> {
-        let store = Store::new(&self.engine, ());
-        ReplayInstance::<()>::from_environment_and_store(self.clone(), store, reader)
+    pub fn instantiate(&self, reader: impl ReplayReader + 'static) -> Result<ReplayInstance> {
+        self.instantiate_with(reader, |_| Ok(()), |_| Ok(()), |_| Ok(()))
     }
 
-    /// Like [`Self::instantiate`] but allows providing a custom [`Store`] generator
-    pub fn instantiate_with_store<T>(
+    /// Like [`Self::instantiate`] but allows providing a custom modifier functions for
+    /// [`Store`], [`crate::Linker`], and [`component::Linker`] within the replay
+    pub fn instantiate_with(
         &self,
-        store_gen: impl FnOnce() -> Store<T>,
         reader: impl ReplayReader + 'static,
-    ) -> Result<ReplayInstance<T>> {
-        ReplayInstance::from_environment_and_store(self.clone(), store_gen(), reader)
+        store_fn: impl FnOnce(&mut Store<ReplayHostContext>) -> Result<()>,
+        module_linker_fn: impl FnOnce(&mut crate::Linker<ReplayHostContext>) -> Result<()>,
+        component_linker_fn: impl FnOnce(&mut component::Linker<ReplayHostContext>) -> Result<()>,
+    ) -> Result<ReplayInstance> {
+        let mut store = Store::new(
+            &self.engine,
+            ReplayHostContext {
+                module_instances: BTreeMap::new(),
+            },
+        );
+        store_fn(&mut store)?;
+        store.init_replaying(reader, self.settings.clone())?;
+
+        ReplayInstance::from_environment_and_store(
+            self.clone(),
+            store,
+            module_linker_fn,
+            component_linker_fn,
+        )
+    }
+}
+
+/// The host context tied to the store during replay.
+///
+/// This context encapsulates the state from the replay environment that are
+/// required to be accessible within the Store. This is an opaque type from the
+/// public API perspective.
+pub struct ReplayHostContext {
+    /// A tracker of instantiated modules.
+    ///
+    /// Core wasm modules can be re-entrant and invoke methods from other instances, and this
+    /// needs to be accessible within host functions
+    module_instances: BTreeMap<InstanceId, crate::Instance>,
+}
+
+impl ReplayHostContext {
+    /// Insert a module instance into the context's tracking map.
+    fn insert_module_instance(&mut self, id: InstanceId, instance: crate::Instance) {
+        self.module_instances.insert(id, instance);
+    }
+
+    /// Get a module instance from the context's tracking map
+    ///
+    /// This is necessary for core wasm to identify re-entrant calls during replay.
+    pub(crate) fn get_module_instance(
+        &self,
+        id: InstanceId,
+    ) -> Result<&crate::Instance, ReplayError> {
+        self.module_instances
+            .get(&id)
+            .ok_or(ReplayError::MissingModuleInstance(id.as_u32()))
     }
 }
 
-/// A [`ReplayInstance`] is an object providing a opaquely managed, replayable [`Store`]
+/// A [`ReplayInstance`] is an object providing a opaquely managed, replayable [`Store`].
 ///
 /// Debugger capabilities in the future will interact with this object for
-/// inserting breakpoints, snapshotting, and restoring state
+/// inserting breakpoints, snapshotting, and restoring state.
 ///
 /// # Example
 ///
@@ -122,32 +170,36 @@ impl ReplayEnvironment {
 /// # Ok(())
 /// # }
 /// ```
-pub struct ReplayInstance<T: 'static> {
+pub struct ReplayInstance {
     env: Arc<ReplayEnvironment>,
-    store: Store<T>,
-    component_linker: component::Linker<T>,
-    module_linker: crate::Linker<T>,
+    store: Store<ReplayHostContext>,
+    component_linker: component::Linker<ReplayHostContext>,
+    module_linker: crate::Linker<ReplayHostContext>,
     module_instances: BTreeMap<InstanceId, crate::Instance>,
     component_instances: BTreeMap<ComponentInstanceId, component::Instance>,
 }
 
-impl<T: 'static> ReplayInstance<T> {
+impl ReplayInstance {
     fn from_environment_and_store(
         env: ReplayEnvironment,
-        mut store: Store<T>,
-        reader: impl ReplayReader + 'static,
+        store: Store<ReplayHostContext>,
+        module_linker_fn: impl FnOnce(&mut crate::Linker<ReplayHostContext>) -> Result<()>,
+        component_linker_fn: impl FnOnce(&mut component::Linker<ReplayHostContext>) -> Result<()>,
     ) -> Result<Self> {
         let env = Arc::new(env);
-        store.init_replaying(reader, env.settings.clone())?;
-        let mut module_linker = crate::Linker::<T>::new(&env.engine);
+        let mut module_linker = crate::Linker::<ReplayHostContext>::new(&env.engine);
         // Replays shouldn't use any imports, so stub them all out as traps
         for module in env.modules.values() {
             module_linker.define_unknown_imports_as_traps(module)?;
         }
-        let mut component_linker = component::Linker::<T>::new(&env.engine);
+        module_linker_fn(&mut module_linker)?;
+
+        let mut component_linker = component::Linker::<ReplayHostContext>::new(&env.engine);
         for component in env.components.values() {
             component_linker.define_unknown_imports_as_traps(component)?;
         }
+        component_linker_fn(&mut component_linker)?;
+
         Ok(Self {
             env,
             store,
@@ -158,23 +210,23 @@ impl<T: 'static> ReplayInstance<T> {
         })
     }
 
-    /// Obtain a reference to the internal [`Store`]
-    pub fn store(&self) -> &Store<T> {
+    /// Obtain a reference to the internal [`Store`].
+    pub fn store(&self) -> &Store<ReplayHostContext> {
         &self.store
     }
 
-    /// Consume the [`ReplayInstance`] and extract the internal [`Store`]
-    pub fn extract_store(self) -> Store<T> {
+    /// Consume the [`ReplayInstance`] and extract the internal [`Store`].
+    pub fn extract_store(self) -> Store<ReplayHostContext> {
         self.store
     }
 
-    /// Run a single top-level event from the instance
+    /// Run a single top-level event from the instance.
     ///
     /// "Top-level" events are those explicitly invoked events, namely:
     /// * Instantiation events (component/module)
     /// * Wasm function begin events (`ComponentWasmFuncBegin` for components and `CoreWasmFuncEntry` for core)
     ///
-    /// All other events are transparently dispatched under the context of these top-level events
+    /// All other events are transparently dispatched under the context of these top-level events.
     fn run_single_top_level_event(&mut self, rr_event: RREvent) -> Result<()> {
         match rr_event {
             RREvent::ComponentInstantiation(event) => {
@@ -281,7 +333,12 @@ impl<T: 'static> ReplayInstance<T> {
                     instance: instance.id(),
                 })?;
 
+                // Insert into host context tracking as well
                 self.module_instances.insert(instance.id(), instance);
+                self.store
+                    .as_context_mut()
+                    .data_mut()
+                    .insert_module_instance(instance.id(), instance);
             }
             RREvent::CoreWasmFuncEntry(event) => {
                 // Grab the correct module instance
@@ -321,12 +378,9 @@ impl<T: 'static> ReplayInstance<T> {
         Ok(())
     }
 
-    /// Exactly like [`Self::run_single_top_level_event`] but uses async stores and calls
+    /// Exactly like [`Self::run_single_top_level_event`] but uses async stores and calls.
     #[cfg(feature = "async")]
-    async fn run_single_top_level_event_async(&mut self, rr_event: RREvent) -> Result<()>
-    where
-        T: Send,
-    {
+    async fn run_single_top_level_event_async(&mut self, rr_event: RREvent) -> Result<()> {
         match rr_event {
             RREvent::ComponentInstantiation(event) => {
                 // Find matching component from environment to instantiate
@@ -441,6 +495,10 @@ impl<T: 'static> ReplayInstance<T> {
                 })?;
 
                 self.module_instances.insert(instance.id(), instance);
+                self.store
+                    .as_context_mut()
+                    .data_mut()
+                    .insert_module_instance(instance.id(), instance);
             }
             RREvent::CoreWasmFuncEntry(event) => {
                 // Grab the correct module instance
@@ -500,10 +558,7 @@ impl<T: 'static> ReplayInstance<T> {
 
     /// Exactly like [`Self::run_to_completion`] but uses async stores and calls
     #[cfg(feature = "async")]
-    pub async fn run_to_completion_async(&mut self) -> Result<()>
-    where
-        T: Send,
-    {
+    pub async fn run_to_completion_async(&mut self) -> Result<()> {
         while let Some(rr_event) = self
             .store
             .as_context_mut()

crates/wasmtime/src/runtime/store.rs

@@ -1965,24 +1965,6 @@ impl StoreOpaque {
         }
     }
 
-    /// Process the next replay event from the store's replay buffer
-    ///
-    /// Convenience wrapper around [`Replayer::next_event_and`]
-    #[cfg(feature = "rr")]
-    #[inline]
-    pub(crate) fn next_replay_event_and<T, F>(&mut self, f: F) -> Result<(), ReplayError>
-    where
-        T: TryFrom<RREvent>,
-        ReplayError: From<<T as TryFrom<RREvent>>::Error>,
-        F: FnOnce(T) -> Result<(), ReplayError>,
-    {
-        if let Some(buf) = self.replay_buffer_mut() {
-            buf.next_event_and(f)
-        } else {
-            Ok(())
-        }
-    }
-
     /// Process the next replay event as a validation event from the store's replay buffer
     /// and if validation is enabled on replay, and run the validation check
     ///