Support single thread execution on MemoryPool::reclaim (#362)

* Remove high usage callback

* Executor not found when resuming a single-thread-execution task during memory reclaiming procedure

* Use memory arbitrator to manage pool shrinking as well as growing

* Allow using customized memory arbitrator when creating memory manager instance

* Deadlock if throwing exception from MemoryManager::addRootPool(...)

* Assertion failed, expected "!spillFinalized_" if reclaiming HashBuild operator during finishing

* UT Fix: MemoryManagerTest.Ctor

* Assertion failed, expected "spillPartitionEntry.second->numFiles() > 0" after HashBuild::reclaim() was called on a HashBuild that has no inputs

Author: zhztheplayer

velox/common/memory/Memory.cpp

@@ -28,15 +28,7 @@ constexpr folly::StringPiece kDefaultLeafName("__default_leaf__");
 MemoryManager::MemoryManager(const Options& options)
     : capacity_{options.capacity},
       allocator_{options.allocator->shared_from_this()},
-      arbitrator_(MemoryArbitrator::create(MemoryArbitrator::Config{
-          .kind = options.arbitratorConfig.kind,
-          .capacity = capacity_,
-          .initMemoryPoolCapacity =
-              options.arbitratorConfig.initMemoryPoolCapacity,
-          .minMemoryPoolCapacityTransferSize =
-              options.arbitratorConfig.minMemoryPoolCapacityTransferSize,
-          .retryArbitrationFailure =
-              options.arbitratorConfig.retryArbitrationFailure})),
+      arbitrator_{options.arbitratorFactory()},
       alignment_(std::max(MemoryAllocator::kMinAlignment, options.alignment)),
       checkUsageLeak_(options.checkUsageLeak),
       debugEnabled_(options.debugEnabled),
@@ -59,6 +51,9 @@ MemoryManager::MemoryManager(const Options& options)
               .debugEnabled = options.debugEnabled})} {
   VELOX_CHECK_NOT_NULL(allocator_);
   VELOX_USER_CHECK_GE(capacity_, 0);
+  if (arbitrator_ != nullptr) {
+    VELOX_CHECK_EQ(arbitrator_->capacity(), capacity_);
+  }
   MemoryAllocator::alignmentCheck(0, alignment_);
   defaultRoot_->grow(defaultRoot_->maxCapacity());
   const size_t numSharedPools =
@@ -119,19 +114,22 @@ std::shared_ptr<MemoryPool> MemoryManager::addRootPool(
   options.checkUsageLeak = checkUsageLeak_;
   options.debugEnabled = debugEnabled_;
 
-  folly::SharedMutex::WriteHolder guard{mutex_};
-  if (pools_.find(poolName) != pools_.end()) {
-    VELOX_FAIL("Duplicate root pool name found: {}", poolName);
+  std::shared_ptr<MemoryPool> pool;
+  {
+    folly::SharedMutex::WriteHolder guard{mutex_};
+    if (pools_.find(poolName) != pools_.end()) {
+      VELOX_FAIL("Duplicate root pool name found: {}", poolName);
+    }
+    pool = std::make_shared<MemoryPoolImpl>(
+        this,
+        poolName,
+        MemoryPool::Kind::kAggregate,
+        nullptr,
+        std::move(reclaimer),
+        poolDestructionCb_,
+        options);
+    pools_.emplace(poolName, pool);
   }
-  auto pool = std::make_shared<MemoryPoolImpl>(
-      this,
-      poolName,
-      MemoryPool::Kind::kAggregate,
-      nullptr,
-      std::move(reclaimer),
-      poolDestructionCb_,
-      options);
-  pools_.emplace(poolName, pool);
   VELOX_CHECK_EQ(pool->capacity(), 0);
   if (arbitrator_ != nullptr) {
     arbitrator_->reserveMemory(pool.get(), capacity);
@@ -154,6 +152,14 @@ std::shared_ptr<MemoryPool> MemoryManager::addLeafPool(
   return defaultRoot_->addLeafChild(poolName, threadSafe, nullptr);
 }
 
+uint64_t MemoryManager::shrinkPool(MemoryPool* pool, uint64_t decrementBytes) {
+  VELOX_CHECK_NOT_NULL(pool);
+  if (arbitrator_ == nullptr) {
+    return pool->shrink(decrementBytes);
+  }
+  return arbitrator_->releaseMemory(pool, decrementBytes);
+}
+
 bool MemoryManager::growPool(MemoryPool* pool, uint64_t incrementBytes) {
   VELOX_CHECK_NOT_NULL(pool);
   VELOX_CHECK_NE(pool->capacity(), kMaxMemory);
@@ -168,14 +174,16 @@ bool MemoryManager::growPool(MemoryPool* pool, uint64_t incrementBytes) {
 
 void MemoryManager::dropPool(MemoryPool* pool) {
   VELOX_CHECK_NOT_NULL(pool);
-  folly::SharedMutex::WriteHolder guard{mutex_};
-  auto it = pools_.find(pool->name());
-  if (it == pools_.end()) {
-    VELOX_FAIL("The dropped memory pool {} not found", pool->name());
+  {
+    folly::SharedMutex::WriteHolder guard{mutex_};
+    auto it = pools_.find(pool->name());
+    if (it == pools_.end()) {
+      VELOX_FAIL("The dropped memory pool {} not found", pool->name());
+    }
+    pools_.erase(it);
   }
-  pools_.erase(it);
   if (arbitrator_ != nullptr) {
-    arbitrator_->releaseMemory(pool);
+    arbitrator_->releaseMemory(pool, 0);
   }
 }
 

velox/common/memory/Memory.h

@@ -86,8 +86,9 @@ class IMemoryManager {
     /// Specifies the backing memory allocator.
     MemoryAllocator* allocator{MemoryAllocator::getInstance()};
 
-    /// Specifies the memory arbitration config.
-    MemoryArbitrator::Config arbitratorConfig{};
+    /// Specifies the memory arbitrator
+    std::function<std::unique_ptr<MemoryArbitrator>()> arbitratorFactory{
+        []() { return MemoryArbitrator::create({}); }};
   };
 
   virtual ~IMemoryManager() = default;
@@ -118,6 +119,10 @@ class IMemoryManager {
       const std::string& name = "",
       bool threadSafe = true) = 0;
 
+  /// Invoked to shrink a memory pool's free capacity with up to
+  /// 'decrementBytes'.
+  virtual uint64_t shrinkPool(MemoryPool* pool, uint64_t decrementBytes) = 0;
+
   /// Invoked to grows a memory pool's free capacity with at least
   /// 'incrementBytes'. The function returns true on success, otherwise false.
   virtual bool growPool(MemoryPool* pool, uint64_t incrementBytes) = 0;
@@ -196,6 +201,7 @@ class MemoryManager final : public IMemoryManager {
       const std::string& name = "",
       bool threadSafe = true) final;
 
+  uint64_t shrinkPool(MemoryPool* pool, uint64_t decrementBytes) final;
   bool growPool(MemoryPool* pool, uint64_t incrementBytes) final;
 
   MemoryPool& deprecatedSharedLeafPool() final;

velox/common/memory/MemoryArbitrator.cpp

@@ -50,6 +50,10 @@ std::unique_ptr<MemoryArbitrator> MemoryArbitrator::create(
   }
 }
 
+uint64_t MemoryArbitrator::capacity() {
+  return capacity_;
+}
+
 std::unique_ptr<MemoryReclaimer> MemoryReclaimer::create() {
   return std::unique_ptr<MemoryReclaimer>(new MemoryReclaimer());
 }

velox/common/memory/MemoryArbitrator.h

@@ -105,9 +105,10 @@ class MemoryArbitrator {
   /// the memory arbitration on demand when actual memory allocation happens.
   virtual void reserveMemory(MemoryPool* pool, uint64_t bytes) = 0;
 
-  /// Invoked by the memory manager to return back all the reserved memory
-  /// capacity of a destroying memory pool.
-  virtual void releaseMemory(MemoryPool* pool) = 0;
+  /// Invoked by the memory manager to return back the specified amount of
+  /// reserved memory capacity of a destroying memory pool. If 0 is specified,
+  /// release all reserve memory. Returns the actually released amount of bytes.
+  virtual uint64_t releaseMemory(MemoryPool* pool, uint64_t bytes) = 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
@@ -126,6 +127,7 @@ class MemoryArbitrator {
       const std::vector<std::shared_ptr<MemoryPool>>& candidatePools,
       uint64_t targetBytes) = 0;
 
+  uint64_t capacity();
   /// The internal execution stats of the memory arbitrator.
   struct Stats {
     /// The number of arbitration requests.

velox/common/memory/MemoryPool.cpp

@@ -577,18 +577,6 @@ int64_t MemoryPoolImpl::capacity() const {
   return capacity_;
 }
 
-bool MemoryPoolImpl::highUsage() {
-  if (parent_ != nullptr) {
-    return parent_->highUsage();
-  }
-
-  if (highUsageCallback_ != nullptr) {
-    return highUsageCallback_(*this);
-  }
-
-  return false;
-}
-
 std::shared_ptr<MemoryPool> MemoryPoolImpl::genChild(
     std::shared_ptr<MemoryPool> parent,
     const std::string& name,
@@ -743,6 +731,16 @@ bool MemoryPoolImpl::incrementReservationThreadSafe(
           capacityToString(manager_->capacity()))));
 }
 
+uint64_t MemoryPoolImpl::shrinkManaged(
+    MemoryPool* requestor,
+    uint64_t targetBytes) {
+  if (parent_ != nullptr) {
+    return parent_->shrinkManaged(requestor, targetBytes);
+  }
+  VELOX_CHECK_NULL(parent_);
+  return manager_->shrinkPool(requestor, targetBytes);
+};
+
 bool MemoryPoolImpl::maybeIncrementReservation(uint64_t size) {
   std::lock_guard<std::mutex> l(mutex_);
   return maybeIncrementReservationLocked(size);

velox/common/memory/MemoryPool.h

@@ -105,7 +105,6 @@ constexpr int64_t kMaxMemory = std::numeric_limits<int64_t>::max();
 /// be merged into memory pool object later.
 class MemoryPool : public std::enable_shared_from_this<MemoryPool> {
  public:
-  using HighUsageCallBack = std::function<bool(MemoryPool& Pool)>;
   /// Defines the kinds of a memory pool.
   enum class Kind {
     /// The leaf memory pool is used for memory allocation. User can allocate
@@ -312,14 +311,6 @@ class MemoryPool : public std::enable_shared_from_this<MemoryPool> {
   /// 'capacity()' is fixed and set to 'maxCapacity()' on creation.
   virtual int64_t capacity() const = 0;
 
-  virtual bool highUsage() = 0;
-
-  virtual void setHighUsageCallback(HighUsageCallBack func) {
-    VELOX_CHECK_NULL(
-        parent_, "Only root memory pool allows to set high-usage callback");
-    highUsageCallback_ = func;
-  }
-
   /// Returns the currently used memory in bytes of this memory pool.
   virtual int64_t currentBytes() const = 0;
 
@@ -358,6 +349,12 @@ class MemoryPool : public std::enable_shared_from_this<MemoryPool> {
   /// without actually freeing the used memory.
   virtual uint64_t freeBytes() const = 0;
 
+  /// Try shrinking up to the specified amount of free memory via memory
+  /// manager.
+  virtual uint64_t shrinkManaged(
+      MemoryPool* requestor,
+      uint64_t targetBytes = 0) = 0;
+
   /// Invoked to free up to the specified amount of free memory by reducing
   /// this memory pool's capacity without actually freeing any used memory. The
   /// function returns the actually freed memory capacity in bytes. If
@@ -522,8 +519,6 @@ class MemoryPool : public std::enable_shared_from_this<MemoryPool> {
   // visitChildren() cost as we don't have to upgrade the weak pointer and copy
   // out the upgraded shared pointers.git
   std::unordered_map<std::string, std::weak_ptr<MemoryPool>> children_;
-
-  HighUsageCallBack highUsageCallback_{};
 };
 
 std::ostream& operator<<(std::ostream& out, MemoryPool::Kind kind);
@@ -571,8 +566,6 @@ class MemoryPoolImpl : public MemoryPool {
 
   int64_t capacity() const override;
 
-  bool highUsage() override;
-
   int64_t currentBytes() const override {
     std::lock_guard<std::mutex> l(mutex_);
     return currentBytesLocked();
@@ -611,6 +604,8 @@ class MemoryPoolImpl : public MemoryPool {
 
   uint64_t reclaim(uint64_t targetBytes) override;
 
+  uint64_t shrinkManaged(MemoryPool* requestor, uint64_t targetBytes) override;
+
   uint64_t shrink(uint64_t targetBytes = 0) override;
 
   uint64_t grow(uint64_t bytes) noexcept override;

velox/common/memory/SharedArbitrator.cpp

@@ -124,10 +124,11 @@ void SharedArbitrator::reserveMemory(MemoryPool* pool, uint64_t /*unused*/) {
   pool->grow(reserveBytes);
 }
 
-void SharedArbitrator::releaseMemory(MemoryPool* pool) {
+uint64_t SharedArbitrator::releaseMemory(MemoryPool* pool, uint64_t bytes) {
   std::lock_guard<std::mutex> l(mutex_);
-  const uint64_t freedBytes = pool->shrink(0);
+  const uint64_t freedBytes = pool->shrink(bytes);
   incrementFreeCapacityLocked(freedBytes);
+  return freedBytes;
 }
 
 std::vector<SharedArbitrator::Candidate> SharedArbitrator::getCandidateStats(
@@ -246,54 +247,59 @@ bool SharedArbitrator::arbitrateMemory(
     uint64_t targetBytes) {
   VELOX_CHECK(!requestor->aborted());
 
-  const uint64_t growTarget = std::min(
-      maxGrowBytes(*requestor),
-      std::max(minMemoryPoolCapacityTransferSize_, targetBytes));
-  uint64_t freedBytes = decrementFreeCapacity(growTarget);
-  if (freedBytes >= targetBytes) {
-    requestor->grow(freedBytes);
-    return true;
-  }
-  VELOX_CHECK_LT(freedBytes, growTarget);
+  const uint64_t growTarget =
+      std::max(minMemoryPoolCapacityTransferSize_, targetBytes);
+  uint64_t unusedFreedBytes = decrementFreeCapacity(growTarget);
 
   auto freeGuard = folly::makeGuard([&]() {
     // Returns the unused freed memory capacity back to the arbitrator.
-    if (freedBytes > 0) {
-      incrementFreeCapacity(freedBytes);
+    if (unusedFreedBytes > 0) {
+      incrementFreeCapacity(unusedFreedBytes);
     }
   });
 
-  freedBytes +=
-      reclaimFreeMemoryFromCandidates(candidates, growTarget - freedBytes);
-  if (freedBytes >= targetBytes) {
-    const uint64_t bytesToGrow = std::min(growTarget, freedBytes);
-    requestor->grow(bytesToGrow);
-    freedBytes -= bytesToGrow;
+  if (unusedFreedBytes >= targetBytes) {
+    requestor->grow(unusedFreedBytes);
+    unusedFreedBytes = 0;
+    return true;
+  }
+  VELOX_CHECK_LT(unusedFreedBytes, growTarget);
+
+  reclaimFreeMemoryFromCandidates(candidates, growTarget - unusedFreedBytes);
+  unusedFreedBytes += decrementFreeCapacity(growTarget - unusedFreedBytes);
+  if (unusedFreedBytes >= targetBytes) {
+    requestor->grow(unusedFreedBytes);
+    unusedFreedBytes = 0;
     return true;
   }
 
-  VELOX_CHECK_LT(freedBytes, growTarget);
-  freedBytes += reclaimUsedMemoryFromCandidates(
-      requestor, candidates, growTarget - freedBytes);
+  VELOX_CHECK_LT(unusedFreedBytes, growTarget);
+  reclaimUsedMemoryFromCandidates(
+      requestor, candidates, growTarget - unusedFreedBytes);
+  unusedFreedBytes += decrementFreeCapacity(growTarget - unusedFreedBytes);
   if (requestor->aborted()) {
     ++numFailures_;
     VELOX_MEM_POOL_ABORTED(requestor);
   }
 
   VELOX_CHECK(!requestor->aborted());
 
-  if (freedBytes < targetBytes) {
+  if (unusedFreedBytes < targetBytes) {
     VELOX_MEM_LOG(WARNING)
         << "Failed to arbitrate sufficient memory for memory pool "
         << requestor->name() << ", request " << succinctBytes(targetBytes)
-        << ", only " << succinctBytes(freedBytes)
+        << ", only " << succinctBytes(unusedFreedBytes)
         << " has been freed, Arbitrator state: " << toString();
     return false;
   }
 
-  const uint64_t bytesToGrow = std::min(freedBytes, growTarget);
-  requestor->grow(bytesToGrow);
-  freedBytes -= bytesToGrow;
+  if (unusedFreedBytes > growTarget) {
+    requestor->grow(growTarget);
+    unusedFreedBytes -= growTarget;
+    return true;
+  }
+  requestor->grow(unusedFreedBytes);
+  unusedFreedBytes = 0;
   return true;
 }
 
@@ -314,7 +320,9 @@ uint64_t SharedArbitrator::reclaimFreeMemoryFromCandidates(
     if (bytesToShrink <= 0) {
       break;
     }
-    freedBytes += candidate.pool->shrink(bytesToShrink);
+    uint64_t shrunk = candidate.pool->shrink(bytesToShrink);
+    incrementFreeCapacity(shrunk);
+    freedBytes += shrunk;
     if (freedBytes >= targetBytes) {
       break;
     }
@@ -354,6 +362,7 @@ uint64_t SharedArbitrator::reclaim(
   uint64_t freedBytes{0};
   try {
     freedBytes = pool->shrink(targetBytes);
+    incrementFreeCapacity(freedBytes);
     if (freedBytes < targetBytes) {
       pool->reclaim(targetBytes - freedBytes);
     }
@@ -363,7 +372,7 @@ uint64_t SharedArbitrator::reclaim(
     abort(pool);
     // Free up all the free capacity from the aborted pool as the associated
     // query has failed at this point.
-    pool->shrink();
+    incrementFreeCapacity(pool->shrink());
   }
   const uint64_t newCapacity = pool->capacity();
   VELOX_CHECK_GE(oldCapacity, newCapacity);

velox/common/memory/SharedArbitrator.h

@@ -33,7 +33,7 @@ class SharedArbitrator : public MemoryArbitrator {
 
   void reserveMemory(MemoryPool* pool, uint64_t /*unused*/) final;
 
-  void releaseMemory(MemoryPool* pool) final;
+  uint64_t releaseMemory(MemoryPool* pool, uint64_t bytes) final;
 
   bool growMemory(
       MemoryPool* pool,