Add supprt for new c23 deallocators

This commit introduces 2 new functions to add support for
`free_sized` and `free_aligned_sized` deallocators introduced in
c23.

We ensure that `free_sized` and `free_aligned_sized` are skipped if
unavailable otherwise we get a `FREE` record recorded for them.

Signed-off-by: Nana Adjei Manu <[email protected]>

Author: claeusdev

src/memray/_memray/hooks.cpp

@@ -196,6 +196,20 @@ free(void* ptr) noexcept
     }
 }
 
+#if defined(__GLIBC__)
+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);
+}
+#endif
+
 void*
 realloc(void* ptr, size_t size) noexcept
 {

src/memray/_memray/hooks.h

@@ -31,6 +31,8 @@
         FOR_EACH_HOOKED_FUNCTION(prctl)                                                                 \
         FOR_EACH_HOOKED_FUNCTION(pvalloc)                                                               \
         FOR_EACH_HOOKED_FUNCTION(mmap64)
+        FOR_EACH_HOOKED_FUNCTION(free_sized)                                                            \
+        FOR_EACH_HOOKED_FUNCTION(free_aligned_sized)
 #else
 #    define MEMRAY_PLATFORM_HOOKED_FUNCTIONS                                                            \
         FOR_EACH_HOOKED_FUNCTION(memalign)                                                              \
@@ -199,6 +201,12 @@ mmap(void* addr, size_t length, int prot, int flags, int fd, off_t offset) noexc
 #if defined(__GLIBC__)
 void*
 mmap64(void* addr, size_t length, int prot, int flags, int fd, off64_t offset) noexcept;
+
+void
+free_sized(void* ptr, size_t size) noexcept;
+
+void
+free_aligned_sized(void* ptr, size_t alignment, size_t size) noexcept;
 #endif
 
 int

tests/integration/free_sized_extension/free_sized_test.cpp

@@ -0,0 +1,111 @@
+#define PY_SSIZE_T_CLEAN
+#include <Python.h>
+
+#include <cassert>
+#include <cstdlib>
+
+#ifdef __linux__
+#include <malloc.h>
+#endif
+
+namespace {  // unnamed
+
+extern "C" void
+test_free_sized()
+{
+    void* ptr = malloc(1024);
+    assert(ptr != nullptr);
+    
+    #ifdef __GLIBC__
+    extern void free_sized(void* ptr, size_t size);
+    free_sized(ptr, 1024);
+    #else
+    free(ptr);
+    #endif
+}
+
+extern "C" void
+test_free_aligned_sized()
+{
+    void* ptr = aligned_alloc(64, 1024);
+    assert(ptr != nullptr);
+    
+    #ifdef __GLIBC__
+    extern void free_aligned_sized(void* ptr, size_t alignment, size_t size);
+    free_aligned_sized(ptr, 64, 1024);
+    #else
+    free(ptr);
+    #endif
+}
+
+extern "C" void
+test_both_free_functions()
+{
+    void* ptr1 = malloc(512);
+    assert(ptr1 != nullptr);
+    #ifdef __GLIBC__
+    extern void free_sized(void* ptr, size_t size);
+    free_sized(ptr1, 512);
+    #else
+    free(ptr1);
+    #endif
+    
+    void* ptr2 = aligned_alloc(32, 256);
+    assert(ptr2 != nullptr);
+    #ifdef __GLIBC__
+    extern void free_aligned_sized(void* ptr, size_t alignment, size_t size);
+    free_aligned_sized(ptr2, 32, 256);
+    #else
+    free(ptr2);
+    #endif
+}
+
+PyObject*
+run_free_sized_test(PyObject*, PyObject*)
+{
+    test_free_sized();
+    Py_RETURN_NONE;
+}
+
+PyObject*
+run_free_aligned_sized_test(PyObject*, PyObject*)
+{
+    test_free_aligned_sized();
+    Py_RETURN_NONE;
+}
+
+PyObject*
+run_both_tests(PyObject*, PyObject*)
+{
+    test_both_free_functions();
+    Py_RETURN_NONE;
+}
+
+}  // unnamed namespace
+
+static PyMethodDef
+free_sized_methods[] = {
+    {"run_free_sized_test", run_free_sized_test, METH_NOARGS, "Test free_sized function"},
+    {"run_free_aligned_sized_test", run_free_aligned_sized_test, METH_NOARGS, "Test free_aligned_sized function"},
+    {"run_both_tests", run_both_tests, METH_NOARGS, "Test both free functions"},
+    {nullptr, nullptr, 0, nullptr}
+};
+
+static PyModuleDef
+free_sized_module = {
+    PyModuleDef_HEAD_INIT,
+    "free_sized_test",
+    "Test module for free_sized and free_aligned_sized functions",
+    -1,
+    free_sized_methods,
+    nullptr,
+    nullptr,
+    nullptr,
+    nullptr
+};
+
+PyMODINIT_FUNC
+PyInit_free_sized_test()
+{
+    return PyModule_Create(&free_sized_module);
+}

tests/integration/free_sized_extension/setup.py

@@ -0,0 +1,12 @@
+from setuptools import Extension, setup
+
+setup(
+    name="free_sized_extension",
+    ext_modules=[
+        Extension(
+            "free_sized_test",
+            sources=["free_sized_test.cpp"],
+            language="c++",
+        )
+    ],
+)

tests/integration/test_extensions.py

@@ -15,6 +15,7 @@
 TEST_MULTITHREADED_EXTENSION = HERE / "multithreaded_extension"
 TEST_MISBEHAVING_EXTENSION = HERE / "misbehaving_extension"
 TEST_RPATH_EXTENSION = HERE / "rpath_extension"
+TEST_FREE_SIZED_EXTENSION = HERE / "free_sized_extension"
 
 
 @pytest.mark.valgrind
@@ -386,3 +387,49 @@ def test_dlopen_with_rpath(tmpdir, monkeypatch):
         # THEN
         with Tracker(output):
             hello_world()
+
+@pytest.mark.skipif(
+    not hasattr(ctypes.CDLL(None), "free_sized"),
+    reason="free_sized not available on this system"
+)
+def test_free_sized_extension(tmpdir, monkeypatch):
+    """Test tracking allocations in a native extension which uses free_sized and free_aligned_sized."""
+    # GIVEN
+    output = Path(tmpdir) / "test.bin"
+    extension_name = "free_sized_extension"
+    extension_path = tmpdir / extension_name
+    shutil.copytree(TEST_FREE_SIZED_EXTENSION, extension_path)
+    subprocess.run(
+        [sys.executable, str(extension_path / "setup.py"), "build_ext", "--inplace"],
+        check=True,
+        cwd=extension_path,
+        capture_output=True,
+    )
+
+    # WHEN
+    with monkeypatch.context() as ctx:
+        ctx.setattr(sys, "path", [*sys.path, str(extension_path)])
+        from free_sized_test import run_both_tests  # type: ignore
+
+        with Tracker(output):
+            run_both_tests()
+
+    # THEN
+    records = list(FileReader(output).get_allocation_records())
+    assert records
+
+    # Check that at least 2 allocations from malloc and aligned_alloc
+    mallocs = [record for record in records if record.allocator == AllocatorType.MALLOC]
+    assert len(mallocs) >= 2
+
+    # Check that corresponding FREE records - this verifies hooks are working!
+    mallocs_addr = {record.address for record in mallocs}
+    frees = [
+        record
+        for record in records
+        if record.address in mallocs_addr and record.allocator == AllocatorType.FREE
+    ]
+    assert len(frees) >= 2
+    
+    assert all(len(malloc.stack_trace()) == 0 for malloc in mallocs)
+    assert all(len(free.stack_trace()) == 0 for free in frees)