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

pitrou · Apr 10, 2024 · 6a53c52 · 6a53c52
1 parent cd607d0
commit 6a53c52
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 21 deletions.
diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt
@@ -428,7 +428,15 @@ if(ARROW_JEMALLOC)
                               PROPERTIES SKIP_PRECOMPILE_HEADERS ON
                                          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_PRECOMPILE_HEADERS ON
+                                         SKIP_UNITY_BUILD_INCLUSION ON)
+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
@@ -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 {
 
@@ -394,15 +389,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) {
@@ -417,7 +412,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");
@@ -429,7 +424,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
@@ -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
@@ -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;
 
@@ -48,8 +45,16 @@ class JemallocAllocator {
 
 #endif  // defined(ARROW_JEMALLOC)
 
-}  // namespace internal
+}  // namespace arrow::memory_pool::internal
+
+#ifdef ARROW_MIMALLOC
+
+extern "C" {
 
-}  // namespace memory_pool
+ARROW_NOINLINE void* arrow_mi_malloc_aligned(size_t size, size_t alignment);
+ARROW_NOINLINE void* arrow_mi_realloc_aligned(void* p, size_t new_size, size_t alignment);
+ARROW_NOINLINE void arrow_mi_free(void* p);
+ARROW_NOINLINE void arrow_mi_collect(bool force);
+}
 
-}  // namespace arrow
+#endif  // defined(ARROW_MIMALLOC)
diff --git a/cpp/src/arrow/memory_pool_test.cc b/cpp/src/arrow/memory_pool_test.cc
@@ -77,7 +77,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
@@ -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