Handle more failures deallocating pooled memroy (#12888)

* Handle more failures deallocating pooled memroy

This commit replaces a few panics in the pooling allocator with
error-handling of what happens at runtime. This is a defense-in-depth
measure to ensure that the pooling allocator doesn't panic at runtime
and instead handles errors where possible.

The first path fixed is in `deallocate_memory` where resetting a slot
could result in an error being returned on non-Linux platforms, and if
this happened it would cause a panic. The error is instead gracefully
handled by continuing slot deallocation but avoiding putting the image
itself back into memory. This leaves the slot in an `Unknown` state
which is already handled by resetting the state upon reuse. The main
consequence here is that future statistics about resident bytes won't be
accurate, but these are already inaccurate on non-Linux platforms
anyway, so there's no loss.

The second path fixes is in flushing a `DecommitQueue` where
`decommit_pages` was asserted to succeed. Instead now the error is
handled by dropping all images and leaving slots in an `Unknown` state,
similar to `deallocate_memory`.

* Review comments

Author: alexcrichton

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

@@ -721,33 +721,52 @@ unsafe impl InstanceAllocator for PoolingInstanceAllocator {
         let prev = self.live_memories.fetch_sub(1, Ordering::Relaxed);
         debug_assert!(prev > 0);
 
-        // Reset the image slot. If there is any error clearing the
-        // image, just drop it here, and let the drop handler for the
-        // slot unmap in a way that retains the address space
-        // reservation.
+        // Reset the image slot. Depending on whether this is successful or not
+        // the `image` is preserved for future use. On success it's queued up to
+        // get deallocated later, and on failure the slot is deallocated
+        // immediately without preserving the image.
         let mut image = memory.unwrap_static_image();
         let mut queue = DecommitQueue::default();
-        let bytes_resident = image
-            .clear_and_remain_ready(
-                self.pagemap.as_ref(),
-                self.memories.keep_resident,
-                |ptr, len| {
-                    // SAFETY: the memory in `image` won't be used until this
-                    // decommit queue is flushed, and by definition the memory is
-                    // not in use when calling this function.
-                    unsafe {
-                        queue.push_raw(ptr, len);
-                    }
-                },
-            )
-            .expect("failed to reset memory image");
+        let bytes_resident = image.clear_and_remain_ready(
+            self.pagemap.as_ref(),
+            self.memories.keep_resident,
+            |ptr, len| {
+                // SAFETY: the memory in `image` won't be used until this
+                // decommit queue is flushed, and by definition the memory is
+                // not in use when calling this function.
+                unsafe {
+                    queue.push_raw(ptr, len);
+                }
+            },
+        );
 
-        // SAFETY: this image is not in use and its memory regions were enqueued
-        // with `push_raw` above.
-        unsafe {
-            queue.push_memory(allocation_index, image, bytes_resident);
+        match bytes_resident {
+            Ok(bytes_resident) => {
+                // SAFETY: this image is not in use and its memory regions were enqueued
+                // with `push_raw` above.
+                unsafe {
+                    queue.push_memory(allocation_index, image, bytes_resident);
+                }
+                self.merge_or_flush(queue);
+            }
+            Err(e) => {
+                log::warn!("ignoring clear_and_remain_ready error {e}");
+                // SAFETY: `allocation_index` comes from this pool, as an unsafe
+                // contract of this function itself, and it's guaranteed to be no
+                // longer in use so safe to deallocate. The slot couldn't be
+                // preserved so it's dropped here.
+                //
+                // Note that at this point it's not clear how many bytes are
+                // resident in memory, so it's inevitably going to leave statistics
+                // a little off. Also note though that non-Linux platforms don't
+                // keep track of resident bytes anyway, and this path is only
+                // reachable on non-Linux platforms because Linux can't return an
+                // error.
+                unsafe {
+                    self.memories.deallocate(allocation_index, None, 0);
+                }
+            }
         }
-        self.merge_or_flush(queue);
     }
 
     fn allocate_table<'a, 'b: 'a, 'c: 'a>(

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

@@ -11,6 +11,7 @@
 use super::PoolingInstanceAllocator;
 use crate::vm::{MemoryAllocationIndex, MemoryImageSlot, Table, TableAllocationIndex};
 use smallvec::SmallVec;
+use std::io;
 
 #[cfg(feature = "async")]
 use wasmtime_fiber::FiberStack;
@@ -155,18 +156,14 @@ impl DecommitQueue {
         self.stacks.push((SendSyncStack(stack), bytes_resident));
     }
 
-    fn decommit_all_raw(&mut self) {
+    /// Returns if any decommit call failed.
+    fn decommit_all_raw(&mut self) -> io::Result<()> {
         for iovec in self.raw.drain(..) {
             unsafe {
-                crate::vm::sys::vm::decommit_pages(iovec.0.iov_base.cast(), iovec.0.iov_len)
-                    .unwrap_or_else(|e| {
-                        panic!(
-                            "failed to decommit ptr={:#p}, len={:#x}: {e}",
-                            iovec.0.iov_base, iovec.0.iov_len
-                        )
-                    });
+                crate::vm::sys::vm::decommit_pages(iovec.0.iov_base.cast(), iovec.0.iov_len)?;
             }
         }
+        Ok(())
     }
 
     /// Flush this queue, decommitting all enqueued regions in batch.
@@ -175,14 +172,26 @@ impl DecommitQueue {
     /// the associated free lists; `false` if the queue was empty.
     pub fn flush(mut self, pool: &PoolingInstanceAllocator) -> bool {
         // First, do the raw decommit syscall(s).
-        self.decommit_all_raw();
+        let decommit_succeeded = self.decommit_all_raw().is_ok();
 
         // Second, restore the various entities to their associated pools' free
         // lists. This is safe, and they are ready for reuse, now that their
         // memory regions have been decommitted.
+        //
+        // Note that for memory images the images are all dropped here and
+        // ignored if any decommits failed. This signifies how the state of the
+        // slot is unknown and needs to be paved over in the future. Also note
+        // that `bytes_resident` is probably too low, but there's no other
+        // precise way to know, so it's left here as-is and it'll get reset when
+        // the slot is reused.
         let mut deallocated_any = false;
         for (allocation_index, image, bytes_resident) in self.memories {
             deallocated_any = true;
+            let image = if decommit_succeeded {
+                Some(image)
+            } else {
+                None
+            };
             unsafe {
                 pool.memories
                     .deallocate(allocation_index, image, bytes_resident);

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

@@ -456,18 +456,23 @@ impl MemoryPool {
 
     /// Deallocate a previously-allocated memory.
     ///
+    /// If `image` is `None` then the state of this memory's slot is left
+    /// unknown. Otherwise `image` is used to retain information about the state
+    /// of this slot.
+    ///
     /// # Safety
     ///
     /// The memory must have been previously allocated from this pool and
     /// assigned the given index, must currently be in an allocated state, and
     /// must never be used again.
     ///
     /// The caller must have already called `clear_and_remain_ready` on the
-    /// memory's image and flushed any enqueued decommits for this memory.
+    /// memory's image and flushed any enqueued decommits for this memory. Note
+    /// that if `image` is `None` then this is not required.
     pub unsafe fn deallocate(
         &self,
         allocation_index: MemoryAllocationIndex,
-        image: MemoryImageSlot,
+        image: Option<MemoryImageSlot>,
         bytes_resident: usize,
     ) {
         self.return_memory_image_slot(allocation_index, image);
@@ -520,7 +525,7 @@ impl MemoryPool {
                     let index = MemoryAllocationIndex(id.0);
                     if let Ok(mut slot) = self.take_memory_image_slot(index) {
                         if slot.remove_image().is_ok() {
-                            self.return_memory_image_slot(index, slot);
+                            self.return_memory_image_slot(index, Some(slot));
                         }
                     }
 
@@ -587,16 +592,23 @@ impl MemoryPool {
     }
 
     /// Return ownership of the given image slot.
+    ///
+    /// If `slot` is not provided then it's reset with `Unknown` meaning a
+    /// future allocation will need to pave over it to use it.
     fn return_memory_image_slot(
         &self,
         allocation_index: MemoryAllocationIndex,
-        slot: MemoryImageSlot,
+        slot: Option<MemoryImageSlot>,
     ) {
-        assert!(!slot.is_dirty());
-
         let prev = mem::replace(
             &mut *self.image_slots[allocation_index.index()].lock().unwrap(),
-            ImageSlot::PreviouslyUsed(slot),
+            match slot {
+                Some(slot) => {
+                    assert!(!slot.is_dirty());
+                    ImageSlot::PreviouslyUsed(slot)
+                }
+                None => ImageSlot::Unknown,
+            },
         );
         assert!(matches!(prev, ImageSlot::Unknown));
     }