wasmtime/runtime/module/
registry.rs

1//! Implements a registry of modules for a store.
2
3use crate::code::{EngineCode, EngineCodePC, ModuleWithCode, StoreCode, StoreCodePC};
4#[cfg(feature = "component-model")]
5use crate::component::Component;
6use crate::runtime::vm::VMWasmCallFunction;
7use crate::sync::{OnceLock, RwLock};
8use crate::vm::CompiledModuleId;
9use crate::{Engine, FrameInfo, Module, code_memory::CodeMemory, prelude::*};
10use alloc::sync::Arc;
11#[cfg(not(feature = "debug"))]
12use core::marker::PhantomData;
13use core::ops::Range;
14use core::ptr::NonNull;
15use wasmtime_environ::{
16    CompiledFunctionsTable, FuncKey, StaticModuleIndex, VMSharedTypeIndex,
17    collections::btree_map::Entry,
18};
19
20/// Used for registering modules with a store.
21///
22/// There are two basic purposes that this registry serves:
23///
24/// - It keeps all modules and their metadata alive as long as the
25///   store exists.
26/// - It owns the [`StoreCode`], i.e. possibly-private-copy of machine
27///   code, for all modules that execute in this store.
28///
29/// The registry allows for translation of EngineCode to StoreCode,
30/// deduplicating by the start address of the EngineCode; and allows
31/// for looking up modules by "registered module ID", and looking up
32/// StoreCode and Modules by PC.
33///
34/// Note that multiple modules may be backed by a single
35/// `StoreCode`. This is specifically the case for components in
36/// general. When a component is first instantiated, the component
37/// itself is registered (which loads the StoreCode into the
38/// registry), then each individual module within that component is
39/// registered and added to the data structures.
40///
41/// A brief overview of the kinds of compiled object and their
42/// relationships:
43///
44/// - `Module` is a Wasm module. It owns a `CompiledModule`.
45/// - `CompiledModule` contains metadata about the module (e.g., a map
46///   from Wasm function indices to locations in the machine code),
47///   and also owns an `EngineCode`.
48/// - `EngineCode` holds an `Arc` to a `CodeMemory` with the canonical
49///   copy of machine code, as well as some lower-level metadata
50///   (signatures and types). It is instantiated by this registry into
51///   `StoreCode`.
52/// - `StoreCode` owns either another `Arc` to the same `CodeMemory`
53///   as `EngineCode`, or if guest debugging is enabled and causes us
54///   to clone private copies of code for patching per store, owns its
55///   own private `CodeMemory` at a different address.
56/// - Instances hold a `RegisteredModuleId` to be able to look up their modules.
57#[derive(Default)]
58pub struct ModuleRegistry {
59    /// StoreCode and Modules associated with it.
60    ///
61    /// Keyed by the start address of the `StoreCode`. We maintain the
62    /// invariant of no overlaps on insertion. We use a range query to
63    /// find the StoreCode for a given PC: take the range `0..=pc`,
64    /// then take the last element of the range. That picks the
65    /// highest start address <= the query, and we can check whether
66    /// it contains the address.
67    loaded_code: TryBTreeMap<StoreCodePC, LoadedCode>,
68
69    /// Map from EngineCodePC start to StoreCodePC start. We use this
70    /// to memoize the store-code creation process: each EngineCode is
71    /// instantiated to a StoreCode only once per store.
72    store_code: TryBTreeMap<EngineCodePC, StoreCodePC>,
73
74    /// Modules instantiated in this registry.
75    ///
76    /// Every module is placed in this map, but not every module will
77    /// be in a LoadedCode entry, because the module may have no text.
78    modules: TryBTreeMap<RegisteredModuleId, Module>,
79}
80
81struct LoadedCode {
82    /// The StoreCode in this range.
83    code: StoreCode,
84
85    /// The index of compiled functions in `code`.
86    ///
87    /// This index is used to map from program counters back to a `FuncKey`.
88    /// Primarily used in [`ModuleRegistry::lookup_frame_info`] at this time
89    /// below.
90    index: Arc<CompiledFunctionsTable>,
91
92    /// A mapping from the [`StaticModuleIndex`] values returned by the `index`
93    /// field above to a module within this registry.
94    ///
95    /// This is lazily populated as modules are registered and instantiated from
96    /// a component.
97    modules: TrySecondaryMap<StaticModuleIndex, Option<RegisteredModuleId>>,
98}
99
100/// An identifier of a module that has previously been inserted into a
101/// `ModuleRegistry`.
102///
103/// This is just a newtype around `CompiledModuleId`, which is unique
104/// within the Engine.
105#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
106pub struct RegisteredModuleId(CompiledModuleId);
107
108fn assert_no_overlap(
109    loaded_code: &TryBTreeMap<StoreCodePC, LoadedCode>,
110    range: Range<StoreCodePC>,
111) {
112    if let Some((start, _)) = loaded_code.range(range.start..).next() {
113        assert!(start >= range.end);
114    }
115    if let Some((_, code)) = loaded_code.range(..range.end).next_back() {
116        assert!(code.code.text_range().end <= range.start);
117    }
118}
119
120#[cfg(feature = "debug")]
121pub struct RegisterBreakpointState<'a>(pub(crate) &'a crate::runtime::debug::BreakpointState);
122#[cfg(not(feature = "debug"))]
123pub struct RegisterBreakpointState<'a>(pub(crate) PhantomData<&'a ()>);
124
125impl<'a> RegisterBreakpointState<'a> {
126    #[cfg(feature = "debug")]
127    fn update(&self, code: &mut StoreCode, module: &Module) -> Result<()> {
128        self.0.patch_new_module(code, module)
129    }
130    #[cfg(not(feature = "debug"))]
131    fn update(&self, _code: &mut StoreCode, _module: &Module) -> Result<()> {
132        Ok(())
133    }
134}
135
136enum ModuleOrComponent<'a> {
137    Module(&'a Module),
138    #[cfg(feature = "component-model")]
139    Component(&'a Component),
140}
141
142impl<'a> ModuleOrComponent<'a> {
143    fn engine(&self) -> &'a Engine {
144        match self {
145            ModuleOrComponent::Module(module) => module.engine(),
146            #[cfg(feature = "component-model")]
147            ModuleOrComponent::Component(component) => component.engine(),
148        }
149    }
150}
151
152impl ModuleRegistry {
153    /// Get a previously-registered module by id.
154    pub fn module_by_id(&self, id: RegisteredModuleId) -> Option<&Module> {
155        self.modules.get(id)
156    }
157
158    /// Get a module by CompiledModuleId, if present.
159    pub fn module_by_compiled_id(&self, id: CompiledModuleId) -> Option<&Module> {
160        self.modules.get(RegisteredModuleId(id))
161    }
162
163    /// Looks up `pc`, an absolute program counter address, to see if it
164    /// corresponds to any loaded code within this [`ModuleRegistry`].
165    ///
166    /// Returns `None` if nothing is found, and otherwise returns the
167    /// corresponding [`LoadedCode`] as well as a relative pc from the start of
168    /// the code's text section if found.
169    fn loaded_code_by_pc(&self, pc: usize) -> Option<(&LoadedCode, usize)> {
170        let (_, code) = self
171            .loaded_code
172            .range(..=StoreCodePC::from_raw(pc))
173            .next_back()?;
174        let offset = StoreCodePC::offset_of(code.code.text_range(), pc)?;
175        Some((code, offset))
176    }
177
178    /// Consults this [`ModuleRegistry`] to see if `pc` corrseponds to any
179    /// previously-registered block of code.
180    ///
181    /// Upon success returns the [`StoreCode`] as well as a relative pc from the
182    /// start of the text section.
183    pub fn store_code_by_pc(&self, pc: usize) -> Option<(&StoreCode, usize)> {
184        let (code, pc) = self.loaded_code_by_pc(pc)?;
185        Some((&code.code, pc))
186    }
187
188    /// Fetches the `StoreCode` for a given `EngineCode`.
189    pub fn store_code(&self, engine_code: &EngineCode) -> Option<&StoreCode> {
190        let store_code_pc = self.store_code_base(engine_code)?;
191        let (_, code) = self.loaded_code.range(store_code_pc..).next()?;
192        Some(&code.code)
193    }
194
195    /// Fetches the base `StoreCodePC` for a given `EngineCode`.
196    pub fn store_code_base(&self, engine_code: &EngineCode) -> Option<StoreCodePC> {
197        self.store_code.get(engine_code.text_range().start).cloned()
198    }
199
200    /// Fetches the base `StoreCodePC` for a given `EngineCode` with
201    /// `Module`, registering the module if not already registered.
202    pub fn store_code_base_or_register(
203        &mut self,
204        module: &Module,
205        breakpoint_state: RegisterBreakpointState,
206    ) -> Result<StoreCodePC> {
207        let key = module.engine_code().text_range().start;
208        if !self.store_code.contains_key(key) {
209            let engine = module.engine().clone();
210            self.register_module(module, &engine, breakpoint_state)?;
211        }
212        Ok(*self.store_code.get(key).unwrap())
213    }
214
215    /// Fetches a mutable `StoreCode` for a given base `StoreCodePC`.
216    pub fn store_code_mut(&mut self, store_code_base: StoreCodePC) -> Option<&mut StoreCode> {
217        let (_, code) = self.loaded_code.range_mut(store_code_base..).next()?;
218        assert_eq!(code.code.text_range().start, store_code_base);
219        Some(&mut code.code)
220    }
221
222    /// Gets an iterator over all modules in the registry.
223    #[cfg(any(feature = "coredump", feature = "debug", feature = "gc"))]
224    pub fn all_modules(&self) -> impl Iterator<Item = (RegisteredModuleId, &'_ Module)> + '_ {
225        self.modules.iter()
226    }
227
228    /// Registers a new module with the registry.
229    ///
230    /// Returns an error if `module` was not compiled by `engine`.
231    pub fn register_module(
232        &mut self,
233        module: &Module,
234        engine: &Engine,
235        breakpoint_state: RegisterBreakpointState,
236    ) -> Result<RegisteredModuleId> {
237        self.register(ModuleOrComponent::Module(module), engine, breakpoint_state)
238            .map(|id| id.unwrap())
239    }
240
241    /// Registers a new component with the registry.
242    ///
243    /// Returns an error if `component` was not compiled by `engine`.
244    #[cfg(feature = "component-model")]
245    pub fn register_component(
246        &mut self,
247        component: &Component,
248        engine: &Engine,
249        breakpoint_state: RegisterBreakpointState,
250    ) -> Result<()> {
251        self.register(
252            ModuleOrComponent::Component(component),
253            engine,
254            breakpoint_state,
255        )?;
256        Ok(())
257    }
258
259    /// Registers a new module with the registry.
260    ///
261    /// Returns an error if `module_or_component` was not compiled by `engine`.
262    /// All code and metadata enters a store through here, so callers
263    /// registering on a store's behalf pass that store's engine to keep foreign
264    /// code out of it.
265    fn register(
266        &mut self,
267        module_or_component: ModuleOrComponent<'_>,
268        engine: &Engine,
269        breakpoint_state: RegisterBreakpointState,
270    ) -> Result<Option<RegisteredModuleId>> {
271        // Nothing below here may run for a foreign module or component: the
272        // type indices it carries mean something else entirely in this
273        // engine.
274        ensure!(
275            Engine::same(engine, module_or_component.engine()),
276            "cross-`Engine` usage is not supported"
277        );
278        let compiled_id = match module_or_component {
279            ModuleOrComponent::Module(module) => module.id(),
280            #[cfg(feature = "component-model")]
281            ModuleOrComponent::Component(component) => component.id(),
282        };
283        let code = match module_or_component {
284            ModuleOrComponent::Module(module) => module.engine_code(),
285            #[cfg(feature = "component-model")]
286            ModuleOrComponent::Component(component) => component.engine_code(),
287        };
288        // Register the module, if any.
289        let id = match module_or_component {
290            ModuleOrComponent::Module(module) => {
291                let id = RegisteredModuleId(compiled_id);
292                self.modules.entry(id).or_insert_with(|| module.clone())?;
293                Some(id)
294            }
295            #[cfg(feature = "component-model")]
296            ModuleOrComponent::Component(_) => None,
297        };
298
299        // Create a StoreCode if one does not already exist.
300        let store_code_pc = match self.store_code.entry(code.text_range().start) {
301            Entry::Vacant(v) => {
302                let store_code = StoreCode::new(engine, code)?;
303                let store_code_pc = store_code.text_range().start;
304                assert_no_overlap(&self.loaded_code, store_code.text_range());
305                let index = match module_or_component {
306                    ModuleOrComponent::Module(module) => module.index(),
307                    #[cfg(feature = "component-model")]
308                    ModuleOrComponent::Component(component) => component.index(),
309                };
310                self.loaded_code.insert(
311                    store_code_pc,
312                    LoadedCode {
313                        code: store_code,
314                        index: index.clone(),
315                        modules: Default::default(),
316                    },
317                )?;
318                *v.insert(store_code_pc)?
319            }
320            Entry::Occupied(o) => *o.get(),
321        };
322
323        // Add this module to the LoadedCode if not present.
324        if let (ModuleOrComponent::Module(module), Some(id)) = (module_or_component, id) {
325            let loaded_code = self
326                .loaded_code
327                .get_mut(store_code_pc)
328                .expect("loaded_code must have entry for StoreCodePC");
329            loaded_code
330                .modules
331                .insert(module.env_module().module_index, Some(id))?;
332            breakpoint_state.update(&mut loaded_code.code, module)?;
333        }
334
335        Ok(id)
336    }
337
338    /// Fetches frame information about a program counter in a backtrace.
339    ///
340    /// Returns an object if this `pc` is known to some previously registered
341    /// module, or returns `None` if no information can be found. The first
342    /// boolean returned indicates whether the original module has unparsed
343    /// debug information due to the compiler's configuration. The second
344    /// boolean indicates whether the engine used to compile this module is
345    /// using environment variables to control debuginfo parsing.
346    pub(crate) fn lookup_frame_info<'a>(
347        &'a self,
348        pc: usize,
349    ) -> Option<(FrameInfo, ModuleWithCode<'a>)> {
350        let (code, text_offset) = self.loaded_code_by_pc(pc)?;
351        let module_index = match code
352            .index
353            .func_by_text_offset(u32::try_from(text_offset).ok()?)?
354        {
355            FuncKey::DefinedWasmFunction(module, _) => module,
356            _ => return None,
357        };
358        let module_id = (*code.modules.get(module_index)?)?;
359        let module = self
360            .modules
361            .get(module_id)
362            .expect("referenced module ID not found");
363        let info = FrameInfo::new(module.clone(), text_offset)?;
364        let module_with_code = ModuleWithCode::from_raw(module, &code.code);
365        Some((info, module_with_code))
366    }
367
368    pub fn wasm_to_array_trampoline(
369        &self,
370        sig: VMSharedTypeIndex,
371    ) -> Option<NonNull<VMWasmCallFunction>> {
372        // TODO: We are doing a linear search over each module. This is fine for
373        // now because we typically have very few modules per store (almost
374        // always one, in fact). If this linear search ever becomes a
375        // bottleneck, we could avoid it by incrementally and lazily building a
376        // `VMSharedSignatureIndex` to `SignatureIndex` map.
377        //
378        // See also the comment in `ModuleInner::wasm_to_native_trampoline`.
379        for module in self.modules.values() {
380            if let Some(trampoline) = module.wasm_to_array_trampoline(sig) {
381                return Some(trampoline);
382            }
383        }
384        None
385    }
386}
387
388// This is the global code registry that stores information for all loaded code
389// objects that are currently in use by any `Store` in the current process.
390//
391// The purpose of this map is to be called from signal handlers to determine
392// whether a program counter is a wasm trap or not. Specifically macOS has
393// no contextual information about the thread available, hence the necessity
394// for global state rather than using thread local state.
395//
396// This is similar to `ModuleRegistry` except that it has less information and
397// supports removal. Any time anything is registered with a `ModuleRegistry`
398// it is also automatically registered with the singleton global module
399// registry. When a `ModuleRegistry` is destroyed then all of its entries
400// are removed from the global registry.
401fn global_code() -> &'static RwLock<GlobalRegistry> {
402    static GLOBAL_CODE: OnceLock<RwLock<GlobalRegistry>> = OnceLock::new();
403    GLOBAL_CODE.get_or_init(Default::default)
404}
405
406type GlobalRegistry = TryBTreeMap<usize, (usize, Arc<CodeMemory>)>;
407
408/// Find which registered region of code contains the given program counter, and
409/// what offset that PC is within that module's code.
410pub fn lookup_code(pc: usize) -> Option<(Arc<CodeMemory>, usize)> {
411    let all_modules = global_code().read();
412    let (_end, (start, module)) = all_modules.range(pc..).next()?;
413    let text_offset = pc.checked_sub(*start)?;
414    Some((module.clone(), text_offset))
415}
416
417/// Registers a new region of code.
418///
419/// Must not have been previously registered and must be `unregister`'d to
420/// prevent leaking memory.
421///
422/// This is required to enable traps to work correctly since the signal handler
423/// will lookup in the `GLOBAL_CODE` list to determine which a particular pc
424/// is a trap or not.
425pub fn register_code(image: &Arc<CodeMemory>, address: Range<usize>) -> Result<(), OutOfMemory> {
426    if address.is_empty() {
427        return Ok(());
428    }
429    let start = address.start;
430    let end = address.end - 1;
431    let prev = global_code().write().insert(end, (start, image.clone()))?;
432    assert!(prev.is_none());
433    Ok(())
434}
435
436/// Unregisters a code mmap from the global map.
437///
438/// Must have been previously registered with `register`.
439pub fn unregister_code(address: Range<usize>) {
440    if address.is_empty() {
441        return;
442    }
443    let end = address.end - 1;
444    let code = global_code().write().remove(end);
445    assert!(code.is_some());
446}
447
448#[test]
449#[cfg_attr(miri, ignore)]
450fn test_frame_info() -> Result<(), crate::Error> {
451    use crate::*;
452
453    let mut store = Store::<()>::default();
454    let module = Module::new(
455        store.engine(),
456        r#"
457            (module
458                (func (export "add") (param $x i32) (param $y i32) (result i32) (i32.add (local.get $x) (local.get $y)))
459                (func (export "sub") (param $x i32) (param $y i32) (result i32) (i32.sub (local.get $x) (local.get $y)))
460                (func (export "mul") (param $x i32) (param $y i32) (result i32) (i32.mul (local.get $x) (local.get $y)))
461                (func (export "div_s") (param $x i32) (param $y i32) (result i32) (i32.div_s (local.get $x) (local.get $y)))
462                (func (export "div_u") (param $x i32) (param $y i32) (result i32) (i32.div_u (local.get $x) (local.get $y)))
463                (func (export "rem_s") (param $x i32) (param $y i32) (result i32) (i32.rem_s (local.get $x) (local.get $y)))
464                (func (export "rem_u") (param $x i32) (param $y i32) (result i32) (i32.rem_u (local.get $x) (local.get $y)))
465            )
466         "#,
467    )?;
468    // Create an instance to ensure the frame information is registered.
469    Instance::new(&mut store, &module, &[])?;
470
471    // Look for frame info for each function. Assume that StoreCode
472    // does not actually clone in the default configuration.
473    for (i, range) in module.compiled_module().finished_function_ranges() {
474        let base = module.engine_code().text_range().start.raw();
475        let start = base + range.start;
476        let end = base + range.end;
477        for pc in start..end {
478            let (frame, _) = store
479                .as_context()
480                .0
481                .modules()
482                .lookup_frame_info(pc)
483                .unwrap();
484            assert!(
485                frame.func_index() == i.as_u32(),
486                "lookup of {:#x} returned {}, expected {}",
487                pc,
488                frame.func_index(),
489                i.as_u32()
490            );
491        }
492    }
493    Ok(())
494}
wasmtime/runtime/module/registry.rs

wasmtime/runtime/module/
registry.rs
