From fe2a4e3a0cf4b3337c220ba45030e656589d7aaf Mon Sep 17 00:00:00 2001 From: Antoine Pitrou Date: Tue, 2 Jun 2026 14:45:41 +0200 Subject: [PATCH] [C++] Access mimalloc through dynamically-resolved symbols --- cpp/src/arrow/CMakeLists.txt | 17 ++++++++++++ cpp/src/arrow/memory_pool.cc | 19 +++++-------- cpp/src/arrow/memory_pool_benchmark.cc | 1 + cpp/src/arrow/memory_pool_internal.h | 26 ++++++++++++------ cpp/src/arrow/memory_pool_mimalloc.cc | 38 ++++++++++++++++++++++++++ cpp/src/arrow/memory_pool_test.cc | 5 +++- cpp/src/arrow/memory_pool_test.h | 8 ++++++ 7 files changed, 93 insertions(+), 21 deletions(-) create mode 100644 cpp/src/arrow/memory_pool_mimalloc.cc diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index 8d3cf9682afd..453a06f9a8c3 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/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) diff --git a/cpp/src/arrow/memory_pool.cc b/cpp/src/arrow/memory_pool.cc index 1c77a60ba0e2..0b4843bec350 100644 --- a/cpp/src/arrow/memory_pool.cc +++ b/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( - mi_malloc_aligned(static_cast(size), static_cast(alignment))); + *out = reinterpret_cast(arrow_mi_malloc_aligned( + static_cast(size), static_cast(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( - mi_realloc_aligned(previous_ptr, static_cast(new_size), alignment)); + arrow_mi_realloc_aligned(previous_ptr, static_cast(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); } } diff --git a/cpp/src/arrow/memory_pool_benchmark.cc b/cpp/src/arrow/memory_pool_benchmark.cc index c2e55314b56f..6426c43100b5 100644 --- a/cpp/src/arrow/memory_pool_benchmark.cc +++ b/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" diff --git a/cpp/src/arrow/memory_pool_internal.h b/cpp/src/arrow/memory_pool_internal.h index dc90d5680b66..4dfe30168022 100644 --- a/cpp/src/arrow/memory_pool_internal.h +++ b/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) diff --git a/cpp/src/arrow/memory_pool_mimalloc.cc b/cpp/src/arrow/memory_pool_mimalloc.cc new file mode 100644 index 000000000000..05f67d6e84bf --- /dev/null +++ b/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 + +// 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); } + +// } diff --git a/cpp/src/arrow/memory_pool_test.cc b/cpp/src/arrow/memory_pool_test.cc index 0af1ed2d9eca..1511f90ac7dc 100644 --- a/cpp/src/arrow/memory_pool_test.cc +++ b/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); diff --git a/cpp/src/arrow/memory_pool_test.h b/cpp/src/arrow/memory_pool_test.h index 32f1cc5d1d31..87f4f2a152ec 100644 --- a/cpp/src/arrow/memory_pool_test.h +++ b/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