Add reserved memory capacity in arbitrator

velox/common/base/Counters.cpp

@@ -135,6 +135,10 @@ void registerVeloxMetrics() {
   DEFINE_METRIC(
       kMetricArbitratorFreeCapacityBytes, facebook::velox::StatType::AVG);
 
+  DEFINE_METRIC(
+      kMetricArbitratorFreeReservedCapacityBytes,
+      facebook::velox::StatType::AVG);
+
   // Tracks the leaf memory pool usage leak in bytes.
   DEFINE_METRIC(
       kMetricMemoryPoolUsageLeakBytes, facebook::velox::StatType::SUM);

velox/common/base/Counters.h

@@ -94,6 +94,9 @@ constexpr folly::StringPiece kMetricArbitratorArbitrationTimeMs{
 constexpr folly::StringPiece kMetricArbitratorFreeCapacityBytes{
     "velox.arbitrator_free_capacity_bytes"};
 
+constexpr folly::StringPiece kMetricArbitratorFreeReservedCapacityBytes{
+    "velox.arbitrator_free_reserved_capacity_bytes"};
+
 constexpr folly::StringPiece kMetricDriverYieldCount{
     "velox.driver_yield_count"};
 

velox/common/memory/Memory.cpp

@@ -61,6 +61,7 @@ std::unique_ptr<MemoryArbitrator> createArbitrator(
       {.kind = options.arbitratorKind,
        .capacity =
            std::min(options.arbitratorCapacity, options.allocatorCapacity),
+       .reservedCapacity = options.arbitratorReservedCapacity,
        .memoryPoolTransferCapacity = options.memoryPoolTransferCapacity,
        .memoryReclaimWaitMs = options.memoryReclaimWaitMs,
        .arbitrationStateCheckCb = options.arbitrationStateCheckCb,
@@ -209,9 +210,15 @@ std::shared_ptr<MemoryPool> MemoryManager::addRootPool(
   pools_.emplace(poolName, pool);
   VELOX_CHECK_EQ(pool->capacity(), 0);
   arbitrator_->growCapacity(
-      pool.get(), std::min<uint64_t>(poolInitCapacity_, maxCapacity));
+      pool.get(),
+      //      std::min<uint64_t>b(poolInitCapacity_, maxCapacity),
+      64ULL << 20,
+      /*useReserve*/ true);
+  const int capacity = pool->capacity();
+  VELOX_MEM_LOG(INFO) << "memory pool " << pool->name() << " initial capacity "
+                      << succinctBytes(capacity);
   RECORD_HISTOGRAM_METRIC_VALUE(
-      kMetricMemoryPoolInitialCapacityBytes, pool->capacity());
+      kMetricMemoryPoolInitialCapacityBytes, capacity);
   return pool;
 }
 

velox/common/memory/Memory.h

@@ -145,6 +145,8 @@ struct MemoryManagerOptions {
   /// reservation capacity for system usage.
   int64_t arbitratorCapacity{kMaxMemory};
 
+  int64_t arbitratorReservedCapacity{0};
+
   /// The string kind of memory arbitrator used in the memory manager.
   ///
   /// NOTE: the arbitrator will only be created if its kind is set explicitly.

velox/common/memory/MemoryArbitrator.cpp

@@ -93,7 +93,8 @@ class NoopArbitrator : public MemoryArbitrator {
 
   // Noop arbitrator has no memory capacity limit so no operation needed for
   // memory pool capacity reserve.
-  uint64_t growCapacity(MemoryPool* pool, uint64_t /*unused*/) override {
+  uint64_t growCapacity(MemoryPool* pool, uint64_t /*unused*/, bool /*unused*/)
+      override {
     pool->grow(pool->maxCapacity());
     return pool->maxCapacity();
   }
@@ -199,7 +200,10 @@ bool MemoryReclaimer::reclaimableBytes(
   if (pool.kind() == MemoryPool::Kind::kLeaf) {
     return false;
   }
-  bool reclaimable{false};
+  if (pool.capacity() <= pool.minCapacity_) {
+    return false;
+  }
+   bool reclaimable{false};
   pool.visitChildren([&](MemoryPool* pool) {
     auto reclaimableBytesOpt = pool->reclaimableBytes();
     reclaimable |= reclaimableBytesOpt.has_value();
@@ -218,6 +222,9 @@ uint64_t MemoryReclaimer::reclaim(
   if (pool->kind() == MemoryPool::Kind::kLeaf) {
     return 0;
   }
+  if (pool->capacity() <= pool->minCapacity_) {
+    return 0;
+  }
 
   // Sort the child pools based on their reserved memory and reclaim from the
   // child pool with most reservation first.

velox/common/memory/MemoryArbitrator.h

@@ -57,9 +57,14 @@ class MemoryArbitrator {
     /// manager.
     int64_t capacity;
 
+    int64_t reservedCapacity{6LL << 30};
+
+    /// The min reserved capacity for each query.
+    uint64_t minReservedCapacity{64UL << 20};
+
     /// The minimal memory capacity to transfer out of or into a memory pool
     /// during the memory arbitration.
-    uint64_t memoryPoolTransferCapacity{32 << 20};
+    uint64_t memoryPoolTransferCapacity{128 << 20};
 
     /// Specifies the max time to wait for memory reclaim by arbitration. The
     /// memory reclaim might fail if the max time has exceeded. This prevents
@@ -123,7 +128,8 @@ class MemoryArbitrator {
   /// grow the memory pool's capacity based on the free available memory
   /// capacity in the arbitrator, and returns the actual growed capacity in
   /// bytes.
-  virtual uint64_t growCapacity(MemoryPool* pool, uint64_t bytes) = 0;
+  virtual uint64_t
+  growCapacity(MemoryPool* pool, uint64_t bytes, bool useReserve) = 0;
 
   /// Invoked by the memory manager to grow a memory pool's capacity.
   /// 'pool' is the memory pool to request to grow. 'candidates' is a list
@@ -239,12 +245,19 @@ class MemoryArbitrator {
  protected:
   explicit MemoryArbitrator(const Config& config)
       : capacity_(config.capacity),
+        minReservedCapacity_(config.minReservedCapacity),
+        reservedCapacity_(config.reservedCapacity),
         memoryPoolTransferCapacity_(config.memoryPoolTransferCapacity),
         memoryReclaimWaitMs_(config.memoryReclaimWaitMs),
         arbitrationStateCheckCb_(config.arbitrationStateCheckCb),
-        checkUsageLeak_(config.checkUsageLeak) {}
+        checkUsageLeak_(config.checkUsageLeak) {
+    VELOX_CHECK_LE(reservedCapacity_, capacity_);
+    VELOX_CHECK_EQ(reservedCapacity_ % minReservedCapacity_, 0);
+  }
 
   const uint64_t capacity_;
+  const uint64_t minReservedCapacity_;
+  const uint64_t reservedCapacity_;
   const uint64_t memoryPoolTransferCapacity_;
   const uint64_t memoryReclaimWaitMs_;
   const MemoryArbitrationStateCheckCB arbitrationStateCheckCb_;