GH-50083: [C++] Access mimalloc through dynamically-resolved symbols

cpp/src/arrow/CMakeLists.txt

@@ -459,7 +459,24 @@ if(ARROW_JEMALLOC)
   set_source_files_properties(memory_pool_jemalloc.cc
                               PROPERTIES SKIP_UNITY_BUILD_INCLUSION ON)
 endif()
+if(ARROW_MIMALLOC)
+  list(APPEND ARROW_MEMORY_POOL_SRCS memory_pool_mimalloc.cc)
+  set_source_files_properties(memory_pool_mimalloc.cc
+                              PROPERTIES SKIP_UNITY_BUILD_INCLUSION ON)
+  # GH-50083: make sure some allocation-related public symbols are interposable,
+  # for the benefit of memory profilers and other runtime analyzers.
+  # Ideally we would scope a specific subset of symbols, but it doesn't seem
+  # easily doable. See `memory_pool_internal.h`.
+  if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU"
+     OR CMAKE_CXX_COMPILER_ID STREQUAL "AppleClang"
+     OR CMAKE_CXX_COMPILER_ID STREQUAL "Clang")
+    set_source_files_properties(${ARROW_MEMORY_POOL_SRCS}
+                                PROPERTIES COMPILE_OPTIONS "-fsemantic-interposition")
+  endif()
+endif()
+
 arrow_add_object_library(ARROW_MEMORY_POOL ${ARROW_MEMORY_POOL_SRCS})
+
 if(ARROW_JEMALLOC)
   foreach(ARROW_MEMORY_POOL_TARGET ${ARROW_MEMORY_POOL_TARGETS})
     target_link_libraries(${ARROW_MEMORY_POOL_TARGET} PRIVATE jemalloc::jemalloc)

cpp/src/arrow/memory_pool.cc

@@ -54,16 +54,11 @@
 #endif
 
 namespace arrow {
-
-namespace memory_pool {
-
-namespace internal {
+namespace memory_pool::internal {
 
 alignas(kDefaultBufferAlignment) int64_t zero_size_area[1] = {kDebugXorSuffix};
 
-}  // namespace internal
-
-}  // namespace memory_pool
+}  // namespace memory_pool::internal
 
 namespace {
 
@@ -400,15 +395,15 @@ class MimallocAllocator {
       *out = memory_pool::internal::kZeroSizeArea;
       return Status::OK();
     }
-    *out = reinterpret_cast<uint8_t*>(
-        mi_malloc_aligned(static_cast<size_t>(size), static_cast<size_t>(alignment)));
+    *out = reinterpret_cast<uint8_t*>(arrow_mi_malloc_aligned(
+        static_cast<size_t>(size), static_cast<size_t>(alignment)));
     if (*out == NULL) {
       return Status::OutOfMemory("malloc of size ", size, " failed");
     }
     return Status::OK();
   }
 
-  static void ReleaseUnused() { mi_collect(true); }
+  static void ReleaseUnused() { arrow_mi_collect(true); }
 
   static Status ReallocateAligned(int64_t old_size, int64_t new_size, int64_t alignment,
                                   uint8_t** ptr) {
@@ -423,7 +418,7 @@ class MimallocAllocator {
       return Status::OK();
     }
     *ptr = reinterpret_cast<uint8_t*>(
-        mi_realloc_aligned(previous_ptr, static_cast<size_t>(new_size), alignment));
+        arrow_mi_realloc_aligned(previous_ptr, static_cast<size_t>(new_size), alignment));
     if (*ptr == NULL) {
       *ptr = previous_ptr;
       return Status::OutOfMemory("realloc of size ", new_size, " failed");
@@ -435,7 +430,7 @@ class MimallocAllocator {
     if (ptr == memory_pool::internal::kZeroSizeArea) {
       DCHECK_EQ(size, 0);
     } else {
-      mi_free(ptr);
+      arrow_mi_free(ptr);
     }
   }
 

cpp/src/arrow/memory_pool_benchmark.cc

@@ -17,6 +17,7 @@
 
 #include "arrow/memory_pool.h"
 #include "arrow/result.h"
+#include "arrow/util/config.h"
 #include "arrow/util/logging.h"
 
 #include "benchmark/benchmark.h"

cpp/src/arrow/memory_pool_internal.h

@@ -19,12 +19,9 @@
 
 #include "arrow/memory_pool.h"
 #include "arrow/util/config.h"
+#include "arrow/util/macros.h"
 
-namespace arrow {
-
-namespace memory_pool {
-
-namespace internal {
+namespace arrow::memory_pool::internal {
 
 static constexpr int64_t kDebugXorSuffix = -0x181fe80e0b464188LL;
 
@@ -49,8 +46,21 @@ class JemallocAllocator {
 
 #endif  // defined(ARROW_JEMALLOC)
 
-}  // namespace internal
+}  // namespace arrow::memory_pool::internal
+
+#ifdef ARROW_MIMALLOC
 
-}  // namespace memory_pool
+extern "C" {
+// GH-50083: expose public symbols with well-known names for memory profilers
+// to be able to intercept our mimalloc calls and better make sense of Arrow's
+// memory profile.
+// These symbols need to be interposable (using e.g. `LD_PRELOAD`), which is
+// ensured using `-fsemantic-interposition` in CMakeLists.txt.
+ARROW_EXPORT ARROW_NOINLINE void* arrow_mi_malloc_aligned(size_t size, size_t alignment);
+ARROW_EXPORT ARROW_NOINLINE void* arrow_mi_realloc_aligned(void* p, size_t new_size,
+                                                           size_t alignment);
+ARROW_EXPORT ARROW_NOINLINE void arrow_mi_free(void* p);
+ARROW_EXPORT ARROW_NOINLINE void arrow_mi_collect(bool force);
+}
 
-}  // namespace arrow
+#endif  // defined(ARROW_MIMALLOC)

cpp/src/arrow/memory_pool_mimalloc.cc

@@ -0,0 +1,38 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/memory_pool_internal.h"
+#include "arrow/util/io_util.h"
+#include "arrow/util/logging.h"  // IWYU pragma: keep
+
+#include <mimalloc.h>
+
+// extern "C" {
+
+void arrow_mi_free(void* p) { mi_free(p); }
+
+void* arrow_mi_malloc_aligned(size_t size, size_t alignment) {
+  return mi_malloc_aligned(size, alignment);
+}
+
+void* arrow_mi_realloc_aligned(void* p, size_t new_size, size_t alignment) {
+  return mi_realloc_aligned(p, new_size, alignment);
+}
+
+void arrow_mi_collect(bool force) { mi_collect(force); }
+
+// }

cpp/src/arrow/memory_pool_test.cc

@@ -78,7 +78,10 @@ TYPED_TEST_P(TestMemoryPool, Reallocate) { this->TestReallocate(); }
 
 TYPED_TEST_P(TestMemoryPool, Alignment) { this->TestAlignment(); }
 
-REGISTER_TYPED_TEST_SUITE_P(TestMemoryPool, MemoryTracking, OOM, Reallocate, Alignment);
+TYPED_TEST_P(TestMemoryPool, ReleaseUnused) { this->TestReleaseUnused(); }
+
+REGISTER_TYPED_TEST_SUITE_P(TestMemoryPool, MemoryTracking, OOM, Reallocate, Alignment,
+                            ReleaseUnused);
 
 INSTANTIATE_TYPED_TEST_SUITE_P(Default, TestMemoryPool, DefaultMemoryPoolFactory);
 INSTANTIATE_TYPED_TEST_SUITE_P(System, TestMemoryPool, SystemMemoryPoolFactory);

cpp/src/arrow/memory_pool_test.h

@@ -106,6 +106,14 @@ class TestMemoryPoolBase : public ::testing::Test {
       pool->Free(data512, 10, 512);
     }
   }
+
+  void TestReleaseUnused() {
+    auto pool = memory_pool();
+    const int64_t nbytes = pool->bytes_allocated();
+    pool->ReleaseUnused();
+    // Unfortunately there's not much that we can assert here
+    ASSERT_EQ(nbytes, pool->bytes_allocated());
+  }
 };
 
 }  // namespace arrow