fix: Add back CustomMemoryResource with bug fix

velox/common/base/tests/StatsReporterTest.cpp

@@ -344,6 +344,10 @@ class TestMemoryPool : public memory::MemoryPool {
     return nullptr;
   }
 
+  memory::MemoryArbitrator* arbitrator() const override {
+    return nullptr;
+  }
+
   void enterArbitration() override {}
 
   void leaveArbitration() noexcept override {}

velox/common/memory/CMakeLists.txt

@@ -22,6 +22,7 @@ velox_add_library(
   ArbitrationOperation.cpp
   ArbitrationParticipant.cpp
   ByteStream.cpp
+  CustomMemoryResource.cpp
   HashStringAllocator.cpp
   MallocAllocator.cpp
   Memory.cpp
@@ -40,6 +41,7 @@ velox_add_library(
   ArbitrationParticipant.h
   ByteStream.h
   CompactDoubleList.h
+  CustomMemoryResource.h
   HashStringAllocator.h
   MallocAllocator.h
   Memory.h
@@ -69,3 +71,12 @@ velox_link_libraries(
     glog::glog
   PRIVATE velox_test_util re2::re2
 )
+
+velox_add_library(
+  velox_custom_memory_resource_registry
+  CustomMemoryResourceRegistry.cpp
+  HEADERS
+  CustomMemoryResourceRegistry.h
+)
+
+velox_link_libraries(velox_custom_memory_resource_registry velox_memory velox_scoped_registry)

velox/common/memory/CustomMemoryResource.cpp

@@ -0,0 +1,52 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed 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 "velox/common/memory/CustomMemoryResource.h"
+
+#include <utility>
+
+#include "velox/common/base/Exceptions.h"
+#include "velox/common/memory/MemoryArbitrator.h"
+
+namespace facebook::velox::memory {
+
+CustomMemoryResource::CustomMemoryResource(
+    std::string tag,
+    std::shared_ptr<MemoryAllocator> allocator,
+    std::shared_ptr<MemoryArbitrator> arbitrator,
+    ReclaimerFactory reclaimerFactory,
+    int64_t maxCapacity)
+    : tag_(std::move(tag)),
+      maxCapacity_(maxCapacity),
+      allocator_(std::move(allocator)),
+      arbitrator_(std::move(arbitrator)),
+      reclaimerFactory_(std::move(reclaimerFactory)) {
+  VELOX_USER_CHECK(!tag_.empty(), "CustomMemoryResource tag is empty");
+  VELOX_USER_CHECK_NOT_NULL(
+      allocator_, "CustomMemoryResource allocator is null for tag: {}", tag_);
+  VELOX_USER_CHECK_NOT_NULL(
+      arbitrator_, "CustomMemoryResource arbitrator is null for tag: {}", tag_);
+  VELOX_USER_CHECK(
+      reclaimerFactory_ != nullptr,
+      "CustomMemoryResource reclaimerFactory is null for tag: {}",
+      tag_);
+}
+
+std::unique_ptr<MemoryReclaimer> CustomMemoryResource::newReclaimer() const {
+  return reclaimerFactory_();
+}
+
+} // namespace facebook::velox::memory

velox/common/memory/CustomMemoryResource.h

@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed 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.
+ */
+
+#pragma once
+
+#include <cstdint>
+#include <functional>
+#include <limits>
+#include <memory>
+#include <string>
+
+namespace facebook::velox::memory {
+
+class MemoryAllocator;
+class MemoryArbitrator;
+class MemoryReclaimer;
+
+/// Describes an externally-provided memory resource (e.g. a GPU or tiered
+/// memory backend) registered with the memory subsystem and referenced by
+/// 'tag' when building per-query memory pools. The constructor enforces
+/// non-empty tag and non-null allocator, arbitrator, and reclaimerFactory;
+/// once constructed, the resource is immutable.
+class CustomMemoryResource {
+ public:
+  using ReclaimerFactory = std::function<std::unique_ptr<MemoryReclaimer>()>;
+
+  CustomMemoryResource(
+      std::string tag,
+      std::shared_ptr<MemoryAllocator> allocator,
+      std::shared_ptr<MemoryArbitrator> arbitrator,
+      ReclaimerFactory reclaimerFactory,
+      int64_t maxCapacity = std::numeric_limits<int64_t>::max());
+
+  /// Unique identifier for this resource.
+  const std::string& tag() const {
+    return tag_;
+  }
+
+  /// Capacity of the per-query root pool created from this resource.
+  int64_t maxCapacity() const {
+    return maxCapacity_;
+  }
+
+  /// Allocator backing pools tagged with this resource.
+  MemoryAllocator* allocator() const {
+    return allocator_.get();
+  }
+
+  /// Arbitrator routing capacity decisions for pools tagged with this
+  /// resource.
+  MemoryArbitrator* arbitrator() const {
+    return arbitrator_.get();
+  }
+
+  /// Returns a fresh reclaimer for a new pool by invoking the factory
+  /// supplied at construction.
+  std::unique_ptr<MemoryReclaimer> newReclaimer() const;
+
+ private:
+  const std::string tag_;
+  const int64_t maxCapacity_;
+  const std::shared_ptr<MemoryAllocator> allocator_;
+  const std::shared_ptr<MemoryArbitrator> arbitrator_;
+  const ReclaimerFactory reclaimerFactory_;
+};
+
+} // namespace facebook::velox::memory

velox/common/memory/CustomMemoryResourceRegistry.cpp

@@ -0,0 +1,33 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed 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 "velox/common/memory/CustomMemoryResourceRegistry.h"
+
+namespace facebook::velox::memory {
+
+// static
+CustomMemoryResourceRegistry::Registry& CustomMemoryResourceRegistry::global() {
+  static Registry instance;
+  return instance;
+}
+
+// static
+std::shared_ptr<CustomMemoryResourceRegistry::Registry>
+CustomMemoryResourceRegistry::createRegistry(Registry* parent) {
+  return std::make_shared<Registry>(parent);
+}
+
+} // namespace facebook::velox::memory

velox/common/memory/CustomMemoryResourceRegistry.h

@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed 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.
+ */
+
+#pragma once
+
+#include <memory>
+#include <string>
+#include <string_view>
+
+#include "velox/common/ScopedRegistry.h"
+#include "velox/common/memory/CustomMemoryResource.h"
+
+namespace facebook::velox::memory {
+
+/// Key under which a per-QueryCtx scoped CustomMemoryResourceRegistry is
+/// stored on QueryCtx via QueryCtx::setRegistry / QueryCtx::registry. Tasks
+/// use this key to look up resources when building the custom memory pool
+/// hierarchy.
+inline constexpr std::string_view kCustomMemoryResourceRegistryKey{
+    "customMemoryResource"};
+
+/// Entry point for the CustomMemoryResource registry. Provides the
+/// process-global root and a factory for scoped registries. Callers
+/// register, look up, and clear entries directly on the Registry instance
+/// (Registry::insert, Registry::find, Registry::clear). Registry methods
+/// are thread-safe.
+class CustomMemoryResourceRegistry {
+ public:
+  using Registry = ScopedRegistry<std::string, CustomMemoryResource>;
+
+  /// Process-global root registry. The reference is stable for the
+  /// lifetime of the process.
+  static Registry& global();
+
+  /// Creates a scoped registry. Defaults to inheriting from the global
+  /// registry; pass nullptr for isolation mode (no fallback).
+  static std::shared_ptr<Registry> createRegistry(Registry* parent = &global());
+};
+
+} // namespace facebook::velox::memory

velox/common/memory/Memory.cpp

@@ -274,7 +274,7 @@
       std::move(reclaimer),
       options);
   VELOX_CHECK_EQ(pool->capacity(), 0);
-  arbitrator_->addPool(pool);
+  pool->arbitrator()->addPool(pool);
   RECORD_HISTOGRAM_METRIC_VALUE(
       kMetricMemoryPoolInitialCapacityBytes, pool->capacity());
   return pool;
@@ -285,6 +285,36 @@
     int64_t maxCapacity,
     std::unique_ptr<MemoryReclaimer> reclaimer,
     const std::optional<MemoryPool::DebugOptions>& poolDebugOpts) {
+  return addRootPoolImpl(
+      name,
+      maxCapacity,
+      std::move(reclaimer),
+      poolDebugOpts,
+      /*customAllocator=*/nullptr,
+      /*customArbitrator=*/nullptr);
+}
+
+std::shared_ptr<MemoryPool> MemoryManager::addCustomRootPool(
+    const std::string& name,
+    std::shared_ptr<CustomMemoryResource> resource,
+    const std::optional<MemoryPool::DebugOptions>& poolDebugOpts) {
+  VELOX_USER_CHECK_NOT_NULL(resource);
+  return addRootPoolImpl(
+      name,
+      resource->maxCapacity(),
+      resource->newReclaimer(),
+      poolDebugOpts,
+      resource->allocator(),
+      resource->arbitrator());
+}
+
+std::shared_ptr<MemoryPool> MemoryManager::addRootPoolImpl(
+    const std::string& name,
+    int64_t maxCapacity,
+    std::unique_ptr<MemoryReclaimer> reclaimer,
+    const std::optional<MemoryPool::DebugOptions>& poolDebugOpts,
+    MemoryAllocator* customAllocator,
+    MemoryArbitrator* customArbitrator) {
   std::string poolName = name;
   if (poolName.empty()) {
     static std::atomic<int64_t> poolId{0};
@@ -298,6 +328,8 @@
   options.coreOnAllocationFailureEnabled = coreOnAllocationFailureEnabled_;
   options.getPreferredSize = getPreferredSize_;
   options.debugOptions = poolDebugOpts;
+  options.customAllocator = customAllocator;
+  options.customArbitrator = customArbitrator;
 
   auto pool = createRootPool(poolName, reclaimer, options);
   if (!disableMemoryPoolTracking_) {
@@ -308,7 +340,7 @@
       }
       pools_.emplace(poolName, pool);
     } catch (const VeloxRuntimeError&) {
-      arbitrator_->removePool(pool.get());
+      pool->arbitrator()->removePool(pool.get());
       throw;
     }
   }
@@ -340,7 +372,7 @@
 void MemoryManager::dropPool(MemoryPool* pool) {
   VELOX_CHECK_NOT_NULL(pool);
   VELOX_DCHECK_EQ(pool->reservedBytes(), 0);
-  arbitrator_->removePool(pool);
+  pool->arbitrator()->removePool(pool);
   if (disableMemoryPoolTracking_) {
     return;
   }

velox/common/memory/Memory.h

@@ -37,6 +37,7 @@
 #include "velox/common/base/CheckedArithmetic.h"
 #include "velox/common/base/SuccinctPrinter.h"
 #include "velox/common/memory/Allocation.h"
+#include "velox/common/memory/CustomMemoryResource.h"
 #include "velox/common/memory/MemoryAllocator.h"
 #include "velox/common/memory/MemoryPool.h"
 
@@ -223,6 +224,18 @@ class MemoryManager {
       const std::optional<MemoryPool::DebugOptions>& poolDebugOpts =
           std::nullopt);
 
+  /// Creates a root memory pool backed by 'resource'. The pool's capacity
+  /// comes from 'resource->maxCapacity'; its reclaimer comes from
+  /// 'resource->reclaimerFactory()'; its allocator and arbitrator are
+  /// borrowed from 'resource->allocator' and 'resource->arbitrator'. The
+  /// caller (typically via CustomMemoryResourceRegistry) is responsible
+  /// for keeping 'resource' alive while the pool exists.
+  std::shared_ptr<MemoryPool> addCustomRootPool(
+      const std::string& name,
+      std::shared_ptr<CustomMemoryResource> resource,
+      const std::optional<MemoryPool::DebugOptions>& poolDebugOpts =
+          std::nullopt);
+
   /// Creates a leaf memory pool for direct memory allocation use with specified
   /// 'name'. If 'name' is missing, the memory manager generates a default name
   /// internally to ensure uniqueness. The leaf memory pool is created as the
@@ -303,6 +316,16 @@ class MemoryManager {
       std::unique_ptr<MemoryReclaimer>& reclaimer,
       MemoryPool::Options& options);
 
+  // 'customAllocator' and 'customArbitrator' are borrowed pointers; if both
+  // are null, the manager's default tier is used.
+  std::shared_ptr<MemoryPool> addRootPoolImpl(
+      const std::string& name,
+      int64_t maxCapacity,
+      std::unique_ptr<MemoryReclaimer> reclaimer,
+      const std::optional<MemoryPool::DebugOptions>& poolDebugOpts,
+      MemoryAllocator* customAllocator,
+      MemoryArbitrator* customArbitrator);
+
   void dropPool(MemoryPool* pool);
 
   //  Returns the shared references to all the alive memory pools in 'pools_'.