pyembed: track allocations with snmalloc

This addresses some outstanding TODOs and allows us to enable this
allocator! I haven't tested this thouroughly yet. But at least the
test doesn't crash!

Part of #358.

docs/history.rst

@@ -129,9 +129,10 @@ New Features
   the *mem* and *obj* domains as well as ``pymalloc``'s arena allocator.
   (However, the ``pymalloc`` arena allocator customization does not yet work
   due to a bug outside of our control.)
-* The *mimalloc* memory allocator can now be used as Python's *raw allocator*.
-  See documentation for
-  :ref:`PythonInterpreterConfig.raw_allocator <config_type_python_interpreter_config_allocator_backend>`.
+* The *mimalloc* and *snmalloc* memory allocators can now be used as Python's
+  memory allocators. See documentation for
+  :ref:`PythonInterpreterConfig.allocator_backend <config_type_python_interpreter_config_allocator_backend>`.
+  Code contributed by Ryan Clanton in #358.
 
 Other Relevant Changes
 ^^^^^^^^^^^^^^^^^^^^^^

pyembed/src/pyalloc.rs

@@ -139,6 +139,10 @@ impl AllocationTracker {
 
     /// Construct an instance from a pointer owned by someone else.
     fn from_owned_ptr(ptr: *mut c_void) -> BorrowedAllocationTracker {
+        if ptr.is_null() {
+            panic!("must not pass NULL pointer");
+        }
+
         BorrowedAllocationTracker {
             inner: Some(unsafe { Box::from_raw(ptr as *mut AllocationTracker) }),
         }
@@ -241,13 +245,20 @@ extern "C" fn mimalloc_malloc(_ctx: *mut c_void, size: size_t) -> *mut c_void {
 }
 
 #[cfg(feature = "snmalloc-sys")]
-extern "C" fn snmalloc_malloc(_ctx: *mut c_void, size: size_t) -> *mut c_void {
+extern "C" fn snmalloc_malloc(ctx: *mut c_void, size: size_t) -> *mut c_void {
     let size = match size {
         0 => 1,
         val => val,
     };
 
-    unsafe { snmalloc_sys::rust_alloc(SNMALLOC_ALIGNMENT, size) as *mut _ }
+    let mut tracker = AllocationTracker::from_owned_ptr(ctx);
+
+    let layout = unsafe { alloc::Layout::from_size_align_unchecked(size, SNMALLOC_ALIGNMENT) };
+    let res = unsafe { snmalloc_sys::rust_alloc(SNMALLOC_ALIGNMENT, size) as *mut _ };
+
+    tracker.insert_allocation(res, layout);
+
+    res
 }
 
 extern "C" fn rust_calloc(ctx: *mut c_void, nelem: size_t, elsize: size_t) -> *mut c_void {
@@ -287,23 +298,29 @@ extern "C" fn mimalloc_calloc(_ctx: *mut c_void, nelem: size_t, elsize: size_t)
 }
 
 #[cfg(feature = "snmalloc-sys")]
-extern "C" fn snmalloc_calloc(_ctx: *mut c_void, nelem: size_t, elsize: size_t) -> *mut c_void {
+extern "C" fn snmalloc_calloc(ctx: *mut c_void, nelem: size_t, elsize: size_t) -> *mut c_void {
     let size = match nelem * elsize {
         0 => 1,
         val => val,
     };
 
-    let ptr = unsafe { snmalloc_sys::rust_alloc(SNMALLOC_ALIGNMENT, size) };
+    let mut tracker = AllocationTracker::from_owned_ptr(ctx);
+
+    let layout = unsafe { alloc::Layout::from_size_align_unchecked(size, SNMALLOC_ALIGNMENT) };
+
+    let ptr: *mut c_void = unsafe { snmalloc_sys::rust_alloc(SNMALLOC_ALIGNMENT, size) } as *mut _;
     if ptr.is_null() {
-        return ptr as *mut _;
+        return ptr;
     }
 
     // TODO should we use write_volatile() + memory fence to be sure?
     unsafe {
         std::ptr::write_bytes(ptr, 0, size);
     }
 
-    ptr as *mut _
+    tracker.insert_allocation(ptr, layout);
+
+    ptr
 }
 
 extern "C" fn rust_realloc(ctx: *mut c_void, ptr: *mut c_void, new_size: size_t) -> *mut c_void {
@@ -379,8 +396,24 @@ extern "C" fn snmalloc_realloc(
         val => val,
     };
 
-    // TODO pass old_size properly.
-    unsafe { snmalloc_sys::rust_realloc(ptr as *mut _, SNMALLOC_ALIGNMENT, 0, new_size) as *mut _ }
+    let mut tracker = AllocationTracker::from_owned_ptr(ctx);
+
+    let layout = unsafe { alloc::Layout::from_size_align_unchecked(new_size, SNMALLOC_ALIGNMENT) };
+
+    let old_layout = tracker.remove_allocation(ptr);
+
+    let res = unsafe {
+        snmalloc_sys::rust_realloc(
+            ptr as *mut _,
+            SNMALLOC_ALIGNMENT,
+            old_layout.size(),
+            new_size,
+        ) as *mut _
+    };
+
+    tracker.insert_allocation(res, layout);
+
+    res
 }
 
 extern "C" fn rust_free(ctx: *mut c_void, ptr: *mut c_void) {
@@ -420,15 +453,22 @@ extern "C" fn mimalloc_free(_ctx: *mut c_void, ptr: *mut c_void) {
 }
 
 #[cfg(feature = "snmalloc-sys")]
-extern "C" fn snmalloc_free(_ctx: *mut c_void, ptr: *mut c_void) {
+extern "C" fn snmalloc_free(ctx: *mut c_void, ptr: *mut c_void) {
     if ptr.is_null() {
         return;
     }
 
-    // TODO pass size properly.
+    let mut tracker = AllocationTracker::from_owned_ptr(ctx);
+
+    let layout = tracker
+        .get_allocation(ptr)
+        .unwrap_or_else(|| panic!("could not find allocated memory record: {:?}", ptr));
+
     unsafe {
-        snmalloc_sys::rust_dealloc(ptr as *mut _, SNMALLOC_ALIGNMENT, 0);
+        snmalloc_sys::rust_dealloc(ptr as *mut _, SNMALLOC_ALIGNMENT, layout.size());
     }
+
+    tracker.remove_allocation(ptr);
 }
 
 extern "C" fn rust_arena_free(ctx: *mut c_void, ptr: *mut c_void, _size: size_t) {
@@ -468,14 +508,18 @@ extern "C" fn mimalloc_arena_free(_ctx: *mut c_void, ptr: *mut c_void, _size: si
 }
 
 #[cfg(feature = "snmalloc-sys")]
-extern "C" fn snmalloc_arena_free(_ctx: *mut c_void, ptr: *mut c_void, size: size_t) {
+extern "C" fn snmalloc_arena_free(ctx: *mut c_void, ptr: *mut c_void, size: size_t) {
     if ptr.is_null() {
         return;
     }
 
+    let mut tracker = AllocationTracker::from_owned_ptr(ctx);
+
     unsafe {
         snmalloc_sys::rust_dealloc(ptr as *mut _, SNMALLOC_ALIGNMENT, size);
     }
+
+    tracker.remove_allocation(ptr);
 }
 
 /// Represents a `PyMemAllocatorEx` that can be installed as a memory allocator.
@@ -596,24 +640,25 @@ impl PythonMemoryAllocator {
     /// Construct a new instance using snmalloc.
     #[cfg(feature = "snmalloc-sys")]
     pub fn snmalloc() -> Self {
-        panic!("snmalloc is not yet fully implemented");
+        let state = Box::into_raw(AllocationTracker::new());
 
         Self {
             backend: MemoryAllocatorBackend::Snmalloc,
-            instance: AllocatorInstance::Simple(
-                pyffi::PyMemAllocatorEx {
-                    ctx: std::ptr::null_mut(),
+            instance: AllocatorInstance::Tracking(TrackingAllocator {
+                allocator: pyffi::PyMemAllocatorEx {
+                    ctx: state as *mut c_void,
                     malloc: Some(snmalloc_malloc),
                     calloc: Some(snmalloc_calloc),
                     realloc: Some(snmalloc_realloc),
                     free: Some(snmalloc_free),
                 },
-                pyffi::PyObjectArenaAllocator {
+                arena: pyffi::PyObjectArenaAllocator {
                     ctx: std::ptr::null_mut(),
                     alloc: Some(snmalloc_malloc),
                     free: Some(snmalloc_arena_free),
                 },
-            ),
+                _state: unsafe { Box::from_raw(state) },
+            }),
         }
     }
 

pyembed/src/test/interpreter_config.rs

@@ -214,8 +214,6 @@ rusty_fork_test! {
 
     #[cfg(feature = "snmalloc")]
     #[test]
-    // snmalloc not yet implemented.
-    #[should_panic]
     fn test_allocator_snmalloc() {
         let mut config = OxidizedPythonInterpreterConfig::default();
         // Otherwise the Rust arguments are interpreted as Python arguments.