Move linear memory faulted guard page tracking into Memory.

This commit moves the tracking for faulted guard pages in a linear memory into
`Memory`.

crates/runtime/src/instance.rs

@@ -69,11 +69,6 @@ pub(crate) struct Instance {
     /// Hosts can store arbitrary per-instance information here.
     host_state: Box<dyn Any>,
 
-    /// Stores linear memory guard page faults for the pooling allocator with uffd enabled.
-    /// These pages need to be reset after the signal handler generates the out-of-bounds trap.
-    #[cfg(all(feature = "uffd", target_os = "linux"))]
-    guard_page_faults: RefCell<Vec<(*mut u8, usize, fn(*mut u8, usize) -> anyhow::Result<()>)>>,
-
     /// Additional context used by compiled wasm code. This field is last, and
     /// represents a dynamically-sized array that extends beyond the nominal
     /// end of the struct (similar to a flexible array member).
@@ -383,14 +378,6 @@ impl Instance {
     /// Returns `None` if memory can't be grown by the specified amount
     /// of pages.
     pub(crate) fn memory_grow(&self, memory_index: DefinedMemoryIndex, delta: u32) -> Option<u32> {
-        // Reset all guard pages before growing any memory when using the uffd feature.
-        // The uffd feature induces a trap when a fault on a linear memory page is determined to be out-of-bounds.
-        // It does this by temporarily setting the protection level to `NONE` to cause the kernel to signal SIGBUS.
-        // Because instances might still be used after a trap, this resets the page back to the expected protection
-        // level (READ_WRITE) for the uffd implementation.
-        #[cfg(all(feature = "uffd", target_os = "linux"))]
-        self.reset_guard_pages().ok()?;
-
         let result = self
             .memories
             .get(memory_index)
@@ -822,36 +809,6 @@ impl Instance {
             (foreign_table_index, foreign_instance)
         }
     }
-
-    /// Records a faulted guard page.
-    ///
-    /// This is used to track faulted guard pages that need to be reset.
-    #[cfg(all(feature = "uffd", target_os = "linux"))]
-    pub(crate) fn record_guard_page_fault(
-        &self,
-        page_addr: *mut u8,
-        size: usize,
-        reset: fn(*mut u8, usize) -> anyhow::Result<()>,
-    ) {
-        self.guard_page_faults
-            .borrow_mut()
-            .push((page_addr, size, reset));
-    }
-
-    /// Resets previously faulted guard pages.
-    ///
-    /// This is used to reset the protection of any guard pages that were previously faulted.
-    ///
-    /// Resetting the guard pages is required before growing memory.
-    #[cfg(all(feature = "uffd", target_os = "linux"))]
-    pub(crate) fn reset_guard_pages(&self) -> anyhow::Result<()> {
-        let mut faults = self.guard_page_faults.borrow_mut();
-        for (addr, len, reset) in faults.drain(..) {
-            reset(addr, len)?;
-        }
-
-        Ok(())
-    }
 }
 
 /// A handle holding an `Instance` of a WebAssembly module.

crates/runtime/src/instance/allocator.rs

@@ -602,8 +602,6 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
                 )),
                 dropped_data: RefCell::new(EntitySet::with_capacity(req.module.passive_data.len())),
                 host_state,
-                #[cfg(all(feature = "uffd", target_os = "linux"))]
-                guard_page_faults: RefCell::new(Vec::new()),
                 vmctx: VMContext {},
             };
             let layout = instance.alloc_layout();

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

@@ -367,8 +367,6 @@ impl InstancePool {
                     dropped_elements: RefCell::new(EntitySet::new()),
                     dropped_data: RefCell::new(EntitySet::new()),
                     host_state: Box::new(()),
-                    #[cfg(all(feature = "uffd", target_os = "linux"))]
-                    guard_page_faults: RefCell::new(Vec::new()),
                     vmctx: VMContext {},
                 },
             );
@@ -431,9 +429,15 @@ impl InstancePool {
             let memory = mem::take(memory);
             debug_assert!(memory.is_static());
 
+            // Reset any faulted guard pages as the physical memory may be reused for another instance in the future
+            #[cfg(all(feature = "uffd", target_os = "linux"))]
+            memory
+                .reset_guard_pages()
+                .expect("failed to reset guard pages");
+
             let size = (memory.size() * WASM_PAGE_SIZE) as usize;
             drop(memory);
-            decommit_memory_pages(base, size).unwrap();
+            decommit_memory_pages(base, size).expect("failed to decommit linear memory pages");
         }
 
         instance.memories.clear();
@@ -450,7 +454,7 @@ impl InstancePool {
             );
 
             drop(table);
-            decommit_table_pages(base, size).unwrap();
+            decommit_table_pages(base, size).expect("failed to decommit table pages");
         }
 
         instance.tables.clear();
@@ -469,12 +473,6 @@ impl InstancePool {
     ) -> Result<(), InstantiationError> {
         let module = instance.module.as_ref();
 
-        // Reset all guard pages before reusing the instance
-        #[cfg(all(feature = "uffd", target_os = "linux"))]
-        instance
-            .reset_guard_pages()
-            .map_err(|e| InstantiationError::Resource(e.to_string()))?;
-
         debug_assert!(instance.memories.is_empty());
 
         for plan in

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

@@ -313,8 +313,12 @@ unsafe fn handle_page_fault(
                 None => {
                     log::trace!("out of bounds memory access at {:p}", addr);
 
-                    // Record the guard page fault with the instance so it can be reset later.
+                    // Record the guard page fault so the page protection level can be reset later
-                    instance.record_guard_page_fault(page_addr, len, reset_guard_page);
+                    instance.memories[memory_index].record_guard_page_fault(
+                        page_addr,
+                        len,
+                        reset_guard_page,
+                    );
                     wake_guard_page_access(&uffd, page_addr, len)?;
                 }
             }

crates/runtime/src/memory.rs

@@ -180,6 +180,10 @@ enum MemoryStorage {
         size: Cell<u32>,
         maximum: u32,
         make_accessible: fn(*mut u8, usize) -> Result<()>,
+        /// Stores the pages in the linear memory that have faulted as guard pages when using the `uffd` feature.
+        /// These pages need their protection level reset before the memory can grow.
+        #[cfg(all(feature = "uffd", target_os = "linux"))]
+        guard_page_faults: RefCell<Vec<(*mut u8, usize, fn(*mut u8, usize) -> Result<()>)>>,
     },
     Dynamic(Box<dyn RuntimeLinearMemory>),
 }
@@ -209,6 +213,8 @@ impl Memory {
             size: Cell::new(plan.memory.minimum),
             maximum: min(plan.memory.maximum.unwrap_or(maximum), maximum),
             make_accessible,
+            #[cfg(all(feature = "uffd", target_os = "linux"))]
+            guard_page_faults: RefCell::new(Vec::new()),
         }))
     }
 
@@ -242,6 +248,10 @@ impl Memory {
                 make_accessible,
                 ..
             } => {
+                // Reset any faulted guard pages before growing the memory.
+                #[cfg(all(feature = "uffd", target_os = "linux"))]
+                self.reset_guard_pages().ok()?;
+
                 let old_size = size.get();
                 if delta == 0 {
                     return Some(old_size);
@@ -276,6 +286,56 @@ impl Memory {
             MemoryStorage::Dynamic(mem) => mem.vmmemory(),
         }
     }
+
+    /// Records a faulted guard page in a static memory.
+    ///
+    /// This is used to track faulted guard pages that need to be reset for the uffd feature.
+    ///
+    /// This function will panic if called on a dynamic memory.
+    #[cfg(all(feature = "uffd", target_os = "linux"))]
+    pub(crate) fn record_guard_page_fault(
+        &self,
+        page_addr: *mut u8,
+        size: usize,
+        reset: fn(*mut u8, usize) -> Result<()>,
+    ) {
+        match &self.0 {
+            MemoryStorage::Static {
+                guard_page_faults, ..
+            } => {
+                guard_page_faults
+                    .borrow_mut()
+                    .push((page_addr, size, reset));
+            }
+            MemoryStorage::Dynamic(_) => {
+                unreachable!("dynamic memories should not have guard page faults")
+            }
+        }
+    }
+
+    /// Resets the previously faulted guard pages of a static memory.
+    ///
+    /// This is used to reset the protection of any guard pages that were previously faulted.
+    ///
+    /// This function will panic if called on a dynamic memory.
+    #[cfg(all(feature = "uffd", target_os = "linux"))]
+    pub(crate) fn reset_guard_pages(&self) -> Result<()> {
+        match &self.0 {
+            MemoryStorage::Static {
+                guard_page_faults, ..
+            } => {
+                let mut faults = guard_page_faults.borrow_mut();
+                for (addr, len, reset) in faults.drain(..) {
+                    reset(addr, len)?;
+                }
+            }
+            MemoryStorage::Dynamic(_) => {
+                unreachable!("dynamic memories should not have guard page faults")
+            }
+        }
+
+        Ok(())
+    }
 }
 
 // The default memory representation is an empty memory that cannot grow.
@@ -290,6 +350,8 @@ impl Default for Memory {
             size: Cell::new(0),
             maximum: 0,
             make_accessible,
+            #[cfg(all(feature = "uffd", target_os = "linux"))]
+            guard_page_faults: RefCell::new(Vec::new()),
         })
     }
 }