Implement user fault handling with userfaultfd on Linux.

This commit implements the `uffd` feature which turns on support for utilizing
the `userfaultfd` system call on Linux for the pooling instance allocator.

By handling page faults in userland, we are able to detect guard page accesses
without having to constantly change memory page protections.

This should help reduce the number of syscalls as well as kernel lock
contentions when many threads are allocating and deallocating instances.

Additionally, the user fault handler can lazy initialize table and linear
memories of an instance (implementation to come).

Author: peterhuene

.github/workflows/main.yml

@@ -122,6 +122,7 @@ jobs:
     - run: cargo check --manifest-path crates/wasmtime/Cargo.toml --features jitdump
     - run: cargo check --manifest-path crates/wasmtime/Cargo.toml --features cache
     - run: cargo check --manifest-path crates/wasmtime/Cargo.toml --features async
+    - run: cargo check --manifest-path crates/wasmtime/Cargo.toml --features uffd
 
     # Check some feature combinations of the `wasmtime-c-api` crate
     - run: cargo check --manifest-path crates/c-api/Cargo.toml --no-default-features

Cargo.toml

@@ -89,6 +89,7 @@ jitdump = ["wasmtime/jitdump"]
 vtune = ["wasmtime/vtune"]
 wasi-crypto = ["wasmtime-wasi-crypto"]
 wasi-nn = ["wasmtime-wasi-nn"]
+uffd = ["wasmtime/uffd"]
 
 # Try the experimental, work-in-progress new x86_64 backend. This is not stable
 # as of June 2020.

crates/runtime/Cargo.toml

@@ -37,3 +37,9 @@ cc = "1.0"
 
 [badges]
 maintenance = { status = "actively-developed" }
+
+[features]
+default = []
+
+# Enables support for userfaultfd in the pooling allocator when building on Linux
+uffd = ["userfaultfd"]

crates/runtime/src/instance/allocator/pooling.rs

@@ -31,6 +31,11 @@ cfg_if::cfg_if! {
     if #[cfg(windows)] {
         mod windows;
         use windows as imp;
+    } else if #[cfg(all(feature = "uffd", target_os = "linux"))] {
+        mod uffd;
+        use uffd as imp;
+        use imp::{PageFaultHandler, reset_guard_page};
+        use std::sync::atomic::{AtomicBool, Ordering};
     } else if #[cfg(target_os = "linux")] {
         mod linux;
         use linux as imp;
@@ -335,6 +340,9 @@ impl Iterator for BasePointerIterator {
 /// structure depending on the limits used to create the pool.
 ///
 /// The pool maintains a free list for fast instance allocation.
+///
+/// The userfault handler relies on how instances are stored in the mapping,
+/// so make sure the uffd implementation is kept up-to-date.
 #[derive(Debug)]
 struct InstancePool {
     mapping: Mmap,
@@ -472,6 +480,10 @@ impl Drop for InstancePool {
 ///
 /// Each index into the pool returns an iterator over the base addresses
 /// of the instance's linear memories.
+///
+///
+/// The userfault handler relies on how memories are stored in the mapping,
+/// so make sure the uffd implementation is kept up-to-date.
 #[derive(Debug)]
 struct MemoryPool {
     mapping: Mmap,
@@ -524,6 +536,9 @@ impl MemoryPool {
 ///
 /// Each index into the pool returns an iterator over the base addresses
 /// of the instance's tables.
+///
+/// The userfault handler relies on how tables are stored in the mapping,
+/// so make sure the uffd implementation is kept up-to-date.
 #[derive(Debug)]
 struct TablePool {
     mapping: Mmap,
@@ -588,13 +603,18 @@ impl TablePool {
 ///
 /// The top of the stack (starting stack pointer) is returned when a stack is allocated
 /// from the pool.
+///
+/// The userfault handler relies on how stacks are stored in the mapping,
+/// so make sure the uffd implementation is kept up-to-date.
 #[derive(Debug)]
 struct StackPool {
     mapping: Mmap,
     stack_size: usize,
     max_instances: usize,
     page_size: usize,
     free_list: Mutex<Vec<usize>>,
+    #[cfg(all(feature = "uffd", target_os = "linux"))]
+    faulted_guard_pages: Arc<[AtomicBool]>,
 }
 
 impl StackPool {
@@ -623,6 +643,11 @@ impl StackPool {
             max_instances,
             page_size,
             free_list: Mutex::new((0..max_instances).collect()),
+            #[cfg(all(feature = "uffd", target_os = "linux"))]
+            faulted_guard_pages: std::iter::repeat_with(|| false.into())
+                .take(max_instances)
+                .collect::<Vec<_>>()
+                .into(),
         })
     }
 
@@ -647,11 +672,25 @@ impl StackPool {
                 .as_mut_ptr()
                 .add((index * self.stack_size) + self.page_size);
 
-            // Make the stack accessible (excluding the guard page)
-            if !make_accessible(bottom_of_stack, size_without_guard) {
-                return Err(FiberStackError::Resource(
-                    "failed to make instance memory accessible".into(),
-                ));
+            cfg_if::cfg_if! {
+                if #[cfg(all(feature = "uffd", target_os = "linux"))] {
+                    // Check to see if a guard page needs to be reset
+                    if self.faulted_guard_pages[index].swap(false, Ordering::SeqCst) {
+                        if !reset_guard_page(bottom_of_stack.sub(self.page_size), self.page_size) {
+                            return Err(FiberStackError::Resource(
+                                "failed to reset stack guard page".into(),
+                            ));
+                        }
+                    }
+
+                } else {
+                    // Make the stack accessible (excluding the guard page)
+                    if !make_accessible(bottom_of_stack, size_without_guard) {
+                        return Err(FiberStackError::Resource(
+                            "failed to make instance memory accessible".into(),
+                        ));
+                    }
+                }
             }
 
             // The top of the stack should be returned
@@ -697,6 +736,8 @@ pub struct PoolingInstanceAllocator {
     memories: mem::ManuallyDrop<MemoryPool>,
     tables: mem::ManuallyDrop<TablePool>,
     stacks: mem::ManuallyDrop<StackPool>,
+    #[cfg(all(feature = "uffd", target_os = "linux"))]
+    _fault_handler: PageFaultHandler,
 }
 
 impl PoolingInstanceAllocator {
@@ -744,6 +785,9 @@ impl PoolingInstanceAllocator {
         let tables = TablePool::new(&module_limits, &instance_limits)?;
         let stacks = StackPool::new(&instance_limits, stack_size)?;
 
+        #[cfg(all(feature = "uffd", target_os = "linux"))]
+        let _fault_handler = PageFaultHandler::new(&instances, &memories, &tables, &stacks)?;
+
         Ok(Self {
             strategy,
             module_limits,
@@ -752,6 +796,8 @@ impl PoolingInstanceAllocator {
             memories: mem::ManuallyDrop::new(memories),
             tables: mem::ManuallyDrop::new(tables),
             stacks: mem::ManuallyDrop::new(stacks),
+            #[cfg(all(feature = "uffd", target_os = "linux"))]
+            _fault_handler,
         })
     }
 
@@ -800,14 +846,28 @@ impl PoolingInstanceAllocator {
     ) -> Result<(), InstantiationError> {
         let module = instance.module.as_ref();
 
+        // Reset all guard pages before clearing the previous memories
+        #[cfg(all(feature = "uffd", target_os = "linux"))]
+        for (_, m) in instance.memories.iter() {
+            m.reset_guard_pages()
+                .map_err(InstantiationError::Resource)?;
+        }
+
         instance.memories.clear();
 
         for plan in
             (&module.memory_plans.values().as_slice()[module.num_imported_memories..]).iter()
         {
             instance.memories.push(
-                Memory::new_static(plan, memories.next().unwrap(), max_pages, make_accessible)
-                    .map_err(InstantiationError::Resource)?,
+                Memory::new_static(
+                    plan,
+                    memories.next().unwrap(),
+                    max_pages,
+                    make_accessible,
+                    #[cfg(all(feature = "uffd", target_os = "linux"))]
+                    reset_guard_page,
+                )
+                .map_err(InstantiationError::Resource)?,
             );
         }
 
@@ -826,7 +886,6 @@ impl PoolingInstanceAllocator {
         let module = instance.module.as_ref();
 
         instance.tables.clear();
-
         for plan in (&module.table_plans.values().as_slice()[module.num_imported_tables..]).iter() {
             let base = tables.next().unwrap();
 
@@ -852,7 +911,8 @@ impl PoolingInstanceAllocator {
 
 impl Drop for PoolingInstanceAllocator {
     fn drop(&mut self) {
-        // There are manually dropped for the future uffd implementation
+        // Manually drop the pools before the fault handler (if uffd is enabled)
+        // This ensures that any fault handler thread monitoring the pool memory terminates
         unsafe {
             mem::ManuallyDrop::drop(&mut self.instances);
             mem::ManuallyDrop::drop(&mut self.memories);