misc(memory) - Enable preferred memory size override for MemoryPool (#12712)

Summary:
Pull Request resolved: #12712

Enable overrides for getPreferredSize behavior on a per memory pool basis

Reviewed By: xiaoxmeng

Differential Revision: D71427371

fbshipit-source-id: ce4e1c9682de7d00af97ca7eb9a6323c4c8ffd02

Author: joeg

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,

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.