From af59c4d568d487b7efbb49d7d31a861e7c3933a6 Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Fri, 7 Jun 2024 17:24:45 -0500 Subject: [PATCH] Make module ids unique per-process, not per-engine (#8758) * Make module ids unique per-process, not per-engine Currently all instances of `wasmtime::Module` have a unique 64-bit id embedded into them. This ID was originally only used for affinity in the pooling allocator as a quick check of module equality. This use case only required engine-local ids so the initial implementation had an ID allocator per-engine. Later, however, this id was reused for `wasmtime::ModuleExport` which was intended to skip the string lookup for an export at runtime. This also stored a module id but it did not store an engine identifier. This meant that it's possible to mix these up and pass the wrong export to the wrong engine. This behavior can lead to a runtime panic in Wasmtime. This commit fixes this by making the module identifier be global per-process instead of per-engine. This mirrors how store IDs are allocated where they're per-process instead of per-engine. The same logic for why store IDs are unlikely to be exhausted applies here too where this 64-bit space of identifiers is unlikely to be exhausted. * Fix warnings * Fix tests --- crates/wasmtime/src/engine.rs | 8 ---- crates/wasmtime/src/runtime/instantiate.rs | 5 +-- crates/wasmtime/src/runtime/module.rs | 8 +--- crates/wasmtime/src/runtime/store/data.rs | 2 +- crates/wasmtime/src/runtime/vm.rs | 2 +- .../allocator/pooling/index_allocator.rs | 32 +++++++-------- crates/wasmtime/src/runtime/vm/module_id.rs | 40 ++++--------------- tests/all/module.rs | 23 +++++++++++ 8 files changed, 51 insertions(+), 69 deletions(-) diff --git a/crates/wasmtime/src/engine.rs b/crates/wasmtime/src/engine.rs index 3bf8bcb1df..64e1a50cb7 100644 --- a/crates/wasmtime/src/engine.rs +++ b/crates/wasmtime/src/engine.rs @@ -61,8 +61,6 @@ struct EngineInner { signatures: TypeRegistry, #[cfg(feature = "runtime")] epoch: AtomicU64, - #[cfg(feature = "runtime")] - unique_id_allocator: crate::runtime::vm::CompiledModuleIdAllocator, /// One-time check of whether the compiler's settings, if present, are /// compatible with the native host. @@ -129,8 +127,6 @@ impl Engine { signatures: TypeRegistry::new(), #[cfg(feature = "runtime")] epoch: AtomicU64::new(0), - #[cfg(feature = "runtime")] - unique_id_allocator: crate::runtime::vm::CompiledModuleIdAllocator::new(), #[cfg(any(feature = "cranelift", feature = "winch"))] compatible_with_native_host: OnceLock::new(), config, @@ -639,10 +635,6 @@ impl Engine { self.inner.epoch.fetch_add(1, Ordering::Relaxed); } - pub(crate) fn unique_id_allocator(&self) -> &crate::runtime::vm::CompiledModuleIdAllocator { - &self.inner.unique_id_allocator - } - /// Returns a [`std::hash::Hash`] that can be used to check precompiled WebAssembly compatibility. /// /// The outputs of [`Engine::precompile_module`] and [`Engine::precompile_component`] diff --git a/crates/wasmtime/src/runtime/instantiate.rs b/crates/wasmtime/src/runtime/instantiate.rs index 1f135f03eb..fd725ceee8 100644 --- a/crates/wasmtime/src/runtime/instantiate.rs +++ b/crates/wasmtime/src/runtime/instantiate.rs @@ -4,7 +4,7 @@ //! steps. use crate::prelude::*; -use crate::runtime::vm::{CompiledModuleId, CompiledModuleIdAllocator, MmapVec}; +use crate::runtime::vm::{CompiledModuleId, MmapVec}; use crate::{code_memory::CodeMemory, profiling_agent::ProfilingAgent}; use alloc::sync::Arc; use anyhow::Result; @@ -50,7 +50,6 @@ impl CompiledModule { code_memory: Arc, info: CompiledModuleInfo, profiler: &dyn ProfilingAgent, - id_allocator: &CompiledModuleIdAllocator, ) -> Result { let mut ret = Self { module: Arc::new(info.module), @@ -60,7 +59,7 @@ impl CompiledModule { dbg_jit_registration: None, code_memory, meta: info.meta, - unique_id: id_allocator.alloc(), + unique_id: CompiledModuleId::new(), func_names: info.func_names, }; ret.register_debug_and_profiling(profiler)?; diff --git a/crates/wasmtime/src/runtime/module.rs b/crates/wasmtime/src/runtime/module.rs index 70c0e3fb3d..e3c6d88cb9 100644 --- a/crates/wasmtime/src/runtime/module.rs +++ b/crates/wasmtime/src/runtime/module.rs @@ -478,12 +478,8 @@ impl Module { info: CompiledModuleInfo, serializable: bool, ) -> Result { - let module = CompiledModule::from_artifacts( - code.code_memory().clone(), - info, - engine.profiler(), - engine.unique_id_allocator(), - )?; + let module = + CompiledModule::from_artifacts(code.code_memory().clone(), info, engine.profiler())?; // Validate the module can be used with the current instance allocator. let offsets = VMOffsets::new(HostPtr, module.module()); diff --git a/crates/wasmtime/src/runtime/store/data.rs b/crates/wasmtime/src/runtime/store/data.rs index 54083b0cd1..4a119bcd80 100644 --- a/crates/wasmtime/src/runtime/store/data.rs +++ b/crates/wasmtime/src/runtime/store/data.rs @@ -209,7 +209,7 @@ pub struct StoreId(NonZeroU64); impl StoreId { /// Allocates a new unique identifier for a store that has never before been /// used in this process. - fn allocate() -> StoreId { + pub fn allocate() -> StoreId { static NEXT_ID: AtomicU64 = AtomicU64::new(0); // Only allow 2^63 stores at which point we start panicking to prevent diff --git a/crates/wasmtime/src/runtime/vm.rs b/crates/wasmtime/src/runtime/vm.rs index 379436630a..5d2bc52b4a 100644 --- a/crates/wasmtime/src/runtime/vm.rs +++ b/crates/wasmtime/src/runtime/vm.rs @@ -74,7 +74,7 @@ pub use crate::runtime::vm::vmcontext::{ pub use send_sync_ptr::SendSyncPtr; mod module_id; -pub use module_id::{CompiledModuleId, CompiledModuleIdAllocator}; +pub use module_id::CompiledModuleId; mod cow; pub use crate::runtime::vm::cow::{MemoryImage, MemoryImageSlot, ModuleMemoryImages}; diff --git a/crates/wasmtime/src/runtime/vm/instance/allocator/pooling/index_allocator.rs b/crates/wasmtime/src/runtime/vm/instance/allocator/pooling/index_allocator.rs index 72efa65417..2992db7cda 100644 --- a/crates/wasmtime/src/runtime/vm/instance/allocator/pooling/index_allocator.rs +++ b/crates/wasmtime/src/runtime/vm/instance/allocator/pooling/index_allocator.rs @@ -485,7 +485,6 @@ impl List { #[cfg(test)] mod test { use super::*; - use crate::runtime::vm::CompiledModuleIdAllocator; use wasmtime_environ::EntityRef; #[test] @@ -503,9 +502,8 @@ mod test { #[test] fn test_affinity_allocation_strategy() { - let id_alloc = CompiledModuleIdAllocator::new(); - let id1 = MemoryInModule(id_alloc.alloc(), DefinedMemoryIndex::new(0)); - let id2 = MemoryInModule(id_alloc.alloc(), DefinedMemoryIndex::new(0)); + let id1 = MemoryInModule(CompiledModuleId::new(), DefinedMemoryIndex::new(0)); + let id2 = MemoryInModule(CompiledModuleId::new(), DefinedMemoryIndex::new(0)); let state = ModuleAffinityIndexAllocator::new(100, 100); let index1 = state.alloc(Some(id1)).unwrap(); @@ -560,8 +558,7 @@ mod test { #[test] fn clear_affine() { - let id_alloc = CompiledModuleIdAllocator::new(); - let id = id_alloc.alloc(); + let id = CompiledModuleId::new(); let memory_index = DefinedMemoryIndex::new(0); for max_unused_warm_slots in [0, 1, 2] { @@ -589,11 +586,11 @@ mod test { use rand::Rng; let mut rng = rand::thread_rng(); - let id_alloc = CompiledModuleIdAllocator::new(); - let ids = - std::iter::repeat_with(|| MemoryInModule(id_alloc.alloc(), DefinedMemoryIndex::new(0))) - .take(10) - .collect::>(); + let ids = std::iter::repeat_with(|| { + MemoryInModule(CompiledModuleId::new(), DefinedMemoryIndex::new(0)) + }) + .take(10) + .collect::>(); let state = ModuleAffinityIndexAllocator::new(1000, 1000); let mut allocated: Vec = vec![]; let mut last_id = vec![None; 1000]; @@ -634,10 +631,9 @@ mod test { #[test] fn test_affinity_threshold() { - let id_alloc = CompiledModuleIdAllocator::new(); - let id1 = MemoryInModule(id_alloc.alloc(), DefinedMemoryIndex::new(0)); - let id2 = MemoryInModule(id_alloc.alloc(), DefinedMemoryIndex::new(0)); - let id3 = MemoryInModule(id_alloc.alloc(), DefinedMemoryIndex::new(0)); + let id1 = MemoryInModule(CompiledModuleId::new(), DefinedMemoryIndex::new(0)); + let id2 = MemoryInModule(CompiledModuleId::new(), DefinedMemoryIndex::new(0)); + let id3 = MemoryInModule(CompiledModuleId::new(), DefinedMemoryIndex::new(0)); let state = ModuleAffinityIndexAllocator::new(10, 2); // Set some slot affinities @@ -673,7 +669,7 @@ mod test { // LRU is 1, so that should be picked assert_eq!( state.alloc(Some(MemoryInModule( - id_alloc.alloc(), + CompiledModuleId::new(), DefinedMemoryIndex::new(0) ))), Some(SlotId(1)) @@ -682,7 +678,7 @@ mod test { // Pick another LRU entry, this time 2 assert_eq!( state.alloc(Some(MemoryInModule( - id_alloc.alloc(), + CompiledModuleId::new(), DefinedMemoryIndex::new(0) ))), Some(SlotId(2)) @@ -691,7 +687,7 @@ mod test { // This should preserve slot `0` and pick up something new assert_eq!( state.alloc(Some(MemoryInModule( - id_alloc.alloc(), + CompiledModuleId::new(), DefinedMemoryIndex::new(0) ))), Some(SlotId(3)) diff --git a/crates/wasmtime/src/runtime/vm/module_id.rs b/crates/wasmtime/src/runtime/vm/module_id.rs index 899e8f43e5..be5bc513ec 100644 --- a/crates/wasmtime/src/runtime/vm/module_id.rs +++ b/crates/wasmtime/src/runtime/vm/module_id.rs @@ -1,43 +1,19 @@ //! Unique IDs for modules in the runtime. -use core::{ - num::NonZeroU64, - sync::atomic::{AtomicU64, Ordering}, -}; +use core::num::NonZeroU64; /// A unique identifier (within an engine or similar) for a compiled /// module. #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct CompiledModuleId(NonZeroU64); -/// An allocator for compiled module IDs. -pub struct CompiledModuleIdAllocator { - next: AtomicU64, -} - -impl CompiledModuleIdAllocator { - /// Create a compiled-module ID allocator. +impl CompiledModuleId { + /// Allocates a new ID which will be unique within this process. pub fn new() -> Self { - Self { - next: AtomicU64::new(1), - } - } - - /// Allocate a new ID. - pub fn alloc(&self) -> CompiledModuleId { - // Note: why is `Relaxed` OK here? - // - // The only requirement we have is that IDs are unique. We - // don't care how one module's ID compares to another, i.e., - // what order they come in. `Relaxed` means that this - // `fetch_add` operation does not have any particular - // synchronization (ordering) with respect to any other memory - // access in the program. However, `fetch_add` is always - // atomic with respect to other accesses to this variable - // (`self.next`). So we will always hand out separate, unique - // IDs correctly, just in some possibly arbitrary order (which - // is fine). - let id = self.next.fetch_add(1, Ordering::Relaxed); - CompiledModuleId(NonZeroU64::new(id).unwrap()) + // As an implementation detail this is implemented on the same + // allocator as stores. It's ok if there are "holes" in the store id + // space as it's not required to be compact, it's just used for + // uniqueness. + CompiledModuleId(crate::store::StoreId::allocate().as_raw()) } } diff --git a/tests/all/module.rs b/tests/all/module.rs index 752c822624..eb658f69fe 100644 --- a/tests/all/module.rs +++ b/tests/all/module.rs @@ -330,3 +330,26 @@ fn tail_call_defaults() -> Result<()> { } Ok(()) } + +#[test] +fn cross_engine_module_exports() -> Result<()> { + let a_engine = Engine::default(); + let b_engine = Engine::default(); + + let a_module = Module::new(&a_engine, "(module)")?; + let b_module = Module::new( + &b_engine, + r#" + (module + (func (export "x")) + ) + "#, + )?; + + let export = b_module.get_export_index("x").unwrap(); + + let mut store = Store::new(&a_engine, ()); + let instance = Instance::new(&mut store, &a_module, &[])?; + assert!(instance.get_module_export(&mut store, &export).is_none()); + Ok(()) +}