Add a helper always-Sync utility to Wasmtime (bytecodealliance#11453)

* Add a helper always-`Sync` utility to Wasmtime

This commit adds a newtype wrapper to Wasmtime, `AlwaysMut`, which is
unconditionally `Sync` if the stored type is `Send`. This is similar to
a `Mutex<T>` where it promotes a `Send` bound to a `Sync` bound, but
it's unlike `Mutex<T>` in that `AlwaysMut<T>` has no synchronization.
The reason that this is safe is that `AlwaysMut<T>` completely disallows
access to the underlying data through `&self` and requires `&mut self`.
This is similar to how `Mutex::get_mut` is safe, for example.

This type cleans up a preexisting `unsafe impl Sync` block in funcref
management around `SendSyncBump` (`bumpalo::Bump` is `Send`, not `Sync`,
but we only access it through `&mut self`). This then additionally
removes `unsafe impl Sync for StoreFiber` which, upon reflection, is not
sound because we don't ever constraint the store's `T` type to `Sync`,
only `Send`. This is effectively no change throughout Wasmtime, however,
as fibers are only accessed with `&mut`.

Overall this is mostly just internal refactoring to reduce the amount of
`unsafe` inside of Wasmtime and to add a new utility to use in the
future too.

* Fix configured build

* Review comments

* Fix configured build

Author: alexcrichton

crates/wasmtime/src/runtime/fiber.rs

@@ -1,20 +1,21 @@
 use crate::prelude::*;
 use crate::store::{Executor, StoreId, StoreInner, StoreOpaque};
 use crate::vm::mpk::{self, ProtectionMask};
-use crate::vm::{AsyncWasmCallState, VMStore};
+use crate::vm::{AlwaysMut, AsyncWasmCallState, VMStore};
 use crate::{Engine, StoreContextMut};
 use anyhow::{Result, anyhow};
 use core::mem;
 use core::ops::Range;
 use core::pin::Pin;
 use core::ptr::{self, NonNull};
 use core::task::{Context, Poll};
-use wasmtime_fiber::{Fiber, Suspend};
+use wasmtime_fiber::{Fiber, FiberStack, Suspend};
 
 type WasmtimeResume = Result<NonNull<Context<'static>>>;
 type WasmtimeYield = StoreFiberYield;
 type WasmtimeComplete = Result<()>;
 type WasmtimeSuspend = Suspend<WasmtimeResume, WasmtimeYield, WasmtimeComplete>;
+type WasmtimeFiber<'a> = Fiber<'a, WasmtimeResume, WasmtimeYield, WasmtimeComplete>;
 
 /// State related to asynchronous computations stored within a `Store<T>`.
 ///
@@ -434,9 +435,9 @@ pub(crate) struct StoreFiber<'a> {
     ///
     /// Note also that every `StoreFiber` is implicitly granted exclusive access
     /// to the store when it is resumed.
-    fiber: Option<Fiber<'a, WasmtimeResume, WasmtimeYield, WasmtimeComplete>>,
+    fiber: Option<AlwaysMut<RawFiber<'a>>>,
     /// See `FiberResumeState`
-    state: Option<FiberResumeState>,
+    state: Option<AlwaysMut<FiberResumeState>>,
     /// The Wasmtime `Engine` to which this fiber belongs.
     engine: Engine,
     /// The id of the store with which this fiber was created.
@@ -446,9 +447,23 @@ pub(crate) struct StoreFiber<'a> {
     id: StoreId,
 }
 
-impl StoreFiber<'_> {
+struct RawFiber<'a>(WasmtimeFiber<'a>);
+
+impl<'a> StoreFiber<'a> {
+    /// Convenience method to peel off some layers of abstraction around the raw
+    /// `wasmtime_fiber::Fiber`.
+    fn fiber(&mut self) -> Option<&mut WasmtimeFiber<'a>> {
+        Some(&mut self.fiber.as_mut()?.get_mut().0)
+    }
+
+    /// Convenience method take the internal fiber and consume it, yielding its
+    /// original stack.
+    fn take_fiber_stack(&mut self) -> Option<FiberStack> {
+        self.fiber.take().map(|f| f.into_inner().0.into_stack())
+    }
+
     pub(crate) fn dispose(&mut self, store: &mut StoreOpaque) {
-        if let Some(fiber) = &mut self.fiber {
+        if let Some(fiber) = self.fiber() {
             if !fiber.done() {
                 let result = resume_fiber(store, self, Err(anyhow!("future dropped")));
                 debug_assert!(result.is_ok());
@@ -469,16 +484,15 @@ impl Drop for StoreFiber<'_> {
         }
 
         assert!(
-            self.fiber.as_ref().unwrap().done(),
+            self.fiber().unwrap().done(),
             "attempted to drop in-progress fiber without first calling `StoreFiber::dispose`"
         );
 
-        self.state.take().unwrap().dispose();
+        self.state.take().unwrap().into_inner().dispose();
 
         unsafe {
-            self.engine
-                .allocator()
-                .deallocate_fiber_stack(self.fiber.take().unwrap().into_stack());
+            let stack = self.take_fiber_stack().unwrap();
+            self.engine.allocator().deallocate_fiber_stack(stack);
         }
     }
 }
@@ -537,11 +551,7 @@ impl Drop for StoreFiber<'_> {
 // declare the fiber as only containing Send data on its stack, despite not
 // knowing for sure at compile time that this is correct. That's what `unsafe`
 // in Rust is all about, though, right?
-unsafe impl Send for StoreFiber<'_> {}
-// See the docs about the `Send` impl above, which also apply to this `Sync`
-// impl.  `Sync` is needed since we store `StoreFiber`s and switch between them
-// when executing components that export async-lifted functions.
-unsafe impl Sync for StoreFiber<'_> {}
+unsafe impl Send for RawFiber<'_> {}
 
 /// State of the world when a fiber last suspended.
 ///
@@ -596,8 +606,7 @@ impl FiberResumeState {
         let tls = unsafe { self.tls.push() };
         let mpk = swap_mpk_states(self.mpk);
         let async_guard_range = fiber
-            .fiber
-            .as_ref()
+            .fiber()
             .unwrap()
             .stack()
             .guard_range()
@@ -723,25 +732,30 @@ fn resume_fiber<'a>(
 
     impl Drop for Restore<'_, '_> {
         fn drop(&mut self) {
-            self.fiber.state = Some(unsafe { self.state.take().unwrap().replace(self.store) });
+            self.fiber.state =
+                Some(unsafe { self.state.take().unwrap().replace(self.store).into() });
         }
     }
     let result = unsafe {
-        let prev = fiber.state.take().unwrap().replace(store, fiber);
+        let prev = fiber
+            .state
+            .take()
+            .unwrap()
+            .into_inner()
+            .replace(store, fiber);
         let restore = Restore {
             store,
             fiber,
             state: Some(prev),
         };
-        restore.fiber.fiber.as_ref().unwrap().resume(result)
+        restore.fiber.fiber().unwrap().resume(result)
     };
 
     match &result {
         // The fiber has finished, so recycle its stack by disposing of the
         // underlying fiber itself.
         Ok(_) => {
-            let stack = fiber.fiber.take().map(|f| f.into_stack());
-            if let Some(stack) = stack {
+            if let Some(stack) = fiber.take_fiber_stack() {
                 store.deallocate_fiber_stack(stack);
             }
         }
@@ -761,7 +775,7 @@ fn resume_fiber<'a>(
             // pointer Wasmtime uses is not pointing anywhere within the
             // stack. If it is then that's a bug indicating that TLS management
             // in Wasmtime is incorrect.
-            if let Some(range) = fiber.fiber.as_ref().unwrap().stack().range() {
+            if let Some(range) = fiber.fiber().unwrap().stack().range() {
                 AsyncWasmCallState::assert_current_state_not_in_range(range);
             }
         }
@@ -826,19 +840,22 @@ pub(crate) fn make_fiber<'a>(
         fun(reset.0)
     })?;
     Ok(StoreFiber {
-        state: Some(FiberResumeState {
-            tls: crate::runtime::vm::AsyncWasmCallState::new(),
-            mpk: if track_pkey_context_switch {
-                Some(ProtectionMask::all())
-            } else {
-                None
-            },
-            stack_limit: usize::MAX,
-            executor,
-        }),
+        state: Some(
+            FiberResumeState {
+                tls: crate::runtime::vm::AsyncWasmCallState::new(),
+                mpk: if track_pkey_context_switch {
+                    Some(ProtectionMask::all())
+                } else {
+                    None
+                },
+                stack_limit: usize::MAX,
+                executor,
+            }
+            .into(),
+        ),
         engine,
         id,
-        fiber: Some(fiber),
+        fiber: Some(RawFiber(fiber).into()),
     })
 }
 

crates/wasmtime/src/runtime/store/func_refs.rs

@@ -5,7 +5,7 @@ use crate::Definition;
 use crate::module::ModuleRegistry;
 use crate::prelude::*;
 use crate::runtime::HostFunc;
-use crate::runtime::vm::{SendSyncPtr, VMArrayCallHostFuncContext, VMFuncRef};
+use crate::runtime::vm::{AlwaysMut, SendSyncPtr, VMArrayCallHostFuncContext, VMFuncRef};
 use alloc::sync::Arc;
 use core::ptr::NonNull;
 
@@ -18,7 +18,7 @@ use core::ptr::NonNull;
 pub struct FuncRefs {
     /// A bump allocation arena where we allocate `VMFuncRef`s such
     /// that they are pinned and owned.
-    bump: SendSyncBump,
+    bump: AlwaysMut<bumpalo::Bump>,
 
     /// Pointers into `self.bump` for entries that need `wasm_call` field filled
     /// in.
@@ -77,23 +77,6 @@ enum Storage {
     },
 }
 
-use send_sync_bump::SendSyncBump;
-mod send_sync_bump {
-    #[derive(Default)]
-    pub struct SendSyncBump(bumpalo::Bump);
-
-    impl SendSyncBump {
-        pub fn alloc<T>(&mut self, val: T) -> &mut T {
-            self.0.alloc(val)
-        }
-    }
-
-    // Safety: We require `&mut self` on the only public method, which means it
-    // is safe to send `&SendSyncBump` references across threads because they
-    // can't actually do anything with it.
-    unsafe impl Sync for SendSyncBump {}
-}
-
 impl FuncRefs {
     /// Push the given `VMFuncRef` into this arena, returning a
     /// pinned pointer to it.
@@ -109,7 +92,7 @@ impl FuncRefs {
         modules: &ModuleRegistry,
     ) -> NonNull<VMFuncRef> {
         debug_assert!(func_ref.wasm_call.is_none());
-        let func_ref = self.bump.alloc(func_ref);
+        let func_ref = self.bump.get_mut().alloc(func_ref);
         // SAFETY: it's a contract of this function itself that `func_ref` has a
         // valid vmctx field to read.
         let has_hole = unsafe { !try_fill(func_ref, modules) };

crates/wasmtime/src/runtime/vm.rs

@@ -51,6 +51,7 @@ use wasmtime_environ::{
 #[cfg(feature = "gc")]
 use wasmtime_environ::ModuleInternedTypeIndex;
 
+mod always_mut;
 #[cfg(feature = "component-model")]
 pub mod component;
 mod const_expr;
@@ -92,6 +93,7 @@ pub(crate) use interpreter_disabled as interpreter;
 #[cfg(feature = "debug-builtins")]
 pub use wasmtime_jit_debug::gdb_jit_int::GdbJitImageRegistration;
 
+pub use crate::runtime::vm::always_mut::*;
 pub use crate::runtime::vm::export::*;
 pub use crate::runtime::vm::gc::*;
 pub use crate::runtime::vm::imports::Imports;

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

@@ -0,0 +1,52 @@
+use core::fmt;
+
+/// A helper types that is `Send` if `T` is `Sync`.
+///
+/// This structure is a newtype wrapper around the `T` type parameter. What
+/// makes this a utility is the fact that it contains an `unsafe impl Sync`
+/// implementation for the type when `T` is `Send`. This is then coupled with
+/// the fact that there is no ability to access a shared reference, `&T`, from
+/// this type. Instead all access is done through `&mut T`.
+///
+/// This means that accessing the `T` in this data structure always requires
+/// exclusive `&mut self` access. This provides the trivial guarantee that this
+/// type is safe to share across threads because if you do so then you're just
+/// not able to do anything with it.
+#[derive(Default /* Do not derive traits with &self here, that's not sound */)]
+#[repr(transparent)]
+pub struct AlwaysMut<T>(T);
+
+// SAFETY: this is the purpose for existence of this type, meaning that if `T`
+// is `Send` then this type is `Sync` because it statically disallows shared
+// access to `T`.
+unsafe impl<T: Send> Sync for AlwaysMut<T> {}
+
+impl<T> AlwaysMut<T> {
+    /// Creates a new [`AlwaysMut`] with the provided value.
+    pub fn new(value: T) -> AlwaysMut<T> {
+        AlwaysMut(value)
+    }
+
+    /// Return a mutable reference to the underlying data in this [`AlwaysMut`]
+    pub fn get_mut(&mut self) -> &mut T {
+        &mut self.0
+    }
+
+    /// Consume this [`AlwaysMut`], returning the underlying data.
+    #[cfg(feature = "async")]
+    pub fn into_inner(self) -> T {
+        self.0
+    }
+}
+
+impl<T> From<T> for AlwaysMut<T> {
+    fn from(val: T) -> AlwaysMut<T> {
+        AlwaysMut::new(val)
+    }
+}
+
+impl<T> fmt::Debug for AlwaysMut<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("AlwaysMut").finish_non_exhaustive()
+    }
+}

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

@@ -995,6 +995,15 @@ pub(crate) mod tls {
         state: raw::Ptr,
     }
 
+    // SAFETY: This is a relatively unsafe unsafe block and not really all that
+    // well audited. The general idea is that the linked list of activations
+    // owned by `self.state` are safe to send to other threads, but that relies
+    // on everything internally being safe as well as stack variables and such.
+    // This is more-or-less tied to the very large comment in `fiber.rs` about
+    // `unsafe impl Send` there.
+    #[cfg(feature = "async")]
+    unsafe impl Send for AsyncWasmCallState {}
+
     #[cfg(feature = "async")]
     impl AsyncWasmCallState {
         /// Creates new state that initially starts as null.