feat: new deferred_free API

.github/workflows/rust.yml

@@ -13,19 +13,25 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        features:
+        mmap:
           - ""
-          - "--features mmap"
-          - "--features std_mutex"
-          - "--features spin_mutex"
-          - "--features mmap,std_mutex"
-          - "--features mmap,spin_mutex"
+          - "-F mmap"
+        mutex:
+          - ""
+          - "-F std_mutex"
+          - "-F spin_mutex"
+        deferred_free:
+          - ""
+          - "-F deferred_free"
     steps:
     - uses: actions/checkout@v4
     - uses: actions-rust-lang/setup-rust-toolchain@v1
+    - name: Calculate features
+      id: features
+      run: echo "features=${{ matrix.mmap }} ${{ matrix.mutex }} ${{ matrix.deferred_free }}" >> "$GITHUB_OUTPUT"
     - name: Clippy
-      run: cargo clippy ${{ matrix.features }}
+      run: cargo clippy ${{ steps.features.outputs.features }}
     - name: Build
-      run: cargo build ${{ matrix.features }}
+      run: cargo build ${{ steps.features.outputs.features }}
     - name: Run tests
-      run: cargo test ${{ matrix.features }} --release -- --test-threads 1
+      run: cargo test ${{ steps.features.outputs.features }} --release -- --test-threads 1

Cargo.toml

@@ -22,11 +22,16 @@ std = []
 mmap = ["dep:libc"]
 std_mutex = ["std"]
 spin_mutex = ["dep:spin"]
+deferred_free = []
 
 [[test]]
 name = "global_alloc"
 required-features = ["mmap", "spin_mutex"]
 
+[[test]]
+name = "deferred_free"
+required-features = ["deferred_free"]
+
 [package.metadata.docs.rs]
 all-features = true
 rustdoc-args = ["--cfg", "docsrs"]

README.md

@@ -14,6 +14,7 @@ It can be used in `no_std` environments.
 - **mmap** - Provide `MimallocMmap` that uses `mmap` as OS allocator for segments.
 - **std_mutex** - Provide `MimallocMutexWrapper` that wraps `Mimalloc` inside `std::sync::Mutex` and implements `GlobalAlloc`.
 - **spin_mutex** - Provide `MimallocMutexWrapper` that wraps `Mimalloc` inside `spin::Mutex` that can be used in `no_std` environments.
+- **deferred_free** - Enable registering a hook to complete deferred free events. See the documentation of [`mi_register_deferred_free`](https://microsoft.github.io/mimalloc/group__extended.html#ga3460a6ca91af97be4058f523d3cb8ece).
 
 ## Usage
 

src/deferred_free.rs

@@ -0,0 +1,22 @@
+use crate::*;
+
+/// Handle to complete deferred free in [`DeferredFreeHook`].
+pub struct DeferredFreeHandle<'a, A: GlobalAlloc> {
+    pub(crate) heap: &'a mut Heap,
+    pub(crate) os_alloc: &'a A,
+}
+
+impl<A: GlobalAlloc> DeferredFreeHandle<'_, A> {
+    /// Deallocate the block of memory at the given `ptr`.
+    ///
+    /// # Safety
+    ///
+    /// See [`GlobalAlloc::dealloc`].
+    pub unsafe fn free(&mut self, ptr: *mut u8) {
+        self.heap.free(ptr, self.os_alloc);
+    }
+}
+
+/// Hook to complete deferred free when the allocator needs more memory.
+/// See [`DeferredFreeHandle`] and [`Mimalloc::register_deferred_free`].
+pub type DeferredFreeHook<A> = fn(handle: &mut DeferredFreeHandle<A>, force: bool, heartbeat: u64);

src/heap.rs

@@ -5,15 +5,18 @@ use crate::segment::{PageKind, Segment};
 use crate::utils::{
     bin_for_size, wsize_from_size, BLOCK_SIZE_FOR_BIN, WSIZE_RANGE_IN_SAME_SMALL_BIN,
 };
+#[cfg(feature = "deferred_free")]
+use crate::{DeferredFreeHandle, DeferredFreeHook};
 use core::alloc::GlobalAlloc;
 use core::ptr::{null_mut, NonNull};
 
 pub struct Heap {
     pages_free_direct: [NonNull<Page>; MI_SMALL_WSIZE_MAX + 1],
     pages: [LinkedList<Page>; MI_BIN_HUGE + 1],
     small_free_segments: LinkedList<Segment>,
-    deferred_free_hook: Option<fn(bool, u64)>,
+    #[cfg(feature = "deferred_free")]
     heartbeat: u64,
+    #[cfg(feature = "deferred_free")]
     calling_deferred_free: bool,
 }
 
@@ -29,18 +32,37 @@ impl Heap {
             pages_free_direct: [empty_page(); MI_SMALL_WSIZE_MAX + 1],
             pages: [const { LinkedList::new() }; MI_BIN_HUGE + 1],
             small_free_segments: LinkedList::new(),
-            deferred_free_hook: None,
+            #[cfg(feature = "deferred_free")]
             heartbeat: 0,
+            #[cfg(feature = "deferred_free")]
             calling_deferred_free: false,
         }
     }
 
-    pub fn malloc<A: GlobalAlloc>(&mut self, size: usize, os_alloc: &A) -> *mut u8 {
+    #[inline(never)]
+    pub fn malloc<A: GlobalAlloc>(
+        &mut self,
+        size: usize,
+        os_alloc: &A,
+        #[cfg(feature = "deferred_free")] deferred_free_hook: Option<DeferredFreeHook<A>>,
+    ) -> *mut u8 {
         let result = if size <= MI_SMALL_SIZE_MAX {
             let page = self.get_small_free_page(size);
-            Page::malloc_fast(page, self, size, os_alloc)
+            Page::malloc_fast(
+                page,
+                self,
+                size,
+                os_alloc,
+                #[cfg(feature = "deferred_free")]
+                deferred_free_hook,
+            )
         } else {
-            self.malloc_generic(size, os_alloc)
+            self.malloc_generic(
+                size,
+                os_alloc,
+                #[cfg(feature = "deferred_free")]
+                deferred_free_hook,
+            )
         };
         debug_assert!(
             match &result {
@@ -64,11 +86,17 @@ impl Heap {
         size: usize,
         align: usize,
         os_alloc: &A,
+        #[cfg(feature = "deferred_free")] deferred_free_hook: Option<DeferredFreeHook<A>>,
     ) -> *mut u8 {
         debug_assert!(align.is_power_of_two());
 
         if align <= MI_INTPTR_SIZE {
-            return self.malloc(size, os_alloc);
+            return self.malloc(
+                size,
+                os_alloc,
+                #[cfg(feature = "deferred_free")]
+                deferred_free_hook,
+            );
         }
         if size >= usize::MAX - align {
             return null_mut();
@@ -78,17 +106,29 @@ impl Heap {
             let page = self.get_small_free_page(size);
             let free = unsafe { page.as_ref() }.free();
             if !free.is_null() && (free as usize & (align - 1) == 0) {
-                return Page::malloc_fast(page, self, size, os_alloc)
-                    .map_or(null_mut(), |(ptr, _)| ptr.as_ptr());
+                return Page::malloc_fast(
+                    page,
+                    self,
+                    size,
+                    os_alloc,
+                    #[cfg(feature = "deferred_free")]
+                    deferred_free_hook,
+                )
+                .map_or(null_mut(), |(ptr, _)| ptr.as_ptr());
             }
         }
 
-        self.malloc_generic(size + align - 1, os_alloc)
-            .map_or(null_mut(), |(ptr, page)| {
-                page.set_aligned(true);
-                let aligned_addr = (ptr.as_ptr() as usize + align - 1) & !(align - 1);
-                aligned_addr as *mut u8
-            })
+        self.malloc_generic(
+            size + align - 1,
+            os_alloc,
+            #[cfg(feature = "deferred_free")]
+            deferred_free_hook,
+        )
+        .map_or(null_mut(), |(ptr, page)| {
+            page.set_aligned(true);
+            let aligned_addr = (ptr.as_ptr() as usize + align - 1) & !(align - 1);
+            aligned_addr as *mut u8
+        })
     }
 
     pub fn free<A: GlobalAlloc>(&mut self, p: *mut u8, os_alloc: &A) {
@@ -108,28 +148,46 @@ impl Heap {
         &mut self,
         size: usize,
         os_alloc: &A,
+        #[cfg(feature = "deferred_free")] deferred_free_hook: Option<DeferredFreeHook<A>>,
     ) -> Option<(NonNull<u8>, &mut Page)> {
-        self.deferred_free(false);
+        #[cfg(feature = "deferred_free")]
+        self.deferred_free(false, os_alloc, deferred_free_hook);
 
         let page = if size <= MI_LARGE_SIZE_MAX {
             self.find_free_page(size, os_alloc)
         } else {
             self.alloc_huge_page(size, os_alloc)
         };
 
-        NonNull::new(page).and_then(|p| Page::malloc_fast(p, self, size, os_alloc))
+        NonNull::new(page).and_then(|p| {
+            Page::malloc_fast(
+                p,
+                self,
+                size,
+                os_alloc,
+                #[cfg(feature = "deferred_free")]
+                deferred_free_hook,
+            )
+        })
     }
 
-    pub const fn register_deferred_free(&mut self, hook: fn(bool, u64)) {
-        self.deferred_free_hook = Some(hook);
-    }
-
-    fn deferred_free(&mut self, force: bool) {
+    #[cfg(feature = "deferred_free")]
+    fn deferred_free<A: GlobalAlloc>(
+        &mut self,
+        force: bool,
+        os_alloc: &A,
+        deferred_free_hook: Option<DeferredFreeHook<A>>,
+    ) {
         self.heartbeat = self.heartbeat.wrapping_add(1);
-        if let Some(hook) = self.deferred_free_hook {
+        if let Some(hook) = deferred_free_hook {
             if !self.calling_deferred_free {
                 self.calling_deferred_free = true;
-                hook(force, self.heartbeat);
+                let heartbeat = self.heartbeat;
+                let mut handle = DeferredFreeHandle {
+                    heap: self,
+                    os_alloc,
+                };
+                hook(&mut handle, force, heartbeat);
                 self.calling_deferred_free = false;
             }
         }

src/lib.rs

@@ -29,6 +29,8 @@
 //!   [`std::sync::Mutex`] and implements [`GlobalAlloc`].
 //! - **spin_mutex** - Provide [`MimallocMutexWrapper`] that wraps [`Mimalloc`] inside
 //!   [`spin::Mutex`] that can be used in `no_std` environments.
+//! - **deferred_free** - Enable registering a hook to complete deferred free events.
+//!   See the documentation of [`mi_register_deferred_free`](https://microsoft.github.io/mimalloc/group__extended.html#ga3460a6ca91af97be4058f523d3cb8ece).
 
 #![cfg_attr(docsrs, feature(doc_auto_cfg))]
 #![cfg_attr(not(feature = "std"), no_std)]
@@ -55,8 +57,15 @@ use heap::Heap;
 pub struct Mimalloc<A: GlobalAlloc> {
     heap: Heap,
     os_alloc: A,
+    #[cfg(feature = "deferred_free")]
+    deferred_free_hook: Option<DeferredFreeHook<A>>,
 }
 
+#[cfg(feature = "deferred_free")]
+pub mod deferred_free;
+#[cfg(feature = "deferred_free")]
+use deferred_free::*;
+
 unsafe impl<A: GlobalAlloc> Send for Mimalloc<A> {}
 
 impl<A: GlobalAlloc> Mimalloc<A> {
@@ -65,17 +74,20 @@ impl<A: GlobalAlloc> Mimalloc<A> {
         Self {
             heap: Heap::new(),
             os_alloc,
+            #[cfg(feature = "deferred_free")]
+            deferred_free_hook: None,
         }
     }
 
+    #[cfg(feature = "deferred_free")]
     /// Register a hook to complete deferred free when the allocator needs more memory.
     /// A new hook replaces the old one.
     ///
     /// See the documentation of
     /// [`mi_register_deferred_free`](https://microsoft.github.io/mimalloc/group__extended.html#ga3460a6ca91af97be4058f523d3cb8ece)
     /// (the extra `arg` is not supported).
-    pub const fn register_deferred_free(&mut self, hook: fn(force: bool, heartbeat: u64)) {
-        self.heap.register_deferred_free(hook)
+    pub const fn register_deferred_free(&mut self, hook: DeferredFreeHook<A>) {
+        self.deferred_free_hook = Some(hook);
     }
 
     /// Collect free memory.
@@ -89,8 +101,13 @@ impl<A: GlobalAlloc> Mimalloc<A> {
     ///
     /// See [`GlobalAlloc::alloc`].
     pub unsafe fn alloc(&mut self, layout: Layout) -> *mut u8 {
-        self.heap
-            .malloc_aligned(layout.size(), layout.align(), &self.os_alloc)
+        self.heap.malloc_aligned(
+            layout.size(),
+            layout.align(),
+            &self.os_alloc,
+            #[cfg(feature = "deferred_free")]
+            self.deferred_free_hook,
+        )
     }
 
     /// [`GlobalAlloc::dealloc`] but requires a mutable reference `&mut self`.

src/mutex.rs

@@ -16,8 +16,9 @@ impl<A: GlobalAlloc> MimallocMutexWrapper<A> {
         Self(Mutex::new(Mimalloc::with_os_allocator(os_alloc)))
     }
 
+    #[cfg(feature = "deferred_free")]
     /// See [`Mimalloc::register_deferred_free`].
-    pub fn register_deferred_free(&self, hook: fn(force: bool, heartbeat: u64)) {
+    pub fn register_deferred_free(&self, hook: crate::DeferredFreeHook<A>) {
         self.allocator().register_deferred_free(hook);
     }
 

src/page.rs

@@ -5,6 +5,8 @@ use crate::heap::Heap;
 use crate::list::impl_list_item;
 use crate::segment::Segment;
 use crate::utils::bin_for_size;
+#[cfg(feature = "deferred_free")]
+use crate::DeferredFreeHook;
 use core::alloc::GlobalAlloc;
 use core::ptr::{null_mut, NonNull};
 
@@ -46,6 +48,7 @@ impl Page {
         heap: &'a mut Heap,
         size: usize,
         os_alloc: &A,
+        #[cfg(feature = "deferred_free")] deferred_free_hook: Option<DeferredFreeHook<A>>,
     ) -> Option<(NonNull<u8>, &'a mut Page)> {
         debug_assert!(
             page.as_ptr() == empty_page().as_ptr()
@@ -54,7 +57,12 @@ impl Page {
             { unsafe { page.as_ref() }.block_size }
         );
         match unsafe { page.as_ref().free.as_mut() } {
-            None => heap.malloc_generic(size, os_alloc),
+            None => heap.malloc_generic(
+                size,
+                os_alloc,
+                #[cfg(feature = "deferred_free")]
+                deferred_free_hook,
+            ),
             Some(block) => {
                 debug_assert!(
                     (block as *const _ as usize).wrapping_sub(page.as_ptr() as usize)