Debugging: apply single-stepping patches to modules instantiated after setting changes. (#12663)

We currently do a store-wide state-change on all registered modules
when the "single stepping" flag changes, patching in or out all
breakpoints. However, this design didn't account for modules
registered with the store *after* single-stepping is enabled (and new
modules may be registered any time an instantiation occurs). In
particular this is problematic when a debugger, e.g., sets the
single-step flag right at the beginning of execution of a "host main
function" that calls Wasm, before the main instantiation occurs.

This PR threads through the breakpoint state in the registration path,
with the narrow waist at `ModuleRegistry::register` which is the only
place where modules are added to the `LoadedCode`.

Author: cfallin

crates/wasmtime/src/runtime/component/instance.rs

@@ -724,8 +724,8 @@ impl<'a> Instantiator<'a> {
         imports: &'a Arc<PrimaryMap<RuntimeImportIndex, RuntimeImport>>,
     ) -> Result<Instantiator<'a>> {
         let env_component = component.env_component();
-        let (modules, engine) = store.modules_and_engine_mut();
-        modules.register_component(component, engine)?;
+        let (modules, engine, breakpoints) = store.modules_and_engine_and_breakpoints_mut();
+        modules.register_component(component, engine, breakpoints)?;
         let imported_resources: ImportedResources =
             PrimaryMap::with_capacity(env_component.imported_resources.len());
 

crates/wasmtime/src/runtime/debug.rs

@@ -1,6 +1,8 @@
 //! Debugging API.
 
 use super::store::AsStoreOpaque;
+use crate::code::StoreCode;
+use crate::module::RegisterBreakpointState;
 use crate::store::StoreId;
 use crate::vm::{Activation, Backtrace};
 use crate::{
@@ -1007,15 +1009,34 @@ impl BreakpointState {
     pub(crate) fn is_single_step(&self) -> bool {
         self.single_step
     }
+
+    /// Internal helper to patch a new module for
+    /// single-stepping. When a module is newly registered in a
+    /// `Store`, we need to patch all breakpoints into the copy for
+    /// this `Store` if single-stepping is currently enabled.
+    pub(crate) fn patch_new_module(&self, code: &mut StoreCode, module: &Module) -> Result<()> {
+        // Apply single-step state if single-stepping is enabled. Note
+        // that no other individual breakpoints will exist yet (as
+        // this is a newly registered module).
+        if self.single_step {
+            let mem = code.code_memory_mut().unwrap();
+            mem.unpublish()?;
+            BreakpointEdit::apply_single_step(mem, module, true, |_key| false)?;
+            mem.publish()?;
+        }
+        Ok(())
+    }
 }
 
 impl<'a> BreakpointEdit<'a> {
     fn get_code_memory<'b>(
+        breakpoints: &BreakpointState,
         registry: &'b mut ModuleRegistry,
         dirty_modules: &mut BTreeSet<StoreCodePC>,
         module: &Module,
     ) -> Result<&'b mut CodeMemory> {
-        let store_code_pc = registry.store_code_base_or_register(module)?;
+        let store_code_pc =
+            registry.store_code_base_or_register(module, RegisterBreakpointState(breakpoints))?;
         let code_memory = registry
             .store_code_mut(store_code_pc)
             .expect("Just checked presence above")
@@ -1052,7 +1073,8 @@ impl<'a> BreakpointEdit<'a> {
         let key = BreakpointKey::from_raw(module, pc);
         self.state.breakpoints.insert(key);
         log::trace!("patching in breakpoint {key:?}");
-        let mem = Self::get_code_memory(self.registry, &mut self.dirty_modules, module)?;
+        let mem =
+            Self::get_code_memory(self.state, self.registry, &mut self.dirty_modules, module)?;
         let frame_table = module
             .frame_table()
             .expect("Frame table must be present when guest-debug is enabled");
@@ -1069,7 +1091,8 @@ impl<'a> BreakpointEdit<'a> {
         let key = BreakpointKey::from_raw(module, pc);
         self.state.breakpoints.remove(&key);
         if !self.state.single_step {
-            let mem = Self::get_code_memory(self.registry, &mut self.dirty_modules, module)?;
+            let mem =
+                Self::get_code_memory(self.state, self.registry, &mut self.dirty_modules, module)?;
             let frame_table = module
                 .frame_table()
                 .expect("Frame table must be present when guest-debug is enabled");
@@ -1079,6 +1102,26 @@ impl<'a> BreakpointEdit<'a> {
         Ok(())
     }
 
+    fn apply_single_step<F: Fn(&BreakpointKey) -> bool>(
+        mem: &mut CodeMemory,
+        module: &Module,
+        enabled: bool,
+        key_enabled: F,
+    ) -> Result<()> {
+        let table = module
+            .frame_table()
+            .expect("Frame table must be present when guest-debug is enabled");
+        for (wasm_pc, patch) in table.breakpoint_patches() {
+            let key = BreakpointKey::from_raw(&module, wasm_pc);
+            let this_enabled = enabled || key_enabled(&key);
+            log::trace!(
+                "single_step: enabled {enabled} key {key:?} -> this_enabled {this_enabled}"
+            );
+            Self::patch(core::iter::once(patch), mem, this_enabled);
+        }
+        Ok(())
+    }
+
     /// Turn on or off single-step mode.
     ///
     /// In single-step mode, a breakpoint event is emitted at every
@@ -1095,18 +1138,11 @@ impl<'a> BreakpointEdit<'a> {
         }
         let modules = self.registry.all_modules().cloned().collect::<Vec<_>>();
         for module in modules {
-            let mem = Self::get_code_memory(self.registry, &mut self.dirty_modules, &module)?;
-            let table = module
-                .frame_table()
-                .expect("Frame table must be present when guest-debug is enabled");
-            for (wasm_pc, patch) in table.breakpoint_patches() {
-                let key = BreakpointKey::from_raw(&module, wasm_pc);
-                let this_enabled = enabled || self.state.breakpoints.contains(&key);
-                log::trace!(
-                    "single_step: enabled {enabled} key {key:?} -> this_enabled {this_enabled}"
-                );
-                Self::patch(core::iter::once(patch), mem, this_enabled);
-            }
+            let mem =
+                Self::get_code_memory(self.state, self.registry, &mut self.dirty_modules, &module)?;
+            Self::apply_single_step(mem, &module, enabled, |key| {
+                self.state.breakpoints.contains(key)
+            })?;
         }
 
         self.state.single_step = enabled;

crates/wasmtime/src/runtime/instance.rs

@@ -225,8 +225,8 @@ impl Instance {
         // Note that under normal operation this shouldn't do much as the list
         // of funcs-with-holes should generally be empty. As a result the
         // process of filling this out is not super optimized at this point.
-        let (modules, engine) = store.modules_and_engine_mut();
-        modules.register_module(module, engine)?;
+        let (modules, engine, breakpoints) = store.modules_and_engine_and_breakpoints_mut();
+        modules.register_module(module, engine, breakpoints)?;
         let (funcrefs, modules) = store.func_refs_and_modules();
         funcrefs.fill(modules);
 
@@ -305,8 +305,8 @@ impl Instance {
 
         // Register the module just before instantiation to ensure we keep the module
         // properly referenced while in use by the store.
-        let (modules, engine) = store.modules_and_engine_mut();
-        let module_id = modules.register_module(module, engine)?;
+        let (modules, engine, breakpoints) = store.modules_and_engine_and_breakpoints_mut();
+        let module_id = modules.register_module(module, engine, breakpoints)?;
 
         // The first thing we do is issue an instance allocation request
         // to the instance allocator. This, on success, will give us an
@@ -966,8 +966,8 @@ fn pre_instantiate_raw(
 ) -> Result<OwnedImports> {
     // Register this module and use it to fill out any funcref wasm_call holes
     // we can. For more comments on this see `typecheck_externs`.
-    let (modules, engine) = store.modules_and_engine_mut();
-    modules.register_module(module, engine)?;
+    let (modules, engine, breakpoints) = store.modules_and_engine_and_breakpoints_mut();
+    modules.register_module(module, engine, breakpoints)?;
     let (funcrefs, modules) = store.func_refs_and_modules();
     funcrefs.fill(modules);
 

crates/wasmtime/src/runtime/module/registry.rs

@@ -10,6 +10,8 @@ use crate::{Engine, prelude::*};
 use crate::{FrameInfo, Module, code_memory::CodeMemory};
 use alloc::collections::btree_map::{BTreeMap, Entry};
 use alloc::sync::Arc;
+#[cfg(not(feature = "debug"))]
+use core::marker::PhantomData;
 use core::ops::Range;
 use core::ptr::NonNull;
 use wasmtime_environ::VMSharedTypeIndex;
@@ -100,6 +102,22 @@ fn assert_no_overlap(loaded_code: &BTreeMap<StoreCodePC, LoadedCode>, range: Ran
     }
 }
 
+#[cfg(feature = "debug")]
+pub struct RegisterBreakpointState<'a>(pub(crate) &'a crate::runtime::debug::BreakpointState);
+#[cfg(not(feature = "debug"))]
+pub struct RegisterBreakpointState<'a>(pub(crate) PhantomData<&'a ()>);
+
+impl<'a> RegisterBreakpointState<'a> {
+    #[cfg(feature = "debug")]
+    fn update(&self, code: &mut StoreCode, module: &Module) -> Result<()> {
+        self.0.patch_new_module(code, module)
+    }
+    #[cfg(not(feature = "debug"))]
+    fn update(&self, _code: &mut StoreCode, _module: &Module) -> Result<()> {
+        Ok(())
+    }
+}
+
 impl ModuleRegistry {
     /// Get a previously-registered module by id.
     pub fn module_by_id(&self, id: RegisteredModuleId) -> Option<&Module> {
@@ -140,11 +158,15 @@ impl ModuleRegistry {
 
     /// Fetches the base `StoreCodePC` for a given `EngineCode` with
     /// `Module`, registering the module if not already registered.
-    pub fn store_code_base_or_register(&mut self, module: &Module) -> Result<StoreCodePC> {
+    pub fn store_code_base_or_register(
+        &mut self,
+        module: &Module,
+        breakpoint_state: RegisterBreakpointState,
+    ) -> Result<StoreCodePC> {
         let key = module.engine_code().text_range().start;
         if !self.store_code.contains_key(&key) {
             let engine = module.engine().clone();
-            self.register_module(module, &engine)?;
+            self.register_module(module, &engine, breakpoint_state)?;
         }
         Ok(*self.store_code.get(&key).unwrap())
     }
@@ -167,14 +189,32 @@ impl ModuleRegistry {
         &mut self,
         module: &Module,
         engine: &Engine,
+        breakpoint_state: RegisterBreakpointState,
     ) -> Result<RegisteredModuleId> {
-        self.register(module.id(), module.engine_code(), Some(module), engine)
-            .map(|id| id.unwrap())
+        self.register(
+            module.id(),
+            module.engine_code(),
+            Some(module),
+            engine,
+            breakpoint_state,
+        )
+        .map(|id| id.unwrap())
     }
 
     #[cfg(feature = "component-model")]
-    pub fn register_component(&mut self, component: &Component, engine: &Engine) -> Result<()> {
-        self.register(component.id(), component.engine_code(), None, engine)?;
+    pub fn register_component(
+        &mut self,
+        component: &Component,
+        engine: &Engine,
+        breakpoint_state: RegisterBreakpointState,
+    ) -> Result<()> {
+        self.register(
+            component.id(),
+            component.engine_code(),
+            None,
+            engine,
+            breakpoint_state,
+        )?;
         Ok(())
     }
 
@@ -185,6 +225,7 @@ impl ModuleRegistry {
         code: &Arc<EngineCode>,
         module: Option<&Module>,
         engine: &Engine,
+        breakpoint_state: RegisterBreakpointState,
     ) -> Result<Option<RegisteredModuleId>> {
         // Register the module, if any.
         let id = module.map(|module| {
@@ -219,6 +260,7 @@ impl ModuleRegistry {
                     .get_mut(&store_code_pc)
                     .expect("loaded_code must have entry for StoreCodePC");
                 loaded_code.modules.insert(range.start, id);
+                breakpoint_state.update(&mut loaded_code.code, module)?;
             }
         }
 

crates/wasmtime/src/runtime/store.rs

@@ -84,7 +84,7 @@ use crate::ThrownException;
 use crate::error::OutOfMemory;
 #[cfg(feature = "async")]
 use crate::fiber;
-use crate::module::RegisteredModuleId;
+use crate::module::{RegisterBreakpointState, RegisteredModuleId};
 use crate::prelude::*;
 #[cfg(feature = "gc")]
 use crate::runtime::vm::GcRootsList;
@@ -1670,8 +1670,15 @@ impl StoreOpaque {
     }
 
     #[inline]
-    pub(crate) fn modules_and_engine_mut(&mut self) -> (&mut ModuleRegistry, &Engine) {
-        (&mut self.modules, &self.engine)
+    pub(crate) fn modules_and_engine_and_breakpoints_mut(
+        &mut self,
+    ) -> (&mut ModuleRegistry, &Engine, RegisterBreakpointState<'_>) {
+        #[cfg(feature = "debug")]
+        let breakpoints = RegisterBreakpointState(&self.breakpoints);
+        #[cfg(not(feature = "debug"))]
+        let breakpoints = RegisterBreakpointState(core::marker::PhantomData);
+
+        (&mut self.modules, &self.engine, breakpoints)
     }
 
     pub(crate) fn func_refs_and_modules(&mut self) -> (&mut FuncRefs, &ModuleRegistry) {

tests/all/debug.rs

@@ -1318,3 +1318,57 @@ fn debug_ids() -> wasmtime::Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+#[cfg_attr(miri, ignore)]
+async fn single_step_before_instantiation() -> wasmtime::Result<()> {
+    let _ = env_logger::try_init();
+
+    let mut config = Config::default();
+    config.guest_debug(true);
+    let engine = Engine::new(&config)?;
+    let module = Module::new(
+        &engine,
+        r#"
+    (module
+      (func (export "main") (param i32 i32) (result i32)
+        local.get 0
+        local.get 1
+        i32.add))
+    "#,
+    )?;
+    let mut store = Store::new(&engine, ());
+
+    // Enable single-stepping *before* instantiation. The module has not
+    // been registered with this store yet.
+    store.edit_breakpoints().unwrap().single_step(true).unwrap();
+    assert!(store.is_single_step());
+
+    #[derive(Clone)]
+    struct CountingHandler(Arc<AtomicUsize>);
+    impl DebugHandler for CountingHandler {
+        type Data = ();
+        async fn handle(&self, _store: StoreContextMut<'_, ()>, event: DebugEvent<'_>) {
+            match event {
+                DebugEvent::Breakpoint => {
+                    self.0.fetch_add(1, Ordering::Relaxed);
+                }
+                _ => {}
+            }
+        }
+    }
+
+    let counter = Arc::new(AtomicUsize::new(0));
+    store.set_debug_handler(CountingHandler(counter.clone()));
+
+    let instance = Instance::new_async(&mut store, &module, &[]).await?;
+    let func = instance.get_func(&mut store, "main").unwrap();
+    let mut results = [Val::I32(0)];
+    func.call_async(&mut store, &[Val::I32(1), Val::I32(2)], &mut results)
+        .await?;
+    assert_eq!(results[0].unwrap_i32(), 3);
+
+    assert_eq!(counter.load(Ordering::Relaxed), 4);
+
+    Ok(())
+}