Track Apache Arrow mimalloc allocations

.github/workflows/coverage.yml

@@ -57,7 +57,7 @@ jobs:
           echo /home/runner/node_modules/.bin >>"$GITHUB_PATH"
       - name: Install Python dependencies
         run: |
-          ${PYTHON} -m pip install --upgrade pip scikit-build-core cython pkgconfig
+          ${PYTHON} -m pip install --upgrade pip scikit-build-core cython pkgconfig "pyarrow>=25.0"
           make test-install
       - name: Disable ptrace security restrictions
         run: |

news/943.feature.rst

@@ -0,0 +1 @@
+Track individual allocations made by Apache Arrow 25.0 and newer when it uses mimalloc on Linux.

src/memray/_memray/elf_shenanigans.cpp

@@ -41,9 +41,17 @@ patch_symbol(
         const Hook& hook,
         typename Hook::signature_t intercept,
         const char* symname,
+        const SymbolTable& symbols,
         Addr addr,
         bool restore_original)
 {
+    const bool is_arrow_mimalloc_hook = strncmp(symname, "arrow_mi_", 9) == 0;
+    if (is_arrow_mimalloc_hook
+        && (!hook || symbols.getSymbolAddress(symname) != reinterpret_cast<Addr>(hook.d_original)))
+    {
+        return;
+    }
+
     // Make sure that we can read and write to the page where the address that we are trying to
     // patch;
     if (unprotect_page(addr) < 0) {
@@ -89,6 +97,7 @@ overwrite_elf_table(
                 MEMRAY_ORIG(hookname),                                                                  \
                 &intercept::hookname,                                                                   \
                 symname,                                                                                \
+                symbols,                                                                                \
                 symbol_addr,                                                                            \
                 restore_original);                                                                      \
         continue;                                                                                       \

src/memray/_memray/hooks.cpp

@@ -102,6 +102,19 @@ ensureAllHooksAreValid()
 #undef FOR_EACH_HOOKED_FUNCTION
 }
 
+void
+ensureArrowHooksAreValid()
+{
+#if defined(__linux__)
+    MEMRAY_ORIG(arrow_mi_malloc_aligned).ensureValidOriginalSymbol();
+    if (!MEMRAY_ORIG(arrow_mi_malloc_aligned)) {
+        return;
+    }
+    MEMRAY_ORIG(arrow_mi_realloc_aligned).ensureValidOriginalSymbol();
+    MEMRAY_ORIG(arrow_mi_free).ensureValidOriginalSymbol();
+#endif
+}
+
 }  // namespace memray::hooks
 
 namespace memray::intercept {
@@ -432,6 +445,56 @@ aligned_alloc(size_t alignment, size_t size) noexcept
 
 #if defined(__linux__)
 
+// Arrow's statically linked mimalloc is otherwise visible only as mmap arenas.
+void*
+arrow_mi_malloc_aligned(size_t size, size_t alignment) noexcept
+{
+    assert(MEMRAY_ORIG(arrow_mi_malloc_aligned));
+
+    void* ret;
+    {
+        tracking_api::RecursionGuard guard;
+        ret = MEMRAY_ORIG(arrow_mi_malloc_aligned)(size, alignment);
+    }
+    if (ret) {
+        tracking_api::Tracker::trackAllocation(ret, size, hooks::Allocator::ALIGNED_ALLOC);
+    }
+    return ret;
+}
+
+void*
+arrow_mi_realloc_aligned(void* ptr, size_t new_size, size_t alignment) noexcept
+{
+    assert(MEMRAY_ORIG(arrow_mi_realloc_aligned));
+
+    void* ret;
+    {
+        tracking_api::RecursionGuard guard;
+        ret = MEMRAY_ORIG(arrow_mi_realloc_aligned)(ptr, new_size, alignment);
+    }
+    if (ret) {
+        if (ptr != nullptr) {
+            tracking_api::Tracker::trackDeallocation(ptr, 0, hooks::Allocator::FREE);
+        }
+        tracking_api::Tracker::trackAllocation(ret, new_size, hooks::Allocator::REALLOC);
+    }
+    return ret;
+}
+
+void
+arrow_mi_free(void* ptr) noexcept
+{
+    assert(MEMRAY_ORIG(arrow_mi_free));
+
+    if (ptr != nullptr) {
+        tracking_api::Tracker::trackDeallocation(ptr, 0, hooks::Allocator::FREE);
+    }
+    {
+        tracking_api::RecursionGuard guard;
+        MEMRAY_ORIG(arrow_mi_free)(ptr);
+    }
+}
+
 void*
 memalign(size_t alignment, size_t size) noexcept
 {

src/memray/_memray/hooks.h

@@ -36,6 +36,24 @@ free_sized(void*, size_t) MAYBE_MISSING;
 extern "C" void
 free_aligned_sized(void*, size_t, size_t) MAYBE_MISSING;
 
+#ifdef __linux__
+// Arrow defines these wrappers only when built with mimalloc.
+extern "C" {
+void*
+arrow_mi_malloc_aligned(size_t size, size_t alignment) MAYBE_MISSING;
+void*
+arrow_mi_realloc_aligned(void* ptr, size_t new_size, size_t alignment) MAYBE_MISSING;
+void
+arrow_mi_free(void* ptr) MAYBE_MISSING;
+}
+#    define MEMRAY_ARROW_HOOKED_FUNCTIONS                                                               \
+        FOR_EACH_HOOKED_FUNCTION(arrow_mi_malloc_aligned)                                               \
+        FOR_EACH_HOOKED_FUNCTION(arrow_mi_realloc_aligned)                                              \
+        FOR_EACH_HOOKED_FUNCTION(arrow_mi_free)
+#else
+#    define MEMRAY_ARROW_HOOKED_FUNCTIONS
+#endif
+
 #if defined(__APPLE__)
 #    define MEMRAY_PLATFORM_HOOKED_FUNCTIONS
 #elif defined(__GLIBC__)
@@ -65,6 +83,7 @@ free_aligned_sized(void*, size_t, size_t) MAYBE_MISSING;
     FOR_EACH_HOOKED_FUNCTION(PyGILState_Ensure)                                                         \
     FOR_EACH_HOOKED_FUNCTION(free_sized)                                                                \
     FOR_EACH_HOOKED_FUNCTION(free_aligned_sized)                                                        \
+    MEMRAY_ARROW_HOOKED_FUNCTIONS                                                                       \
     MEMRAY_PLATFORM_HOOKED_FUNCTIONS
 
 namespace memray::hooks {
@@ -100,8 +119,8 @@ struct SymbolHook
         symbol_query query{0, d_symbol, nullptr};
         dl_iterate_phdr(&phdr_symfind_callback, (void*)&query);
         auto symbol_addr = reinterpret_cast<signature_t>(query.address);
-        if (symbol_addr != nullptr) {
-            if (symbol_addr != d_original) {
+        if (symbol_addr != nullptr && symbol_addr != d_original) {
+            if (d_original != nullptr) {
                 LOG(WARNING) << "Correcting symbol for " << d_symbol << " from " << std::hex
                              << reinterpret_cast<void*>(d_original) << " to "
                              << reinterpret_cast<void*>(symbol_addr);
@@ -128,6 +147,9 @@ struct SymbolHook
 void
 ensureAllHooksAreValid();
 
+void
+ensureArrowHooksAreValid();
+
 // Allocators are roughly ordered by the likelihood of any given allocation
 // record using them. The allocators assigned 1 to 7 have a more compact
 // representation within capture files. 0 must remain unassigned (it's used as
@@ -193,6 +215,17 @@ posix_memalign(void** memptr, size_t alignment, size_t size) noexcept;
 void*
 aligned_alloc(size_t alignment, size_t size) noexcept;
 
+#ifdef __linux__
+void*
+arrow_mi_malloc_aligned(size_t size, size_t alignment) noexcept;
+
+void*
+arrow_mi_realloc_aligned(void* ptr, size_t new_size, size_t alignment) noexcept;
+
+void
+arrow_mi_free(void* ptr) noexcept;
+#endif
+
 void*
 memalign(size_t alignment, size_t size) noexcept;
 

src/memray/_memray/tracking_api.cpp

@@ -872,6 +872,8 @@ Tracker::Tracker(
 #endif
     });
 
+    hooks::ensureArrowHooksAreValid();
+
     d_writer->setMainTidAndSkippedFrames(thread_id(), computeMainTidSkip());
     if (!d_writer->writeHeader(false)) {
         throw IoError{"Failed to write output header: " + std::string(strerror(errno))};
@@ -1235,6 +1237,8 @@ Tracker::trackObjectImpl(PyObject* obj, int event, const std::optional<NativeTra
 void
 Tracker::invalidate_module_cache_impl()
 {
+    // libarrow can be loaded after tracking starts.
+    hooks::ensureArrowHooksAreValid();
     d_patcher.overwrite_symbols();
     updateModuleCacheImpl();
 }

tests/integration/test_pyarrow.py

@@ -0,0 +1,68 @@
+import subprocess
+import sys
+
+import pytest
+
+from memray import AllocatorType
+from memray import FileReader
+
+pytestmark = pytest.mark.skipif(
+    sys.platform != "linux", reason="Arrow hooks use ELF interposition"
+)
+
+
+PYARROW_TEST = """
+import sys
+
+from memray import Tracker
+
+output, load_timing = sys.argv[1:]
+
+if load_timing == "before_tracking":
+    import pyarrow
+elif load_timing == "between_trackers":
+    with Tracker(output + ".initial"):
+        pass
+    import pyarrow
+
+with Tracker(output):
+    if load_timing == "while_tracking":
+        import pyarrow
+    assert pyarrow.default_memory_pool().backend_name == "mimalloc"
+    buffer = pyarrow.allocate_buffer(4096, resizable=True)
+    allocated = buffer.address
+    buffer.resize(8192)
+    reallocated = buffer.address
+    del buffer
+
+print(allocated, reallocated)
+"""
+
+
+@pytest.mark.parametrize(
+    "load_timing", ["before_tracking", "while_tracking", "between_trackers"]
+)
+def test_tracks_pyarrow_mimalloc_buffers(tmp_path, load_timing):
+    pytest.importorskip("pyarrow", minversion="25.0.0")
+    output = tmp_path / "test.bin"
+
+    result = subprocess.run(
+        [sys.executable, "-c", PYARROW_TEST, output, load_timing],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+
+    allocated, reallocated = map(int, result.stdout.split())
+    addresses = {allocated, reallocated}
+    records = [
+        (record.address, record.size, record.allocator)
+        for record in FileReader(output).get_allocation_records()
+        if record.address in addresses
+    ]
+    assert records == [
+        (allocated, 4096, AllocatorType.ALIGNED_ALLOC),
+        (allocated, 0, AllocatorType.FREE),
+        (reallocated, 8192, AllocatorType.REALLOC),
+        (reallocated, 0, AllocatorType.FREE),
+    ]