GH-50083: [C++] Access mimalloc through dynamically-resolved symbols (#41128)

### Rationale for this change

The [memray memory profiler](https://github.com/bloomberg/memray) works by interposing certain dynamic symbols in the profiled process to replace them with their own functions that will collect memory allocation data. It will currently, to the best of my knowledge, only recognize system C calls such `malloc`, `mmap`...

When a third-party allocator like mimalloc or jemalloc is being used, such that Arrow does by default, memray does not see the logical allocation calls made through these allocator's APIs (because they are not interposed), but only the raw memory reservations that they issue using system routines.

This can lead people using memray to think that a given Arrow workload (or any workload using such allocators, really) that an inordinate amount of memory is being used, while the reported memory mostly represents non-committed virtual memory that the allocator keeps for performance reasons. Concrete example in GH-40301: we allocate a number of 1kiB buffers from mimalloc, but memray sees a similar number of 64MiB calls to `mmap`.

We [discussed](bloomberg/memray#577) how to enhance memray such as to account for the corresponding logical allocations, and we came to the conclusion that it requires that Arrow exposes API calls that can be dynamically interposed. Since we typically build against a static `libmimalloc.a`, the mimalloc symbols cannot be exposed (at least, I cannot seem to get this to work on Ubuntu). This means we need to define our own symbols wrapping the mimalloc APIs.

### What changes are included in this PR?

Define public, interposable symbols that redirect into the mimalloc APIs that we use.

### Are these changes tested?

Not for now. We could probably test them, at least on Linux, by compiling an almost trivial shared library and interposing it using `LD_PRELOAD`.

### Are there any user-facing changes?

No. There should not be any noticeable performance regression, except perhaps on memory pool micro-benchmarks.

* GitHub Issue: #50083

Authored-by: Antoine Pitrou <antoine@python.org>
Signed-off-by: Antoine Pitrou <antoine@python.org>

Author: pitrou

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