From 8ec478137883da5c76307c006e25b0e6d65732d9 Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Mon, 12 Feb 2024 11:25:58 -0800 Subject: [PATCH] Make `RegisteredType` hold a whole `Engine` Rather than an `Arc>`. This also removes the `Arc` inside `TypeRegistry` and we effectively reuse `Engine`'s internal `Arc` instead. Also, to avoid various `TypeRegistry` methods needing to take an `Engine` to construct their `RegisteredType` results, but then needing to assert that the given engine is this registry's engine, I turned the methods into `TypeRegistry` constructors. These constructors take just an `&Engine` and then access the underlying `TypeRegistry` as needed, effectively making that property hold statically. This is a minor clean up on its own, but is a big help for follow up work I am doing for Wasm GC and typed function references, where being able to grab a reference to the engine that a `FuncType` is registered within will prevent needing to thread in additional engine parameters to various places and then assert that the engine is the engine that the type is registered within, and etc... --- .../src/runtime/component/component.rs | 2 +- crates/wasmtime/src/runtime/module.rs | 2 +- .../wasmtime/src/runtime/trampoline/func.rs | 2 +- crates/wasmtime/src/runtime/type_registry.rs | 144 +++++++++++------- crates/wasmtime/src/runtime/types.rs | 6 +- 5 files changed, 94 insertions(+), 62 deletions(-) diff --git a/crates/wasmtime/src/runtime/component/component.rs b/crates/wasmtime/src/runtime/component/component.rs index 97ca94164afc..326dafcd3b83 100644 --- a/crates/wasmtime/src/runtime/component/component.rs +++ b/crates/wasmtime/src/runtime/component/component.rs @@ -307,7 +307,7 @@ impl Component { // Create a signature registration with the `Engine` for all trampolines // and core wasm types found within this component, both for the // component and for all included core wasm modules. - let signatures = TypeCollection::new_for_module(engine.signatures(), types.module_types()); + let signatures = TypeCollection::new_for_module(engine, types.module_types()); // Assemble the `CodeObject` artifact which is shared by all core wasm // modules as well as the final component. diff --git a/crates/wasmtime/src/runtime/module.rs b/crates/wasmtime/src/runtime/module.rs index 70301e8074ca..cb081b372a56 100644 --- a/crates/wasmtime/src/runtime/module.rs +++ b/crates/wasmtime/src/runtime/module.rs @@ -517,7 +517,7 @@ impl Module { // Note that the unsafety here should be ok since the `trampolines` // field should only point to valid trampoline function pointers // within the text section. - let signatures = TypeCollection::new_for_module(engine.signatures(), &types); + let signatures = TypeCollection::new_for_module(engine, &types); // Package up all our data into a `CodeObject` and delegate to the final // step of module compilation. diff --git a/crates/wasmtime/src/runtime/trampoline/func.rs b/crates/wasmtime/src/runtime/trampoline/func.rs index 6c800f82a942..50c078e9d2a3 100644 --- a/crates/wasmtime/src/runtime/trampoline/func.rs +++ b/crates/wasmtime/src/runtime/trampoline/func.rs @@ -123,7 +123,7 @@ where let native_call = text[native_call_range.start as usize..].as_ptr() as *mut _; let native_call = NonNull::new(native_call).unwrap(); - let sig = engine.signatures().register(ft.as_wasm_func_type()); + let sig = ft.clone().into_registered_type(); unsafe { Ok(VMArrayCallHostFuncContext::new( diff --git a/crates/wasmtime/src/runtime/type_registry.rs b/crates/wasmtime/src/runtime/type_registry.rs index ce491fdbec98..9ddef38b1cdd 100644 --- a/crates/wasmtime/src/runtime/type_registry.rs +++ b/crates/wasmtime/src/runtime/type_registry.rs @@ -3,32 +3,49 @@ //! Helps implement fast indirect call signature checking, reference type //! casting, and etc. +use std::fmt::Debug; use std::{collections::HashMap, sync::RwLock}; use std::{convert::TryFrom, sync::Arc}; use wasmtime_environ::{ModuleInternedTypeIndex, ModuleTypes, PrimaryMap, WasmFuncType}; use wasmtime_runtime::VMSharedTypeIndex; +use crate::Engine; + /// Represents a collection of shared types. /// /// This is used to register shared types with a shared type registry. /// /// The collection will unregister any contained types with the registry /// when dropped. -#[derive(Debug)] pub struct TypeCollection { - registry: Arc>, + engine: Engine, types: PrimaryMap, reverse_types: HashMap, } +impl Debug for TypeCollection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let TypeCollection { + engine: _, + types, + reverse_types: _, + } = self; + f.debug_struct("TypeCollection") + .field("types", types) + .finish_non_exhaustive() + } +} + impl TypeCollection { /// Creates a type collection for a module given the module's types. - pub fn new_for_module(registry: &TypeRegistry, types: &ModuleTypes) -> Self { + pub fn new_for_module(engine: &Engine, types: &ModuleTypes) -> Self { + let engine = engine.clone(); + let registry = engine.signatures(); let types = registry.0.write().unwrap().register_for_module(types); let reverse_types = types.iter().map(|(k, v)| (*v, k)).collect(); Self { - registry: registry.0.clone(), + engine, types, reverse_types, } @@ -58,20 +75,26 @@ impl TypeCollection { impl Drop for TypeCollection { fn drop(&mut self) { if !self.types.is_empty() { - self.registry.write().unwrap().unregister_types(self); + self.engine + .signatures() + .0 + .write() + .unwrap() + .unregister_types(self); } } } /// A Wasm type that has been registered in the engine's `TypeRegistry`. /// -/// Automatically unregisters the type on drop. (If other things are keeping the -/// type registered, it will remain registered). +/// Prevents its associated type from being unregistered while it is alive. +/// +/// Automatically unregisters the type on drop. (Unless other `RegisteredTypes` +/// are keeping the type registered). /// /// Dereferences to its underlying `WasmFuncType`. -#[derive(Debug)] pub struct RegisteredType { - registry: Arc>, + engine: Engine, index: VMSharedTypeIndex, // This field is not *strictly* necessary to have in this type, since we @@ -81,11 +104,25 @@ pub struct RegisteredType { ty: Arc, } +impl Debug for RegisteredType { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let RegisteredType { + engine: _, + index, + ty, + } = self; + f.debug_struct("RegisteredType") + .field("index", index) + .field("ty", ty) + .finish_non_exhaustive() + } +} + impl Clone for RegisteredType { fn clone(&self) -> Self { { let i = usize::try_from(self.index.bits()).unwrap(); - let mut registry = self.registry.write().unwrap(); + let mut registry = self.engine.signatures().0.write().unwrap(); let entry = registry.entries[i].unwrap_occupied_mut(); entry.references += 1; log::trace!( @@ -96,7 +133,7 @@ impl Clone for RegisteredType { } RegisteredType { - registry: Arc::clone(&self.registry), + engine: self.engine.clone(), index: self.index, ty: Arc::clone(&self.ty), } @@ -105,7 +142,9 @@ impl Clone for RegisteredType { impl Drop for RegisteredType { fn drop(&mut self) { - self.registry + self.engine + .signatures() + .0 .write() .unwrap() .unregister_entry(self.index, 1); @@ -127,9 +166,9 @@ impl PartialEq for RegisteredType { if cfg!(debug_assertions) { if eq { assert_eq!(self.index, other.index); - assert!(Arc::ptr_eq(&self.registry, &other.registry)); + assert!(Engine::same(&self.engine, &other.engine)); } else { - assert!(self.index != other.index || !Arc::ptr_eq(&self.registry, &other.registry)); + assert!(self.index != other.index || !Engine::same(&self.engine, &other.engine)); } } @@ -147,26 +186,44 @@ impl std::hash::Hash for RegisteredType { } impl RegisteredType { + /// Constructs a new `RegisteredType`, registering the given type with the + /// engine's `TypeRegistry`. + pub fn new(engine: &Engine, ty: &WasmFuncType) -> RegisteredType { + let (index, ty) = engine.signatures().0.write().unwrap().register_raw(ty); + RegisteredType::from_parts(engine.clone(), index, ty) + } + + /// Create an owning handle to the given index's associated type. + /// + /// This will prevent the associated type from being unregistered as long as + /// the returned `RegisteredType` is kept alive. + /// + /// Returns `None` if `index` is not registered in the given engine's + /// registry. + pub fn root(engine: &Engine, index: VMSharedTypeIndex) -> Option { + let i = usize::try_from(index.bits()).unwrap(); + let ty = { + let mut inner = engine.signatures().0.write().unwrap(); + let e = inner.entries.get_mut(i)?.as_occupied_mut()?; + e.references += 1; + log::trace!("rooting {index:?} (references -> {})", e.references); + Arc::clone(&e.ty) + }; + Some(RegisteredType::from_parts(engine.clone(), index, ty)) + } + /// Construct a new `RegisteredType`. /// /// It is the caller's responsibility to ensure that the entry's reference /// count has already been incremented. - fn new( - registry: Arc>, - index: VMSharedTypeIndex, - ty: Arc, - ) -> Self { + fn from_parts(engine: Engine, index: VMSharedTypeIndex, ty: Arc) -> Self { debug_assert!({ - let registry = registry.read().unwrap(); + let registry = engine.signatures().0.read().unwrap(); let i = usize::try_from(index.bits()).unwrap(); let e = registry.entries[i].as_occupied().unwrap(); e.references > 0 }); - RegisteredType { - registry, - index, - ty, - } + RegisteredType { engine, index, ty } } /// Get this registered type's index. @@ -375,49 +432,24 @@ impl Drop for TypeRegistryInner { /// types, shared by all instances, so that call sites can just do an /// index comparison. #[derive(Debug)] -pub struct TypeRegistry(Arc>); +pub struct TypeRegistry(RwLock); impl TypeRegistry { /// Creates a new shared type registry. pub fn new() -> Self { - Self(Arc::new(RwLock::new(TypeRegistryInner::default()))) - } - - /// Get an owning handle to the given index's associated type. - /// - /// This will keep the type registered as long as the return value is kept - /// alive. - pub fn root(&self, index: VMSharedTypeIndex) -> Option { - let i = usize::try_from(index.bits()).unwrap(); - let registry = Arc::clone(&self.0); - let ty = { - let mut inner = self.0.write().unwrap(); - let e = inner.entries.get_mut(i)?.as_occupied_mut()?; - e.references += 1; - log::trace!("rooting {index:?} (references -> {})", e.references); - Arc::clone(&e.ty) - }; - Some(RegisteredType::new(registry, index, ty)) + Self(RwLock::new(TypeRegistryInner::default())) } /// Looks up a function type from a shared type index. /// - /// This does NOT guarantee that the type remains registered while you use - /// the result. Use the `root` method if you need to ensure that property - /// and you don't have some other mechanism already keeping the type - /// registered. + /// This does *NOT* prevent the type from being unregistered while you are + /// still using the resulting value! Use the `RegisteredType::root` + /// constructor if you need to ensure that property and you don't have some + /// other mechanism already keeping the type registered. pub fn borrow(&self, index: VMSharedTypeIndex) -> Option> { let i = usize::try_from(index.bits()).unwrap(); let inner = self.0.read().unwrap(); let e = inner.entries.get(i)?; Some(e.as_occupied()?.ty.clone()) } - - /// Registers a single function with the collection. - /// - /// Returns the shared type index for the function. - pub fn register(&self, ty: &WasmFuncType) -> RegisteredType { - let (index, ty) = self.0.write().unwrap().register_raw(ty); - RegisteredType::new(Arc::clone(&self.0), index, ty) - } } diff --git a/crates/wasmtime/src/runtime/types.rs b/crates/wasmtime/src/runtime/types.rs index bc4413f2a10a..131a8c7c189d 100644 --- a/crates/wasmtime/src/runtime/types.rs +++ b/crates/wasmtime/src/runtime/types.rs @@ -252,13 +252,13 @@ impl FuncType { } pub(crate) fn from_wasm_func_type(engine: &Engine, ty: &WasmFuncType) -> FuncType { - let ty = engine.signatures().register(ty); + let ty = RegisteredType::new(engine, ty); Self { ty } } pub(crate) fn from_shared_type_index(engine: &Engine, index: VMSharedTypeIndex) -> FuncType { - let ty = engine.signatures().root(index).expect( - "VMSharedTypeIndex is not registered in the Engine. Wrong \ + let ty = RegisteredType::root(engine, index).expect( + "VMSharedTypeIndex is not registered in the Engine! Wrong \ engine? Didn't root the index somewhere?", ); Self { ty }