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,
@@ -208,7 +209,7 @@ std::shared_ptr<MemoryPool> MemoryManager::addRootPool(
       options);
   pools_.emplace(poolName, pool);
   VELOX_CHECK_EQ(pool->capacity(), 0);
-  arbitrator_->growCapacity(
+  arbitrator_->reserveCapacity(
       pool.get(), std::min<uint64_t>(poolInitCapacity_, maxCapacity));
   RECORD_HISTOGRAM_METRIC_VALUE(
       kMetricMemoryPoolInitialCapacityBytes, pool->capacity());
@@ -249,7 +250,7 @@ void MemoryManager::dropPool(MemoryPool* pool) {
   }
   pools_.erase(it);
   VELOX_DCHECK_EQ(pool->currentBytes(), 0);
-  arbitrator_->shrinkCapacity(pool, 0);
+  arbitrator_->releaseCapacity(pool);
 }
 
 MemoryPool& MemoryManager::deprecatedSharedLeafPool() {

velox/common/memory/Memory.h

@@ -145,18 +145,28 @@ struct MemoryManagerOptions {
   /// reservation capacity for system usage.
   int64_t arbitratorCapacity{kMaxMemory};
 
+  /// Memory capacity reserved to ensure that a query has minimal memory
+  /// capacity to run. This capacity should be less than 'arbitratorCapacity'.
+  /// A query's minimal memory capacity is defined by
+  /// 'memoryPoolReservedCapacity'.
+  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.
   /// Otherwise MemoryArbitrator::create returns a nullptr.
   std::string arbitratorKind{};
 
-  /// The initial memory capacity to reserve for a newly created memory pool.
+  /// The initial memory capacity to reserve for a newly created query memory
+  /// pool.
   uint64_t memoryPoolInitCapacity{256 << 20};
 
+  /// The minimal memory capacity reserved for a query memory pool to run.
+  uint64_t memoryPoolReservedCapacity{64 << 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 wait time has exceeded. If it is

velox/common/memory/MemoryArbitrator.cpp

@@ -93,7 +93,7 @@ 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 reserveCapacity(MemoryPool* pool, uint64_t /*unused*/) override {
     pool->grow(pool->maxCapacity());
     return pool->maxCapacity();
   }
@@ -109,7 +109,7 @@ class NoopArbitrator : public MemoryArbitrator {
 
   // Noop arbitrator has no memory capacity limit so no operation needed for
   // memory pool capacity release.
-  uint64_t shrinkCapacity(MemoryPool* pool, uint64_t targetBytes) override {
+  uint64_t releaseCapacity(MemoryPool* pool) override {
     // No-op
     return 0;
   }

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 memoryPoolReservedCapacity{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,13 @@ 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 reserveCapacity(MemoryPool* pool, uint64_t bytes) = 0;
+
+  /// Invoked by the memory manager to shrink up to 'targetBytes' free capacity
+  /// from a memory 'pool', and returns them back to the arbitrator. If
+  /// 'targetBytes' is zero, we shrink all the free capacity from the memory
+  /// pool. The function returns the actual freed capacity from 'pool'.
+  virtual uint64_t releaseCapacity(MemoryPool* pool) = 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
@@ -142,12 +153,6 @@ class MemoryArbitrator {
       const std::vector<std::shared_ptr<MemoryPool>>& candidatePools,
       uint64_t targetBytes) = 0;
 
-  /// Invoked by the memory manager to shrink up to 'targetBytes' free capacity
-  /// from a memory 'pool', and returns them back to the arbitrator. If
-  /// 'targetBytes' is zero, we shrink all the free capacity from the memory
-  /// pool. The function returns the actual freed capacity from 'pool'.
-  virtual uint64_t shrinkCapacity(MemoryPool* pool, uint64_t targetBytes) = 0;
-
   /// Invoked by the memory manager to shrink memory capacity from a given list
   /// of memory pools by reclaiming free and used memory. The freed memory
   /// capacity is given back to the arbitrator.  If 'targetBytes' is zero, then
@@ -239,12 +244,18 @@ class MemoryArbitrator {
  protected:
   explicit MemoryArbitrator(const Config& config)
       : capacity_(config.capacity),
+        reservedCapacity_(config.reservedCapacity),
+        memoryPoolReservedCapacity_(config.memoryPoolReservedCapacity),
         memoryPoolTransferCapacity_(config.memoryPoolTransferCapacity),
         memoryReclaimWaitMs_(config.memoryReclaimWaitMs),
         arbitrationStateCheckCb_(config.arbitrationStateCheckCb),
-        checkUsageLeak_(config.checkUsageLeak) {}
+        checkUsageLeak_(config.checkUsageLeak) {
+    VELOX_CHECK_LE(reservedCapacity_, capacity_);
+  }
 
   const uint64_t capacity_;
+  const uint64_t reservedCapacity_;
+  const uint64_t memoryPoolReservedCapacity_;
   const uint64_t memoryPoolTransferCapacity_;
   const uint64_t memoryReclaimWaitMs_;
   const MemoryArbitrationStateCheckCB arbitrationStateCheckCb_;