Issues calling reclaimer / arbitrator APIs in single-thread execution (5790)

Author: rui-mo

velox/common/memory/Memory.cpp

@@ -129,19 +129,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);
   arbitrator_->reserveMemory(pool.get(), capacity);
   return pool;
@@ -158,6 +161,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);
@@ -176,7 +187,7 @@ void MemoryManager::dropPool(MemoryPool* pool) {
     VELOX_FAIL("The dropped memory pool {} not found", pool->name());
   }
   pools_.erase(it);
-  arbitrator_->releaseMemory(pool);
+  arbitrator_->releaseMemory(pool, 0);
 }
 
 MemoryPool& MemoryManager::deprecatedSharedLeafPool() {

velox/common/memory/Memory.h

@@ -161,6 +161,10 @@ class MemoryManager {
       const std::string& name = "",
       bool threadSafe = true);
 
+  /// Invoked to shrink a memory pool's free capacity with up to
+  /// 'decrementBytes'.
+  uint64_t shrinkPool(MemoryPool* pool, uint64_t decrementBytes);
+
   /// Invoked to grows a memory pool's free capacity with at least
   /// 'incrementBytes'. The function returns true on success, otherwise false.
   bool growPool(MemoryPool* pool, uint64_t incrementBytes);

velox/common/memory/MemoryArbitrator.cpp

@@ -96,8 +96,9 @@ class NoopArbitrator : public MemoryArbitrator {
 
   // Noop arbitrator has no memory capacity limit so no operation needed for
   // memory pool capacity release.
-  void releaseMemory(MemoryPool* /*unused*/) override {
+  uint64_t releaseMemory(MemoryPool* /*unused*/, uint64_t /*unused*/) override {
     // No-op
+    return 0ULL;
   }
 
   // Noop arbitrator has no memory capacity limit so no operation needed for
@@ -161,6 +162,10 @@ void MemoryArbitrator::unregisterAllFactories() {
   SharedArbitrator::unregisterFactory();
 }
 
+uint64_t MemoryArbitrator::capacity() {
+  return capacity_;
+}
+
 std::unique_ptr<MemoryReclaimer> MemoryReclaimer::create() {
   return std::unique_ptr<MemoryReclaimer>(new MemoryReclaimer());
 }

velox/common/memory/MemoryArbitrator.h

@@ -120,9 +120,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
@@ -148,6 +149,7 @@ class MemoryArbitrator {
       const std::vector<std::shared_ptr<MemoryPool>>& pools,
       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

@@ -756,6 +756,16 @@ bool MemoryPoolImpl::incrementReservationThreadSafe(
       treeMemoryUsage()));
 }
 
+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

@@ -363,6 +363,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
@@ -625,6 +631,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

@@ -163,10 +163,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,10 +247,7 @@ bool SharedArbitrator::ensureCapacity(
   if (checkCapacityGrowth(*requestor, targetBytes)) {
     return true;
   }
-  const uint64_t reclaimedBytes = reclaim(requestor, targetBytes);
-  // NOTE: return the reclaimed bytes back to the arbitrator and let the memory
-  // arbitration process to grow the requestor's memory capacity accordingly.
-  incrementFreeCapacity(reclaimedBytes);
+  reclaim(requestor, targetBytes);
   // Check if the requestor has been aborted in reclaim operation above.
   if (requestor->aborted()) {
     ++numFailures_;
@@ -294,51 +292,57 @@ bool SharedArbitrator::arbitrateMemory(
   const uint64_t growTarget = std::min(
       maxGrowBytes(*requestor),
       std::max(memoryPoolTransferCapacity_, targetBytes));
-  uint64_t freedBytes = decrementFreeCapacity(growTarget);
-  if (freedBytes >= targetBytes) {
-    requestor->grow(freedBytes);
-    return true;
-  }
-  VELOX_CHECK_LT(freedBytes, growTarget);
+  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("The requestor pool has been aborted.");
   }
 
   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;
 }
 
@@ -359,7 +363,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;
     }
@@ -399,6 +405,7 @@ uint64_t SharedArbitrator::reclaim(
   uint64_t freedBytes{0};
   try {
     freedBytes = pool->shrink(targetBytes);
+    incrementFreeCapacity(freedBytes);
     if (freedBytes < targetBytes) {
       pool->reclaim(targetBytes - freedBytes);
     }
@@ -408,7 +415,7 @@ uint64_t SharedArbitrator::reclaim(
     abort(pool, std::current_exception());
     // 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

@@ -43,7 +43,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,

velox/common/memory/tests/MemoryArbitratorTest.cpp

@@ -157,13 +157,9 @@ class FakeTestArbitrator : public MemoryArbitrator {
              .memoryPoolTransferCapacity = config.memoryPoolTransferCapacity,
              .retryArbitrationFailure = config.retryArbitrationFailure}) {}
 
-  void reserveMemory(MemoryPool* pool, uint64_t bytes) override {
-    VELOX_NYI();
-  }
+  void reserveMemory(MemoryPool* pool, uint64_t bytes) override{VELOX_NYI()}
 
-  void releaseMemory(MemoryPool* pool) override {
-    VELOX_NYI();
-  }
+  uint64_t releaseMemory(MemoryPool* pool, uint64_t bytes) override{VELOX_NYI()}
 
   std::string kind() const override {
     return "USER";

velox/common/memory/tests/MemoryManagerTest.cpp

@@ -125,7 +125,7 @@ class FakeTestArbitrator : public MemoryArbitrator {
     VELOX_NYI();
   }
 
-  void releaseMemory(MemoryPool* pool) override {
+  uint64_t releaseMemory(MemoryPool* pool, uint64_t bytes) override {
     VELOX_NYI();
   }
 
@@ -209,9 +209,6 @@ TEST_F(MemoryManagerTest, addPoolWithArbitrator) {
   options.allocator = allocator.get();
   options.capacity = kCapacity;
   options.arbitratorKind = arbitratorKind_;
-  // The arbitrator capacity will be overridden by the memory manager's
-  // capacity.
-  options.capacity = options.capacity;
   const uint64_t initialPoolCapacity = options.capacity / 32;
   options.memoryPoolInitCapacity = initialPoolCapacity;
   MemoryManager manager{options};

velox/common/memory/tests/MockSharedArbitratorTest.cpp

@@ -307,7 +307,7 @@ class MockMemoryOperator {
     for (const auto& allocation : allocationsToFree) {
       pool_->free(allocation.buffer, allocation.size);
     }
-    return pool_->shrink(targetBytes);
+    return pool_->shrinkManaged(pool, targetBytes);
   }
 
   void abort(MemoryPool* pool) {

velox/core/QueryCtx.h

@@ -67,11 +67,13 @@ class QueryCtx {
     return cache_;
   }
 
-  folly::Executor* executor() const {
-    if (executor_ != nullptr) {
-      return executor_;
-    }
-    auto executor = executorKeepalive_.get();
+  bool isExecutorSupplied() const {
+    auto executor = executor0();
+    return executor != nullptr;
+  }
+
+  folly::Executor* FOLLY_NONNULL executor() const {
+    auto executor = executor0();
     VELOX_CHECK(executor, "Executor was not supplied.");
     return executor;
   }
@@ -130,6 +132,14 @@ class QueryCtx {
     }
   }
 
+  folly::Executor* executor0() const {
+    if (executor_ != nullptr) {
+      return executor_;
+    }
+    auto executor = executorKeepalive_.get();
+    return executor;
+  }
+
   const std::string queryId_;
 
   std::unordered_map<std::string, std::shared_ptr<Config>> connectorConfigs_;

velox/dwio/dwrf/test/WriterFlushTest.cpp

@@ -208,6 +208,10 @@ class MockMemoryPool : public velox::memory::MemoryPool {
     VELOX_UNSUPPORTED("{} unsupported", __FUNCTION__);
   }
 
+  uint64_t shrinkManaged(MemoryPool* /*unused*/, uint64_t /*unused*/) override {
+    VELOX_UNSUPPORTED("{} unsupported", __FUNCTION__);
+  }
+
   uint64_t grow(uint64_t /*unused*/) noexcept override {
     VELOX_UNSUPPORTED("{} unsupported", __FUNCTION__);
   }