diff --git a/src/memray/_memray/hooks.cpp b/src/memray/_memray/hooks.cpp index 170f029ab6..10115350b7 100644 --- a/src/memray/_memray/hooks.cpp +++ b/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 { diff --git a/src/memray/_memray/hooks.h b/src/memray/_memray/hooks.h index 0b3fef3c2b..4a4d5e9418 100644 --- a/src/memray/_memray/hooks.h +++ b/src/memray/_memray/hooks.h @@ -38,6 +38,14 @@ FOR_EACH_HOOKED_FUNCTION(prctl) #endif +#if defined(__STDC_VERSION__) && __STDC_VERSION__ >= 202311L +# define MEMRAY_C23_HOOKED_FUNCTIONS \ + FOR_EACH_HOOKED_FUNCTION(free_sized) \ + FOR_EACH_HOOKED_FUNCTION(free_aligned_sized) +#else +# define MEMRAY_C23_HOOKED_FUNCTIONS +#endif + #define MEMRAY_HOOKED_FUNCTIONS \ FOR_EACH_HOOKED_FUNCTION(malloc) \ FOR_EACH_HOOKED_FUNCTION(free) \ @@ -51,6 +59,7 @@ FOR_EACH_HOOKED_FUNCTION(dlopen) \ FOR_EACH_HOOKED_FUNCTION(dlclose) \ FOR_EACH_HOOKED_FUNCTION(PyGILState_Ensure) \ + MEMRAY_C23_HOOKED_FUNCTIONS \ MEMRAY_PLATFORM_HOOKED_FUNCTIONS namespace memray::hooks { @@ -202,6 +211,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; diff --git a/tests/integration/free_sized_extension/free_sized_test.c b/tests/integration/free_sized_extension/free_sized_test.c new file mode 100644 index 0000000000..447320fba1 --- /dev/null +++ b/tests/integration/free_sized_extension/free_sized_test.c @@ -0,0 +1,150 @@ +#define PY_SSIZE_T_CLEAN +#include + +#include +#include +#include + +__attribute__((weak)) void free_sized(void* ptr, size_t size); +__attribute__((weak)) void free_aligned_sized(void* ptr, size_t alignment, size_t size); + +// Check if C23 functions are available +static int functions_available = -1; + +static void check_functions_available(void) { + if (functions_available == -1) { + functions_available = (free_sized != NULL && free_aligned_sized != NULL); + } +} + +// Structure to hold allocation info +void* +test_free_sized(void) +{ + void* address; + + check_functions_available(); + if (!functions_available) { + return address; + } + + void* ptr = malloc(1024); + assert(ptr != NULL); + + address = ptr; + + free_sized(ptr, 1024); + return address; +} + +void* +test_free_aligned_sized(void) +{ + void* address; + + check_functions_available(); + if (!functions_available) { + return address; + } + + void* ptr = aligned_alloc(64, 1024); + assert(ptr != NULL); + + address = ptr; + + free_aligned_sized(ptr, 64, 1024); + return address; +} + +void* +test_both_free_functions(void) +{ + void* address; + + check_functions_available(); + if (!functions_available) { + return NULL; + } + + void* ptr1 = malloc(512); + assert(ptr1 != NULL); + free_sized(ptr1, 512); + + void* ptr2 = aligned_alloc(32, 256); + assert(ptr2 != NULL); + free_aligned_sized(ptr2, 32, 256); + + address = ptr2; + + return address; +} + +PyObject* +run_free_sized_test(PyObject* self, PyObject* args) +{ + check_functions_available(); + if (!functions_available) { + Py_RETURN_NONE; + } + + void* address = test_free_sized(); + + // Return address for verification + PyObject* result = Py_BuildValue("(KII)", address); + return result; +} + +PyObject* +run_free_aligned_sized_test(PyObject* self, PyObject* args) +{ + check_functions_available(); + if (!functions_available) { + Py_RETURN_NONE; + } + + void* address = test_free_aligned_sized(); + + PyObject* result = Py_BuildValue("(KII)", address); + return result; +} + +PyObject* +run_both_tests(PyObject* self, PyObject* args) +{ + check_functions_available(); + if (!functions_available) { + Py_RETURN_NONE; // Skip test if functions not available + } + + void* address = test_both_free_functions(); + + PyObject* result = Py_BuildValue("(KII)", address); + return result; +} + +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"}, + {NULL, NULL, 0, NULL} +}; + +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, + NULL, + NULL, + NULL, + NULL +}; + +PyMODINIT_FUNC +PyInit_free_sized_test(void) +{ + return PyModule_Create(&free_sized_module); +} diff --git a/tests/integration/free_sized_extension/setup.py b/tests/integration/free_sized_extension/setup.py new file mode 100644 index 0000000000..597eb7ee62 --- /dev/null +++ b/tests/integration/free_sized_extension/setup.py @@ -0,0 +1,14 @@ +from setuptools import Extension +from setuptools import setup + +setup( + name="free_sized_extension", + ext_modules=[ + Extension( + "free_sized_test", + sources=["free_sized_test.c"], + language="c", + extra_compile_args=["-std=c23"], + ) + ], +) diff --git a/tests/integration/test_extensions.py b/tests/integration/test_extensions.py index c8f6e429ed..27189d965b 100644 --- a/tests/integration/test_extensions.py +++ b/tests/integration/test_extensions.py @@ -1,3 +1,4 @@ +import ctypes import shutil import subprocess import sys @@ -15,6 +16,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 @@ -108,7 +110,7 @@ def allocating_function(): # pragma: no cover func, filename, line = bottom_frame assert func == "allocating_function" assert filename.endswith(__file__) - assert line == 83 + assert line == 85 frees = [ event @@ -171,7 +173,7 @@ def foo2(): func, filename, line = bottom_frame assert func == "test_extension_that_uses_pygilstate_ensure" assert filename.endswith(__file__) - assert line == 154 + assert line == 156 # We should have 2 frames here: this function calling `allocator.valloc`, # and `allocator.valloc` calling the C `valloc`. @@ -186,7 +188,7 @@ def foo2(): func, filename, line = caller assert func == "test_extension_that_uses_pygilstate_ensure" assert filename.endswith(__file__) - assert line == 155 + assert line == 157 frees = [ event @@ -242,7 +244,7 @@ def allocating_function(): func, filename, line = bottom_frame assert func == "test_native_dlopen" assert filename.endswith(__file__) - assert line == 226 + assert line == 228 frees = [ event @@ -386,3 +388,99 @@ 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 allocations in a native extension using 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) + + # Try to build the extension, skip if compilation fails + try: + subprocess.run( + [ + sys.executable, + str(extension_path / "setup.py"), + "build_ext", + "--inplace", + ], + check=True, + cwd=extension_path, + capture_output=True, + ) + except subprocess.CalledProcessError as e: + pytest.skip(f"Failed to compile C23 extension: {e}") + except Exception as e: + pytest.skip(f"Unexpected error building extension: {e}") + + # WHEN + with monkeypatch.context() as ctx: + ctx.setattr(sys, "path", [*sys.path, str(extension_path)]) + + try: + from free_sized_test import run_both_tests # type: ignore + except ImportError as e: + pytest.skip(f"Failed to import compiled extension: {e}") + + with Tracker(output): + # Get allocation info from the extension + result = run_both_tests() + + # Skip test if functions not available (e.g., on macOS) + if result is None: + pytest.skip("C23 functions not available on this system") + + # 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.ALIGNED_ALLOC + ] + 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) == len(mallocs) + + assert all(len(malloc.stack_trace()) == 0 for malloc in mallocs) + assert all(len(free.stack_trace()) == 0 for free in frees) + + # Verify that the specific addresses returned by the extension were tracked + if result is not None: + expected_address = result[0] + + # Find the allocation record for this address + matching_allocs = [ + record + for record in records + if record.address == expected_address + and record.allocator == AllocatorType.ALIGNED_ALLOC + ] + assert ( + len(matching_allocs) >= 1 + ), f"Expected allocation at address {expected_address} not found" + + # Find the corresponding free record + matching_frees = [ + record + for record in records + if record.address == expected_address + and record.allocator == AllocatorType.FREE + ] + assert ( + len(matching_frees) >= 1 + ), f"Expected free at address {expected_address} not found"