Intercept C23 free_sized() and free_aligned_sized()

news/725.feature.rst

@@ -0,0 +1 @@
+Intercept the C23 ``free_sized()`` and ``free_aligned_sized()`` deallocators on Linux. Calls to either are now recorded as ``FREE`` allocations, matching the behavior of plain ``free()``.

src/memray/_memray/hooks.cpp

@@ -196,6 +196,44 @@ free(void* ptr) noexcept
     }
 }
 
+#if defined(__GLIBC__)
+void
+free_sized(void* ptr, size_t size) noexcept
+{
+    if (ptr != nullptr) {
+        tracking_api::Tracker::trackDeallocation(ptr, 0, hooks::Allocator::FREE);
+    }
+
+    {
+        tracking_api::RecursionGuard guard;
+        if (MEMRAY_ORIG(free_sized)) {
+            MEMRAY_ORIG(free_sized)(ptr, size);
+        } else {
+            // libc lacks free_sized: fall back to plain free, which is a
+            // valid implementation per the C23 specification.
+            MEMRAY_ORIG(free)(ptr);
+        }
+    }
+}
+
+void
+free_aligned_sized(void* ptr, size_t alignment, size_t size) noexcept
+{
+    if (ptr != nullptr) {
+        tracking_api::Tracker::trackDeallocation(ptr, 0, hooks::Allocator::FREE);
+    }
+
+    {
+        tracking_api::RecursionGuard guard;
+        if (MEMRAY_ORIG(free_aligned_sized)) {
+            MEMRAY_ORIG(free_aligned_sized)(ptr, alignment, size);
+        } else {
+            MEMRAY_ORIG(free)(ptr);
+        }
+    }
+}
+#endif
+
 void*
 realloc(void* ptr, size_t size) noexcept
 {

src/memray/_memray/hooks.h

@@ -14,12 +14,38 @@
 #        define _GNU_SOURCE
 #    endif
 #    include "elf_utils.h"
+#    include <features.h>
 #    include <malloc.h>
 #    include <sys/prctl.h>
 #endif
 
 #include <dlfcn.h>
 
+// C23 introduced free_sized() and free_aligned_sized(). They may be missing
+// from libc when memray is built against older system headers, but a binary
+// memray runs in may still call them (e.g. when the wheel is built against
+// an old glibc but installed on a newer one). Declare them with weak linkage
+// so the build succeeds either way; at tracker startup we resolve the actual
+// runtime addresses via dl_iterate_phdr.
+#if defined(__linux__)
+#    if defined(__GLIBC__)
+#        if !__GLIBC_PREREQ(2, 42)
+#            define MEMRAY_NEED_FREE_SIZED_WEAK_DECLS 1
+#        endif
+#    else
+#        define MEMRAY_NEED_FREE_SIZED_WEAK_DECLS 1
+#    endif
+#endif
+
+#ifdef MEMRAY_NEED_FREE_SIZED_WEAK_DECLS
+extern "C" {
+void
+free_sized(void* ptr, size_t size) __attribute__((weak));
+void
+free_aligned_sized(void* ptr, size_t alignment, size_t size) __attribute__((weak));
+}
+#endif
+
 #include "alloc.h"
 #include "compat.h"
 #include "logging.h"
@@ -31,7 +57,9 @@
         FOR_EACH_HOOKED_FUNCTION(memalign)                                                              \
         FOR_EACH_HOOKED_FUNCTION(prctl)                                                                 \
         FOR_EACH_HOOKED_FUNCTION(pvalloc)                                                               \
-        FOR_EACH_HOOKED_FUNCTION(mmap64)
+        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)                                                              \
@@ -167,6 +195,14 @@ malloc(size_t size) noexcept;
 void
 free(void* ptr) noexcept;
 
+#if defined(__GLIBC__)
+void
+free_sized(void* ptr, size_t size) noexcept;
+
+void
+free_aligned_sized(void* ptr, size_t alignment, size_t size) noexcept;
+#endif
+
 void*
 realloc(void* ptr, size_t size) noexcept;
 

src/memray/commands/run.py

@@ -62,7 +62,13 @@ def _run_tracker(
                     sys.path[0] = os.getcwd()
                 # run_module will replace argv[0] with the script's path
                 sys.argv = ["", *args.script_args]
-                runpy.run_module(args.script, run_name="__main__", alter_sys=True)
+                try:
+                    runpy.run_module(args.script, run_name="__main__", alter_sys=True)
+                except ImportError as e:
+                    if e.name in (None, args.script):
+                        post_run_message = None  # Initial import failed
+                        raise MemrayCommandError(f"memray: {e}", exit_code=1)
+                    raise
             elif args.run_as_cmd:
                 if _should_modify_sys_path():
                     sys.path[0] = ""
@@ -314,6 +320,8 @@ def validate_target_file(self, args: argparse.Namespace) -> None:
                 "Only valid Python files or commands can be executed under memray",
                 exit_code=1,
             )
+        except (FileNotFoundError, PermissionError) as error:
+            raise MemrayCommandError(str(error), exit_code=1)
 
     def run(self, args: argparse.Namespace, parser: argparse.ArgumentParser) -> None:
         if args.no_compress:

tests/integration/multithreaded_extension/testext.cpp

@@ -3,11 +3,23 @@
 
 #include <assert.h>
 #include <pthread.h>
+#include <stdlib.h>
 
 #ifdef __linux__
+#include <features.h>
 #include <malloc.h>
 #endif
 
+// Weak forward declarations so the test extension builds against pre-C23
+// libc headers but can still call free_sized / free_aligned_sized when the
+// running process is linked against a newer libc.
+#if defined(__linux__) && (!defined(__GLIBC__) || !__GLIBC_PREREQ(2, 42))
+extern "C" {
+void free_sized(void* ptr, size_t size) __attribute__((weak));
+void free_aligned_sized(void* ptr, size_t alignment, size_t size) __attribute__((weak));
+}
+#endif
+
 namespace {  // unnamed
 
 #pragma GCC push_options
@@ -92,13 +104,52 @@ run_valloc_at_exit(PyObject*, PyObject*)
     Py_RETURN_NONE;
 }
 
+PyObject*
+has_free_sized(PyObject*, PyObject*)
+{
+#if defined(__linux__)
+    if (&free_sized != nullptr && &free_aligned_sized != nullptr) {
+        Py_RETURN_TRUE;
+    }
+#endif
+    Py_RETURN_FALSE;
+}
+
+PyObject*
+run_free_sized(PyObject*, PyObject*)
+{
+#if defined(__linux__)
+    if (&free_sized == nullptr || &free_aligned_sized == nullptr) {
+        PyErr_SetString(PyExc_RuntimeError, "libc lacks free_sized / free_aligned_sized");
+        return NULL;
+    }
+    const size_t plain_size = 128;
+    void* p = malloc(plain_size);
+    if (!p) return PyErr_NoMemory();
+    free_sized(p, plain_size);
+
+    const size_t alignment = 64;
+    const size_t aligned_size = 256;
+    void* q = aligned_alloc(alignment, aligned_size);
+    if (!q) return PyErr_NoMemory();
+    free_aligned_sized(q, alignment, aligned_size);
+
+    return Py_BuildValue("(KK)", (unsigned long long)p, (unsigned long long)q);
+#else
+    PyErr_SetString(PyExc_RuntimeError, "free_sized only tested on Linux");
+    return NULL;
+#endif
+}
+
 #pragma GCC pop_options
 
 }  // unnamed namespace
 
 static PyMethodDef methods[] = {
         {"run", run, METH_NOARGS, "Run a bunch of threads"},
         {"run_valloc_at_exit", run_valloc_at_exit, METH_NOARGS, "Run valloc while exiting a thread"},
+        {"has_free_sized", has_free_sized, METH_NOARGS, "Whether libc provides C23 sized deallocators"},
+        {"run_free_sized", run_free_sized, METH_NOARGS, "Allocate and free using C23 sized deallocators"},
         {NULL, NULL, 0, NULL},
 };
 

tests/integration/test_extensions.py

@@ -327,6 +327,48 @@ def test_valloc_at_thread_exit_in_subprocess(tmpdir, monkeypatch):
     assert len(vallocs) == 1
 
 
+def test_free_sized_extension(tmpdir, monkeypatch):
+    """Test tracking of C23 free_sized() and free_aligned_sized()."""
+    output = Path(tmpdir) / "test.bin"
+    extension_name = "multithreaded_extension"
+    extension_path = tmpdir / extension_name
+    shutil.copytree(TEST_MULTITHREADED_EXTENSION, extension_path)
+    subprocess.run(
+        [sys.executable, str(extension_path / "setup.py"), "build_ext", "--inplace"],
+        check=True,
+        cwd=extension_path,
+        capture_output=True,
+    )
+
+    with monkeypatch.context() as ctx:
+        ctx.setattr(sys, "path", [*sys.path, str(extension_path)])
+        from testext import has_free_sized  # type: ignore
+        from testext import run_free_sized  # type: ignore
+
+        if not has_free_sized():
+            pytest.skip(
+                "libc lacks free_sized / free_aligned_sized (needs glibc 2.42+)"
+            )
+
+        with Tracker(output):
+            plain_addr, aligned_addr = run_free_sized()
+
+    records = list(FileReader(output).get_allocation_records())
+
+    plain_frees = [
+        r
+        for r in records
+        if r.address == plain_addr and r.allocator == AllocatorType.FREE
+    ]
+    aligned_frees = [
+        r
+        for r in records
+        if r.address == aligned_addr and r.allocator == AllocatorType.FREE
+    ]
+    assert len(plain_frees) == 1, "expected one FREE record from free_sized"
+    assert len(aligned_frees) == 1, "expected one FREE record from free_aligned_sized"
+
+
 @pytest.mark.parametrize("py_finalize", [True, False])
 def test_hard_exit(tmpdir, py_finalize):
     """Test a program that exits directly under the context manager"""

tests/unit/test_cli.py

@@ -919,3 +919,26 @@ def test_parser_takes_force_flag(self):
         assert namespace.output == "output.html"
         assert namespace.force is True
         assert namespace.format == "gprof2dot"
+
+
+class TestRunCommandReal:
+    def test_run_non_existent_file(self, capsys):
+        # GIVEN
+        # WHEN
+        exit_code = main(["run", "non_existent_file.py"])
+
+        # THEN
+        assert exit_code == 1
+        captured = capsys.readouterr()
+        assert "No such file or directory: 'non_existent_file.py'" in captured.err
+
+    def test_run_non_existent_module(self, capsys):
+        # GIVEN
+        # WHEN
+        exit_code = main(["run", "-m", "non_existent_module"])
+
+        # THEN
+        assert exit_code == 1
+        captured = capsys.readouterr()
+        assert "No module named non_existent_module" in captured.err
+        assert "[memray] Successfully generated profile results." not in captured.out