Simplify handling of caller-vmctx and dyn VMStore pointers (#10542)

* Simplify handling of caller-vmctx and `dyn VMStore` pointers

Just keep a copy of the trait object pointer in `StoreOpaque`, rather than
trying to pluck it out of the default caller.

Clean up `CallThreadState` so that it doesn't need a caller-vmctx.

* preserve comment that accidentally got removed

Author: fitzgen

crates/wasmtime/src/runtime/store.rs

@@ -346,6 +346,7 @@ pub struct StoreOpaque {
     /// `rooted_host_funcs` below. This structure contains pointers which are
     /// otherwise kept alive by the `Arc` references in `rooted_host_funcs`.
     store_data: ManuallyDrop<StoreData>,
+    traitobj: StorePtr,
     default_caller: InstanceHandle,
 
     /// Used to optimized wasm->host calls when the host function is defined with
@@ -517,67 +518,83 @@ impl<T> Store<T> {
     /// tables created to 10,000. This can be overridden with the
     /// [`Store::limiter`] configuration method.
     pub fn new(engine: &Engine, data: T) -> Self {
+        let store_data = StoreData::new();
+        log::trace!("creating new store {:?}", store_data.id());
+
         let pkey = engine.allocator().next_available_pkey();
 
-        let mut inner = Box::new(StoreInner {
-            inner: StoreOpaque {
-                _marker: marker::PhantomPinned,
-                engine: engine.clone(),
-                vm_store_context: Default::default(),
-                instances: Vec::new(),
-                #[cfg(feature = "component-model")]
-                num_component_instances: 0,
-                signal_handler: None,
-                gc_store: None,
-                gc_roots: RootSet::default(),
-                #[cfg(feature = "gc")]
-                gc_roots_list: GcRootsList::default(),
-                #[cfg(feature = "gc")]
-                gc_host_alloc_types: Default::default(),
-                modules: ModuleRegistry::default(),
-                func_refs: FuncRefs::default(),
-                host_globals: Vec::new(),
-                instance_count: 0,
-                instance_limit: crate::DEFAULT_INSTANCE_LIMIT,
-                memory_count: 0,
-                memory_limit: crate::DEFAULT_MEMORY_LIMIT,
-                table_count: 0,
-                table_limit: crate::DEFAULT_TABLE_LIMIT,
-                #[cfg(feature = "async")]
-                async_state: AsyncState::default(),
-                fuel_reserve: 0,
-                fuel_yield_interval: None,
-                store_data: ManuallyDrop::new(StoreData::new()),
-                default_caller: InstanceHandle::null(),
-                hostcall_val_storage: Vec::new(),
-                wasm_val_raw_storage: Vec::new(),
-                rooted_host_funcs: ManuallyDrop::new(Vec::new()),
-                pkey,
-                #[cfg(feature = "component-model")]
-                component_host_table: Default::default(),
-                #[cfg(feature = "component-model")]
-                component_calls: Default::default(),
-                #[cfg(feature = "component-model")]
-                host_resource_data: Default::default(),
-                #[cfg(has_host_compiler_backend)]
-                executor: if cfg!(feature = "pulley") && engine.target().is_pulley() {
-                    Executor::Interpreter(Interpreter::new(engine))
-                } else {
-                    Executor::Native
-                },
-                #[cfg(not(has_host_compiler_backend))]
-                executor: {
-                    debug_assert!(engine.target().is_pulley());
-                    Executor::Interpreter(Interpreter::new(engine))
-                },
+        let inner = StoreOpaque {
+            _marker: marker::PhantomPinned,
+            engine: engine.clone(),
+            vm_store_context: Default::default(),
+            instances: Vec::new(),
+            #[cfg(feature = "component-model")]
+            num_component_instances: 0,
+            signal_handler: None,
+            gc_store: None,
+            gc_roots: RootSet::default(),
+            #[cfg(feature = "gc")]
+            gc_roots_list: GcRootsList::default(),
+            #[cfg(feature = "gc")]
+            gc_host_alloc_types: Default::default(),
+            modules: ModuleRegistry::default(),
+            func_refs: FuncRefs::default(),
+            host_globals: Vec::new(),
+            instance_count: 0,
+            instance_limit: crate::DEFAULT_INSTANCE_LIMIT,
+            memory_count: 0,
+            memory_limit: crate::DEFAULT_MEMORY_LIMIT,
+            table_count: 0,
+            table_limit: crate::DEFAULT_TABLE_LIMIT,
+            #[cfg(feature = "async")]
+            async_state: AsyncState::default(),
+            fuel_reserve: 0,
+            fuel_yield_interval: None,
+            store_data: ManuallyDrop::new(store_data),
+            traitobj: StorePtr::empty(),
+            default_caller: InstanceHandle::null(),
+            hostcall_val_storage: Vec::new(),
+            wasm_val_raw_storage: Vec::new(),
+            rooted_host_funcs: ManuallyDrop::new(Vec::new()),
+            pkey,
+            #[cfg(feature = "component-model")]
+            component_host_table: Default::default(),
+            #[cfg(feature = "component-model")]
+            component_calls: Default::default(),
+            #[cfg(feature = "component-model")]
+            host_resource_data: Default::default(),
+            #[cfg(has_host_compiler_backend)]
+            executor: if cfg!(feature = "pulley") && engine.target().is_pulley() {
+                Executor::Interpreter(Interpreter::new(engine))
+            } else {
+                Executor::Native
             },
+            #[cfg(not(has_host_compiler_backend))]
+            executor: {
+                debug_assert!(engine.target().is_pulley());
+                Executor::Interpreter(Interpreter::new(engine))
+            },
+        };
+        let mut inner = Box::new(StoreInner {
+            inner,
             limiter: None,
             call_hook: None,
             #[cfg(target_has_atomic = "64")]
             epoch_deadline_behavior: None,
             data: ManuallyDrop::new(data),
         });
 
+        // Note the erasure of the lifetime here into `'static`, so in general
+        // usage of this trait object must be strictly bounded to the `Store`
+        // itself, and this is an invariant that we have to maintain throughout
+        // Wasmtime.
+        inner.traitobj = StorePtr::new(unsafe {
+            mem::transmute::<
+                NonNull<dyn crate::runtime::vm::VMStore + '_>,
+                NonNull<dyn crate::runtime::vm::VMStore + 'static>,
+            >(NonNull::from(&mut *inner))
+        });
+
         // Wasmtime uses the callee argument to host functions to learn about
         // the original pointer to the `Store` itself, allowing it to
         // reconstruct a `StoreContextMut<T>`. When we initially call a `Func`,
@@ -589,9 +606,11 @@ impl<T> Store<T> {
             let module = Arc::new(wasmtime_environ::Module::default());
             let shim = ModuleRuntimeInfo::bare(module);
             let allocator = OnDemandInstanceAllocator::default();
+
             allocator
                 .validate_module(shim.env_module(), shim.offsets())
                 .unwrap();
+
             let mut instance = unsafe {
                 allocator
                     .allocate_module(InstanceAllocationRequest {
@@ -605,19 +624,10 @@ impl<T> Store<T> {
                     })
                     .expect("failed to allocate default callee")
             };
-
-            // Note the erasure of the lifetime here into `'static`, so in
-            // general usage of this trait object must be strictly bounded to
-            // the `Store` itself, and this is an invariant that we have to
-            // maintain throughout Wasmtime.
             unsafe {
-                let traitobj = mem::transmute::<
-                    NonNull<dyn crate::runtime::vm::VMStore + '_>,
-                    NonNull<dyn crate::runtime::vm::VMStore + 'static>,
-                >(NonNull::from(&mut *inner));
-                instance.set_store(traitobj);
-                instance
+                instance.set_store(Some(inner.traitobj()));
             }
+            instance
         };
 
         Self {
@@ -1705,7 +1715,12 @@ impl StoreOpaque {
 
     #[inline]
     pub fn traitobj(&self) -> NonNull<dyn crate::runtime::vm::VMStore> {
-        self.default_caller.traitobj(self)
+        self.traitobj.as_raw().unwrap()
+    }
+
+    #[inline]
+    pub fn traitobj_mut(&mut self) -> &mut dyn crate::runtime::vm::VMStore {
+        unsafe { self.traitobj().as_mut() }
     }
 
     /// Takes the cached `Vec<Val>` stored internally across hostcalls to get

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

@@ -1712,8 +1712,8 @@ impl InstanceHandle {
     ///
     /// This is provided for the original `Store` itself to configure the first
     /// self-pointer after the original `Box` has been initialized.
-    pub unsafe fn set_store(&mut self, store: NonNull<dyn VMStore>) {
-        self.instance_mut().set_store(Some(store));
+    pub unsafe fn set_store(&mut self, store: Option<NonNull<dyn VMStore>>) {
+        self.instance_mut().set_store(store);
     }
 
     /// Returns a clone of this instance.

crates/wasmtime/src/runtime/vm/instance/allocator.rs

@@ -94,6 +94,13 @@ pub struct InstanceAllocationRequest<'a> {
 /// InstanceAllocationRequest.
 pub struct StorePtr(Option<NonNull<dyn VMStore>>);
 
+// We can't make `VMStore: Send + Sync` because that requires making all of
+// Wastime's internals generic over the `Store`'s `T`. So instead, we take care
+// in the whole VM layer to only use the `VMStore` in ways that are `Send`- and
+// `Sync`-safe and we have to have these unsafe impls.
+unsafe impl Send for StorePtr {}
+unsafe impl Sync for StorePtr {}
+
 impl StorePtr {
     /// A pointer to no Store.
     pub fn empty() -> Self {

crates/wasmtime/src/runtime/vm/traphandlers.rs

@@ -19,7 +19,7 @@ use crate::prelude::*;
 use crate::runtime::module::lookup_code;
 use crate::runtime::store::{ExecutorRef, StoreOpaque};
 use crate::runtime::vm::sys::traphandlers;
-use crate::runtime::vm::{Instance, InterpreterRef, VMContext, VMStoreContext};
+use crate::runtime::vm::{InterpreterRef, VMContext, VMStoreContext};
 use crate::{StoreContextMut, WasmBacktrace};
 use core::cell::Cell;
 use core::num::NonZeroU32;
@@ -371,7 +371,8 @@ where
     F: FnMut(NonNull<VMContext>, Option<InterpreterRef<'_>>) -> bool,
 {
     let caller = store.0.default_caller();
-    let result = CallThreadState::new(store.0, caller).with(|cx| match store.0.executor() {
+
+    let result = CallThreadState::new(store.0).with(|cx| match store.0.executor() {
         // In interpreted mode directly invoke the host closure since we won't
         // be using host-based `setjmp`/`longjmp` as that's not going to save
         // the context we want.
@@ -474,14 +475,7 @@ mod call_thread_state {
         pub const JMP_BUF_INTERPRETER_SENTINEL: *mut u8 = 1 as *mut u8;
 
         #[inline]
-        pub(super) fn new(store: &mut StoreOpaque, caller: NonNull<VMContext>) -> CallThreadState {
-            let vm_store_context = unsafe {
-                Instance::from_vmctx(caller, |i| i.vm_store_context())
-                    .read()
-                    .unwrap()
-                    .as_non_null()
-            };
-
+        pub(super) fn new(store: &mut StoreOpaque) -> CallThreadState {
             // Don't try to plumb #[cfg] everywhere for this field, just pretend
             // we're using it on miri/windows to silence compiler warnings.
             let _: Range<_> = store.async_guard_range();
@@ -495,18 +489,18 @@ mod call_thread_state {
                 capture_backtrace: store.engine().config().wasm_backtrace,
                 #[cfg(feature = "coredump")]
                 capture_coredump: store.engine().config().coredump_on_trap,
-                vm_store_context,
+                vm_store_context: store.vm_store_context_ptr(),
                 #[cfg(all(has_native_signals, unix))]
                 async_guard_range: store.async_guard_range(),
                 prev: Cell::new(ptr::null()),
                 old_last_wasm_exit_fp: Cell::new(unsafe {
-                    *vm_store_context.as_ref().last_wasm_exit_fp.get()
+                    *store.vm_store_context().last_wasm_exit_fp.get()
                 }),
                 old_last_wasm_exit_pc: Cell::new(unsafe {
-                    *vm_store_context.as_ref().last_wasm_exit_pc.get()
+                    *store.vm_store_context().last_wasm_exit_pc.get()
                 }),
                 old_last_wasm_entry_fp: Cell::new(unsafe {
-                    *vm_store_context.as_ref().last_wasm_entry_fp.get()
+                    *store.vm_store_context().last_wasm_entry_fp.get()
                 }),
             }
         }