Track Apache Arrow mimalloc allocations via arrow_mi_* hooks

Author: pablogsal

src/memray/_memray/hooks.cpp

@@ -418,6 +418,59 @@ aligned_alloc(size_t alignment, size_t size) noexcept
     return ret;
 }
 
+// Apache Arrow mimalloc wrapper hooks (apache/arrow#41128). These wrap
+// mimalloc allocations that are otherwise opaque to memray because mimalloc
+// is statically linked into libarrow and serves allocations out of mmap
+// arenas. Tracked under the existing aligned-alloc / realloc / free families.
+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);
+    }
+}
+
 #if defined(__linux__)
 
 void*

src/memray/_memray/hooks.h

@@ -24,6 +24,16 @@
 #include "compat.h"
 #include "logging.h"
 
+// Apache Arrow mimalloc wrapper symbols (apache/arrow#41128). Declared weak
+// so memray links even when the host binary doesn't depend on libarrow.
+// At runtime memray's phdr scanner resolves the real address from libarrow.so.
+extern "C" {
+__attribute__((weak)) void* arrow_mi_malloc_aligned(size_t size, size_t alignment);
+__attribute__((weak)) void* arrow_mi_realloc_aligned(void* ptr, size_t new_size,
+                                                    size_t alignment);
+__attribute__((weak)) void arrow_mi_free(void* ptr);
+}
+
 #if defined(__APPLE__)
 #    define MEMRAY_PLATFORM_HOOKED_FUNCTIONS
 #elif defined(__GLIBC__)
@@ -51,6 +61,9 @@
     FOR_EACH_HOOKED_FUNCTION(dlopen)                                                                    \
     FOR_EACH_HOOKED_FUNCTION(dlclose)                                                                   \
     FOR_EACH_HOOKED_FUNCTION(PyGILState_Ensure)                                                         \
+    FOR_EACH_HOOKED_FUNCTION(arrow_mi_malloc_aligned)                                                   \
+    FOR_EACH_HOOKED_FUNCTION(arrow_mi_realloc_aligned)                                                  \
+    FOR_EACH_HOOKED_FUNCTION(arrow_mi_free)                                                             \
     MEMRAY_PLATFORM_HOOKED_FUNCTIONS
 
 namespace memray::hooks {
@@ -179,6 +192,15 @@ posix_memalign(void** memptr, size_t alignment, size_t size) noexcept;
 void*
 aligned_alloc(size_t alignment, size_t size) noexcept;
 
+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;
+
 void*
 memalign(size_t alignment, size_t size) noexcept;
 

src/memray/_memray/tracking_api.cpp

@@ -898,6 +898,11 @@ Tracker::~Tracker()
 {
     RecursionGuard guard;
     tracking_api::Tracker::deactivate();
+#ifdef MEMRAY_HAS_GHOST_STACK
+    if (d_fast_unwind) {
+        ghost_stack_dump_stats();
+    }
+#endif
 
     PythonStackTracker::s_native_tracking_enabled = false;
     d_background_thread->stop();
@@ -1233,6 +1238,12 @@ Tracker::trackObjectImpl(PyObject* obj, int event, const std::optional<NativeTra
 void
 Tracker::invalidate_module_cache_impl()
 {
+    // Re-resolve hook originals before patching: hooked symbols defined in
+    // libraries loaded at runtime (e.g. libarrow.so's arrow_mi_*) are not
+    // present when ensureAllHooksAreValid() runs at tracker init; without
+    // this call, the patcher would rewrite GOT entries to point at our
+    // intercepts, which would then call a null d_original and crash.
+    hooks::ensureAllHooksAreValid();
     d_patcher.overwrite_symbols();
     updateModuleCacheImpl();
 }