Add support for the new C23 deallocators

news/725.feature.rst

@@ -0,0 +1 @@
+Recognize deallocations performed with the ``free_sized`` and ``free_aligned_sized`` functions introduced in C23.

src/memray/_memray/hooks.cpp

@@ -196,6 +196,18 @@ free(void* ptr) noexcept
     }
 }
 
+void
+free_sized(void* ptr, size_t size) noexcept
+{
+    memray::intercept::free(ptr);
+}
+
+void
+free_aligned_sized(void* ptr, size_t alignment, size_t size) noexcept
+{
+    memray::intercept::free(ptr);
+}
+
 void*
 realloc(void* ptr, size_t size) noexcept
 {

src/memray/_memray/hooks.h

@@ -24,6 +24,18 @@
 #include "compat.h"
 #include "logging.h"
 
+#ifdef __APPLE__
+#    define MAYBE_MISSING __attribute__((weak_import))
+#else
+#    define MAYBE_MISSING __attribute__((weak))
+#endif
+
+extern "C" void
+free_sized(void*, size_t) MAYBE_MISSING;
+
+extern "C" void
+free_aligned_sized(void*, size_t, size_t) MAYBE_MISSING;
+
 #if defined(__APPLE__)
 #    define MEMRAY_PLATFORM_HOOKED_FUNCTIONS
 #elif defined(__GLIBC__)
@@ -51,6 +63,8 @@
     FOR_EACH_HOOKED_FUNCTION(dlopen)                                                                    \
     FOR_EACH_HOOKED_FUNCTION(dlclose)                                                                   \
     FOR_EACH_HOOKED_FUNCTION(PyGILState_Ensure)                                                         \
+    FOR_EACH_HOOKED_FUNCTION(free_sized)                                                                \
+    FOR_EACH_HOOKED_FUNCTION(free_aligned_sized)                                                        \
     MEMRAY_PLATFORM_HOOKED_FUNCTIONS
 
 namespace memray::hooks {
@@ -202,6 +216,12 @@ void*
 mmap64(void* addr, size_t length, int prot, int flags, int fd, off64_t offset) noexcept;
 #endif
 
+void
+free_sized(void* ptr, size_t size) noexcept;
+
+void
+free_aligned_sized(void* ptr, size_t alignment, size_t size) noexcept;
+
 int
 munmap(void* addr, size_t length) noexcept;
 

src/memray/_memray_test_utils.pyx

@@ -67,6 +67,11 @@ cdef extern from *:
     int set_thread_name_impl(const char* new_name)
 
 
+cdef extern from "hooks.h":
+    void free_sized(void*, size_t)
+    void free_aligned_sized(void*, size_t, size_t)
+
+
 def set_thread_name(new_name):
     return set_thread_name_impl(new_name)
 
@@ -83,6 +88,24 @@ cdef class MemoryAllocator:
         free(self.ptr)
         self.ptr = NULL
 
+    def free_sized(self, size_t size):
+        if self.ptr == NULL:
+            raise RuntimeError("Pointer cannot be NULL")
+        if &free_sized == NULL:
+            return False
+        free_sized(self.ptr, size)
+        self.ptr = NULL
+        return True
+
+    def free_aligned_sized(self, size_t size):
+        if self.ptr == NULL:
+            raise RuntimeError("Pointer cannot be NULL")
+        if &free_aligned_sized == NULL:
+            return False
+        free_aligned_sized(self.ptr, sizeof(void*), size)
+        self.ptr = NULL
+        return True
+
     def malloc(self, size_t size):
         self.ptr = malloc(size)
         return self.ptr != NULL

src/memray/_test.py

@@ -22,6 +22,12 @@ def __init__(self) -> None:
     def free(self) -> None:
         return self.allocator.free()
 
+    def free_sized(self, size: int) -> bool:
+        return self.allocator.free_sized(size)
+
+    def free_aligned_sized(self, size: int) -> bool:
+        return self.allocator.free_aligned_sized(size)
+
     def malloc(self, size: int) -> bool:
         return self.allocator.malloc(size)
 

src/memray/_test_utils.pyi

@@ -5,6 +5,8 @@ from ._memray import PymallocDomain as PymallocDomain
 class MemoryAllocator:
     def __init__(self) -> None: ...
     def free(self) -> None: ...
+    def free_sized(self, size: int) -> bool: ...
+    def free_aligned_sized(self, size: int) -> bool: ...
     def malloc(self, size: int) -> bool: ...
     def calloc(self, size: int) -> bool: ...
     def realloc(self, size: int) -> bool: ...

tests/integration/test_tracing.py

@@ -143,7 +143,7 @@ def test_cython_traceback(tmpdir):
 
     traceback = list(alloc1.stack_trace())
     assert traceback == [
-        ("valloc", sys.modules["memray._test"].__file__, 44),
+        ("valloc", sys.modules["memray._test"].__file__, 50),
         ("test_cython_traceback", __file__, 134),
     ]
 

tests/integration/test_tracking.py

@@ -123,6 +123,71 @@ def test_simple_cpp_allocation_tracking(tmp_path):
     assert len(frees) >= 1
 
 
+def test_intercepting_free_sized(tmp_path):
+    # GIVEN
+    allocator = MemoryAllocator()
+    output = tmp_path / "test.bin"
+
+    # WHEN
+    with Tracker(output):
+        res = allocator.valloc(ALLOC_SIZE)
+        assert res
+        res = allocator.free_sized(ALLOC_SIZE)
+
+    if not res:
+        pytest.skip("Deallocator free_sized is not supported by this platform")
+
+    # THEN
+    allocations = list(FileReader(output).get_allocation_records())
+    allocs = [
+        event
+        for event in allocations
+        if event.size == ALLOC_SIZE and event.allocator == AllocatorType.VALLOC
+    ]
+    assert len(allocs) == 1
+    (alloc,) = allocs
+
+    frees = [
+        event
+        for event in allocations
+        if event.address == alloc.address and event.allocator == AllocatorType.FREE
+    ]
+    assert len(frees) >= 1
+
+
+def test_intercepting_free_aligned_sized(tmp_path):
+    # GIVEN
+    allocator = MemoryAllocator()
+    output = tmp_path / "test.bin"
+
+    # WHEN
+    with Tracker(output):
+        res = allocator.aligned_alloc(ALLOC_SIZE)
+        if not res:
+            pytest.skip("Allocator aligned_alloc is not supported by this platform")
+        res = allocator.free_aligned_sized(ALLOC_SIZE)
+
+    if not res:
+        pytest.skip("Deallocator free_aligned_sized is not supported by this platform")
+
+    # THEN
+    allocations = list(FileReader(output).get_allocation_records())
+    allocs = [
+        event
+        for event in allocations
+        if event.size == ALLOC_SIZE and event.allocator == AllocatorType.ALIGNED_ALLOC
+    ]
+    assert len(allocs) == 1
+    (alloc,) = allocs
+
+    frees = [
+        event
+        for event in allocations
+        if event.address == alloc.address and event.allocator == AllocatorType.FREE
+    ]
+    assert len(frees) >= 1
+
+
 @pytest.mark.parametrize("domain", PYMALLOC_DOMAINS)
 @pytest.mark.parametrize(["allocator_func", "allocator_type"], PYMALLOC_ALLOCATORS)
 def test_simple_pymalloc_allocation_tracking(