From bc177ba1213bd7c02bf19045f6c6e49ecc208b69 Mon Sep 17 00:00:00 2001
From: Joe Giardino <joegiard@meta.com>
Date: Wed, 19 Mar 2025 09:43:41 -0700
Subject: [PATCH] misc(memory) - Enable preferred memory size override for
 MemoryPool (#12712)

Summary:

Enable overrides for getPreferredSize behavior on a per memory pool basis

Reviewed By: xiaoxmeng

Differential Revision: D71427371
---
 velox/common/memory/MemoryPool.cpp | 6 ++++++
 velox/common/memory/MemoryPool.h   | 5 ++++-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/velox/common/memory/MemoryPool.cpp b/velox/common/memory/MemoryPool.cpp
index 7d5c0816b42d..9a833bcf1dc0 100644
--- a/velox/common/memory/MemoryPool.cpp
+++ b/velox/common/memory/MemoryPool.cpp
@@ -426,6 +426,12 @@ size_t MemoryPool::getPreferredSize(size_t size) {
   return lower * 2;
 }
 
+void MemoryPool::setPreferredSize(
+    std::function<size_t(size_t)> getPreferredSizeFunc) {
+  VELOX_CHECK_NOT_NULL(getPreferredSizeFunc);
+  getPreferredSize_ = getPreferredSizeFunc;
+}
+
 MemoryPoolImpl::MemoryPoolImpl(
     MemoryManager* memoryManager,
     const std::string& name,
diff --git a/velox/common/memory/MemoryPool.h b/velox/common/memory/MemoryPool.h
index 699f3c91547c..77cd1110c9c3 100644
--- a/velox/common/memory/MemoryPool.h
+++ b/velox/common/memory/MemoryPool.h
@@ -473,6 +473,9 @@ class MemoryPool : public std::enable_shared_from_this<MemoryPool> {
     return bits::roundUp(size, 8 * kMB);
   }
 
+  // Overrides getPreferredSize to allow specializing behavior for this pool.
+  void setPreferredSize(std::function<size_t(size_t)> getPreferredSizeFunc);
+
  protected:
   static constexpr uint64_t kMB = 1 << 20;
 
@@ -530,7 +533,7 @@ class MemoryPool : public std::enable_shared_from_this<MemoryPool> {
   const bool threadSafe_;
   const bool debugEnabled_;
   const bool coreOnAllocationFailureEnabled_;
-  const std::function<size_t(size_t)> getPreferredSize_;
+  std::function<size_t(size_t)> getPreferredSize_;
 
   /// Indicates if the memory pool has been aborted by the memory arbitrator or
   /// not.