Always use mmap in cranelift-jit (#13180)

* Always use mmap in cranelift-jit

This was previously not done with the theory that using the regular heap
allocator could be faster. We now have an arena memory provider and
using mmap is necessary anyway when SELinux is enabled. I don't think
the benefits outweight the code complexity cost.

* Use memmap2 on Windows too in cranelift-jit

Author: bjorn3

cranelift/filetests/Cargo.toml

@@ -19,7 +19,7 @@ cranelift-frontend = { workspace = true }
 cranelift-interpreter = { workspace = true }
 cranelift-native = { workspace = true }
 cranelift-reader = { workspace = true }
-cranelift-jit = { workspace = true, features = ["selinux-fix", "wasmtime-unwinder"] }
+cranelift-jit = { workspace = true, features = ["wasmtime-unwinder"] }
 cranelift-module = { workspace = true }
 cranelift-control = { workspace = true }
 wasmtime-unwinder = { workspace = true, features = ["cranelift"] }

cranelift/jit/Cargo.toml

@@ -24,7 +24,7 @@ anyhow = { workspace = true }
 region = "3.0.2"
 libc = { workspace = true }
 target-lexicon = { workspace = true }
-memmap2 = { version = "0.2.1", optional = true }
+memmap2 = "0.2.1"
 log = { workspace = true }
 wasmtime-jit-icache-coherence = { workspace = true }
 
@@ -37,7 +37,6 @@ features = [
 ]
 
 [features]
-selinux-fix = ['memmap2']
 default = []
 
 wasmtime-unwinder = ["dep:wasmtime-unwinder"]

cranelift/jit/src/memory/system.rs

@@ -1,21 +1,15 @@
-use cranelift_module::{ModuleError, ModuleResult};
-
-#[cfg(all(not(target_os = "windows"), feature = "selinux-fix"))]
-use memmap2::MmapMut;
-
-#[cfg(not(any(feature = "selinux-fix", windows)))]
-use std::alloc;
 use std::io;
 use std::mem;
 use std::ptr;
 
+use cranelift_module::{ModuleError, ModuleResult};
+use memmap2::MmapMut;
+
 use super::{BranchProtection, JITMemoryKind, JITMemoryProvider};
 
 /// A simple struct consisting of a pointer and length.
 struct PtrLen {
-    #[cfg(all(not(target_os = "windows"), feature = "selinux-fix"))]
     map: Option<MmapMut>,
-
     ptr: *mut u8,
     len: usize,
 }
@@ -24,17 +18,14 @@ impl PtrLen {
     /// Create a new empty `PtrLen`.
     fn new() -> Self {
         Self {
-            #[cfg(all(not(target_os = "windows"), feature = "selinux-fix"))]
             map: None,
-
             ptr: ptr::null_mut(),
             len: 0,
         }
     }
 
     /// Create a new `PtrLen` pointing to at least `size` bytes of memory,
     /// suitably sized and aligned for memory protection.
-    #[cfg(all(not(target_os = "windows"), feature = "selinux-fix"))]
     fn with_size(size: usize) -> io::Result<Self> {
         let alloc_size = region::page::ceil(size as *const ()) as usize;
         MmapMut::map_anon(alloc_size).map(|mut mmap| {
@@ -47,70 +38,8 @@ impl PtrLen {
             }
         })
     }
-
-    #[cfg(all(not(target_os = "windows"), not(feature = "selinux-fix")))]
-    fn with_size(size: usize) -> io::Result<Self> {
-        assert_ne!(size, 0);
-        let page_size = region::page::size();
-        let alloc_size = region::page::ceil(size as *const ()) as usize;
-        let layout = alloc::Layout::from_size_align(alloc_size, page_size).unwrap();
-        // Safety: We assert that the size is non-zero above.
-        let ptr = unsafe { alloc::alloc(layout) };
-
-        if !ptr.is_null() {
-            Ok(Self {
-                ptr,
-                len: alloc_size,
-            })
-        } else {
-            Err(io::Error::from(io::ErrorKind::OutOfMemory))
-        }
-    }
-
-    #[cfg(target_os = "windows")]
-    fn with_size(size: usize) -> io::Result<Self> {
-        use windows_sys::Win32::System::Memory::{
-            MEM_COMMIT, MEM_RESERVE, PAGE_READWRITE, VirtualAlloc,
-        };
-
-        // VirtualAlloc always rounds up to the next multiple of the page size
-        let ptr = unsafe {
-            VirtualAlloc(
-                ptr::null_mut(),
-                size,
-                MEM_COMMIT | MEM_RESERVE,
-                PAGE_READWRITE,
-            )
-        };
-        if !ptr.is_null() {
-            Ok(Self {
-                ptr: ptr as *mut u8,
-                len: region::page::ceil(size as *const ()) as usize,
-            })
-        } else {
-            Err(io::Error::last_os_error())
-        }
-    }
 }
 
-// `MMapMut` from `cfg(feature = "selinux-fix")` already deallocates properly.
-#[cfg(all(not(target_os = "windows"), not(feature = "selinux-fix")))]
-impl Drop for PtrLen {
-    fn drop(&mut self) {
-        if !self.ptr.is_null() {
-            let page_size = region::page::size();
-            let layout = alloc::Layout::from_size_align(self.len, page_size).unwrap();
-            unsafe {
-                region::protect(self.ptr, self.len, region::Protection::READ_WRITE)
-                    .expect("unable to unprotect memory");
-                alloc::dealloc(self.ptr, layout)
-            }
-        }
-    }
-}
-
-// TODO: add a `Drop` impl for `cfg(target_os = "windows")`
-
 /// JIT memory manager. This manages pages of suitably aligned and
 /// accessible memory. Memory will be leaked by default to have
 /// function pointers remain valid for the remainder of the
@@ -201,13 +130,9 @@ impl Memory {
 
     /// Iterates non protected memory allocations that are of not zero bytes in size.
     fn non_protected_allocations_iter(&self) -> impl Iterator<Item = &PtrLen> {
-        let iter = self.allocations[self.already_protected..].iter();
-
-        #[cfg(all(not(target_os = "windows"), feature = "selinux-fix"))]
-        return iter.filter(|&PtrLen { map, len, .. }| *len != 0 && map.is_some());
-
-        #[cfg(any(target_os = "windows", not(feature = "selinux-fix")))]
-        return iter.filter(|&PtrLen { len, .. }| *len != 0);
+        self.allocations[self.already_protected..]
+            .iter()
+            .filter(|&PtrLen { map, len, .. }| *len != 0 && map.is_some())
     }
 
     /// Frees all allocated memory regions that would be leaked otherwise.

crates/wasmtime/src/runtime/func.rs

@@ -1064,6 +1064,10 @@ impl Func {
         self.vm_func_ref(store).as_ptr().cast()
     }
 
+    pub(crate) fn to_val_raw(&self, store: &mut StoreOpaque) -> ValRaw {
+        ValRaw::funcref(self.to_raw_(store))
+    }
+
     /// Invokes this function with the `params` given, returning the results
     /// asynchronously.
     ///

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

@@ -13,7 +13,7 @@ use crate::runtime::vm::vmcontext::{
     VMTableDefinition, VMTableImport, VMTagDefinition, VMTagImport,
 };
 use crate::runtime::vm::{
-    GcStore, HostResult, Imports, ModuleRuntimeInfo, SendSyncPtr, VMGcRef, VMGlobalKind, VMStore,
+    GcStore, HostResult, Imports, ModuleRuntimeInfo, SendSyncPtr, VMGlobalKind, VMStore,
     VMStoreRawPtr, VmPtr, VmSafe, WasmFault, catch_unwind_and_record_trap,
 };
 use crate::store::{
@@ -133,7 +133,7 @@ pub struct Instance {
     //
     // TODO(#12621): This should be a `TrySecondaryMap<PassiveElemIndex, _>`
     // but that type is currently footgun-y / isn't actually OOM-safe yet.
-    passive_elements: TryVec<PassiveElementSegment>,
+    passive_elements: TryVec<Option<(NeedsGcRooting, TryVec<ValRaw>)>>,
 
     /// Stores the dropped passive data segments in this instantiation by index.
     /// If the index is present in the set, the segment has been dropped.
@@ -225,8 +225,8 @@ impl Instance {
         gc_roots: &mut crate::vm::GcRootsList,
     ) {
         for segment in self.passive_elements_mut().iter_mut() {
-            if segment.needs_gc_rooting() == NeedsGcRooting::Yes {
-                for e in segment.elements() {
+            if let Some((wasmtime_environ::NeedsGcRooting::Yes, elems)) = segment {
+                for e in elems {
                     let Some(root) = e.as_vmgc_ref_ptr() else {
                         continue;
                     };
@@ -912,12 +912,16 @@ impl Instance {
             return &[];
         };
 
-        self.passive_elements[passive.index()].elements()
+        let Some((_, seg)) = &self.passive_elements[passive.index()] else {
+            return &[];
+        };
+
+        &**seg
     }
 
     pub(crate) fn passive_elements_mut(
         self: Pin<&mut Self>,
-    ) -> Pin<&mut TryVec<PassiveElementSegment>> {
+    ) -> Pin<&mut TryVec<Option<(NeedsGcRooting, TryVec<ValRaw>)>>> {
         // SAFETY: Not moving data out of `self`.
         Pin::new(&mut unsafe { self.get_unchecked_mut() }.passive_elements)
     }
@@ -992,7 +996,6 @@ impl Instance {
     /// Drop an element.
     pub(crate) fn elem_drop(
         self: Pin<&mut Self>,
-        gc_store: Option<&mut GcStore>,
         elem_index: ElemIndex,
     ) -> Result<(), OutOfMemory> {
         // https://webassembly.github.io/reference-types/core/exec/instructions.html#exec-elem-drop
@@ -1007,7 +1010,7 @@ impl Instance {
             return Ok(());
         };
 
-        self.passive_elements_mut()[passive_index.index()].clear(gc_store);
+        self.passive_elements_mut()[passive_index.index()] = None;
         Ok(())
     }
 
@@ -1919,77 +1922,3 @@ impl<T: InstanceLayout> Drop for OwnedInstance<T> {
         }
     }
 }
-
-#[derive(Debug)]
-pub(crate) struct PassiveElementSegment {
-    needs_gc_rooting: NeedsGcRooting,
-    elements: TryVec<ValRaw>,
-}
-
-impl PassiveElementSegment {
-    /// Create a new passive element segment with the given capacity.
-    pub(crate) fn new(
-        needs_gc_rooting: NeedsGcRooting,
-        capacity: usize,
-    ) -> Result<Self, OutOfMemory> {
-        Ok(Self {
-            needs_gc_rooting,
-            elements: TryVec::with_capacity(capacity)?,
-        })
-    }
-
-    /// Push a value onto this passive element segment.
-    ///
-    /// NB: Does not type check the value, relies on callers to ensure the value
-    /// is of the correct type (generally, due to validation).
-    pub(crate) fn push(&mut self, store: &mut StoreOpaque, val: Val) -> Result<()> {
-        let val = {
-            let mut store = AutoAssertNoGc::new(store);
-            val.to_raw_(&mut store)?
-        };
-        let val = self.clone_gc_ref(store, val);
-        self.elements.push(val)?;
-        Ok(())
-    }
-
-    fn clone_gc_ref(&mut self, store: &mut StoreOpaque, val: ValRaw) -> ValRaw {
-        if let NeedsGcRooting::Yes = self.needs_gc_rooting {
-            let gc_ref = val.get_anyref();
-            if let Some(gc_ref) = VMGcRef::from_raw_u32(gc_ref) {
-                if let Some(gc_store) = store.optional_gc_store_mut() {
-                    return ValRaw::anyref(gc_store.clone_gc_ref(&gc_ref).as_raw_u32());
-                }
-            }
-        }
-        val
-    }
-
-    /// Clear this segment's elements.
-    pub(crate) fn clear(&mut self, mut gc_store: Option<&mut GcStore>) {
-        for val in mem::take(&mut self.elements) {
-            self.drop_gc_ref(&mut gc_store, val);
-        }
-    }
-
-    fn drop_gc_ref(&mut self, gc_store: &mut Option<&mut GcStore>, val: ValRaw) {
-        if let NeedsGcRooting::Yes = self.needs_gc_rooting {
-            let gc_ref = val.get_anyref();
-            if let Some(gc_ref) = VMGcRef::from_raw_u32(gc_ref) {
-                if let Some(gc_store) = gc_store.as_deref_mut() {
-                    let _ = gc_store.drop_gc_ref(gc_ref);
-                }
-            }
-        }
-    }
-
-    /// Whether this segment needs GC rooting and tracing.
-    #[cfg(feature = "gc")]
-    pub(crate) fn needs_gc_rooting(&self) -> NeedsGcRooting {
-        self.needs_gc_rooting
-    }
-
-    /// The elements of this segment.
-    pub(crate) fn elements(&self) -> &[ValRaw] {
-        &self.elements
-    }
-}

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

@@ -6,8 +6,7 @@ use crate::runtime::vm::memory::Memory;
 use crate::runtime::vm::mpk::ProtectionKey;
 use crate::runtime::vm::table::Table;
 use crate::runtime::vm::{CompiledModuleId, ModuleRuntimeInfo};
-use crate::store::{Asyncness, InstanceId, StoreOpaque, StoreResourceLimiter};
-use crate::vm::instance::PassiveElementSegment;
+use crate::store::{Asyncness, AutoAssertNoGc, InstanceId, StoreOpaque, StoreResourceLimiter};
 use crate::{OpaqueRootScope, Val};
 use core::future::Future;
 use core::pin::Pin;
@@ -860,34 +859,40 @@ async fn initialize_passive_elements(
     for (idx, segment) in &module.passive_elements {
         match segment {
             TableSegmentElements::Functions(func_indices) => {
-                let mut segment =
-                    PassiveElementSegment::new(NeedsGcRooting::No, func_indices.len())?;
+                let mut vals = TryVec::with_capacity(func_indices.len())?;
                 for func_idx in func_indices {
                     let (instance, registry) =
                         store.instance_and_module_registry_mut(context.instance);
                     // SAFETY: `store_id` is for the store that owns this instance.
                     let func = unsafe { instance.get_exported_func(registry, store_id, *func_idx) };
-                    segment.push(store, func.into())?;
+                    vals.push(func.to_val_raw(store))?;
                 }
                 let instance = store.instance_mut(context.instance);
                 debug_assert_eq!(instance.passive_elements.len(), idx.index());
-                instance.passive_elements_mut().push(segment)?;
+                instance
+                    .passive_elements_mut()
+                    .push(Some((NeedsGcRooting::No, vals)))?;
             }
             TableSegmentElements::Expressions {
                 needs_gc_rooting,
                 exprs,
             } => {
-                let mut segment = PassiveElementSegment::new(*needs_gc_rooting, exprs.len())?;
+                let mut vals = TryVec::with_capacity(exprs.len())?;
                 for expr in exprs {
                     let mut store = OpaqueRootScope::new(&mut *store);
+
                     let val = const_evaluator
                         .eval(&mut store, limiter.as_deref_mut(), context, expr)
                         .await?;
-                    segment.push(&mut store, *val)?;
+
+                    let mut store = AutoAssertNoGc::new(&mut store);
+                    vals.push(val.to_raw_(&mut store)?)?;
                 }
                 let instance = store.instance_mut(context.instance);
                 debug_assert_eq!(instance.passive_elements.len(), idx.index());
-                instance.passive_elements_mut().push(segment)?;
+                instance
+                    .passive_elements_mut()
+                    .push(Some((*needs_gc_rooting, vals)))?;
             }
         }
     }

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

@@ -528,8 +528,7 @@ fn table_init(
 // Implementation of `elem.drop`.
 fn elem_drop(store: &mut dyn VMStore, instance: InstanceId, elem_index: u32) -> Result<()> {
     let elem_index = ElemIndex::from_u32(elem_index);
-    let (gc_store, instance) = store.optional_gc_store_and_instance_mut(instance);
-    instance.elem_drop(gc_store, elem_index)?;
+    store.instance_mut(instance).elem_drop(elem_index)?;
     Ok(())
 }