Add shrinkPool implementation (facebookincubator#8865)

Summary:
Add shrinkPool implementation to memory manager and arbitrator.

Pull Request resolved: facebookincubator#8865

Reviewed By: xiaoxmeng

Differential Revision: D54274877

Pulled By: tanjialiang

fbshipit-source-id: 631d4cc04f7932bf16823d7ef5c222ce932503b5

velox/common/memory/MemoryArbitrator.cpp

@@ -116,8 +116,8 @@ class NoopArbitrator : public MemoryArbitrator {
   // Noop arbitrator has no memory capacity limit so no operation needed for
   // memory pool capacity shrink.
   uint64_t shrinkCapacity(
-      const std::vector<std::shared_ptr<MemoryPool>>& /*unused*/,
-      uint64_t /*unused*/) override {
+      const std::vector<std::shared_ptr<MemoryPool>>& /* unused */,
+      uint64_t /* unused */) override {
     return 0;
   }
 
@@ -440,7 +440,7 @@ bool MemoryArbitrator::Stats::operator<=(const Stats& other) const {
 }
 
 ScopedMemoryArbitrationContext::ScopedMemoryArbitrationContext(
-    const MemoryPool& requestor)
+    const MemoryPool* requestor)
     : savedArbitrationCtx_(arbitrationCtx),
       currentArbitrationCtx_({.requestor = requestor}) {
   arbitrationCtx = &currentArbitrationCtx_;

velox/common/memory/MemoryArbitrator.h

@@ -148,9 +148,9 @@ class MemoryArbitrator {
   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 used memory. The freed memory capacity is
-  /// given back to the arbitrator. The function returns the actual freed memory
-  /// capacity in bytes.
+  /// of memory pools by reclaiming free and used memory. The freed memory
+  /// capacity is given back to the arbitrator. The function returns the actual
+  /// freed memory capacity in bytes.
   virtual uint64_t shrinkCapacity(
       const std::vector<std::shared_ptr<MemoryPool>>& pools,
       uint64_t targetBytes) = 0;
@@ -361,14 +361,14 @@ class MemoryReclaimer {
 /// the memory reservations during memory arbitration should come from the
 /// spilling memory pool.
 struct MemoryArbitrationContext {
-  const MemoryPool& requestor;
+  const MemoryPool* requestor;
 };
 
 /// Object used to set/restore the memory arbitration context when a thread is
 /// under memory arbitration processing.
 class ScopedMemoryArbitrationContext {
  public:
-  explicit ScopedMemoryArbitrationContext(const MemoryPool& requestor);
+  explicit ScopedMemoryArbitrationContext(const MemoryPool* requestor);
   ~ScopedMemoryArbitrationContext();
 
  private:

velox/common/memory/tests/MemoryArbitratorTest.cpp

@@ -681,29 +681,29 @@ TEST_F(MemoryReclaimerTest, arbitrationContext) {
   ASSERT_FALSE(isSpillMemoryPool(leafChild2.get()));
   ASSERT_TRUE(memoryArbitrationContext() == nullptr);
   {
-    ScopedMemoryArbitrationContext arbitrationContext(*leafChild1);
+    ScopedMemoryArbitrationContext arbitrationContext(leafChild1.get());
     ASSERT_TRUE(memoryArbitrationContext() != nullptr);
-    ASSERT_EQ(&memoryArbitrationContext()->requestor, leafChild1.get());
+    ASSERT_EQ(memoryArbitrationContext()->requestor, leafChild1.get());
   }
   ASSERT_TRUE(memoryArbitrationContext() == nullptr);
   {
-    ScopedMemoryArbitrationContext arbitrationContext(*leafChild2);
+    ScopedMemoryArbitrationContext arbitrationContext(leafChild2.get());
     ASSERT_TRUE(memoryArbitrationContext() != nullptr);
-    ASSERT_EQ(&memoryArbitrationContext()->requestor, leafChild2.get());
+    ASSERT_EQ(memoryArbitrationContext()->requestor, leafChild2.get());
   }
   ASSERT_TRUE(memoryArbitrationContext() == nullptr);
   std::thread nonAbitrationThread([&]() {
     ASSERT_TRUE(memoryArbitrationContext() == nullptr);
     {
-      ScopedMemoryArbitrationContext arbitrationContext(*leafChild1);
+      ScopedMemoryArbitrationContext arbitrationContext(leafChild1.get());
       ASSERT_TRUE(memoryArbitrationContext() != nullptr);
-      ASSERT_EQ(&memoryArbitrationContext()->requestor, leafChild1.get());
+      ASSERT_EQ(memoryArbitrationContext()->requestor, leafChild1.get());
     }
     ASSERT_TRUE(memoryArbitrationContext() == nullptr);
     {
-      ScopedMemoryArbitrationContext arbitrationContext(*leafChild2);
+      ScopedMemoryArbitrationContext arbitrationContext(leafChild2.get());
       ASSERT_TRUE(memoryArbitrationContext() != nullptr);
-      ASSERT_EQ(&memoryArbitrationContext()->requestor, leafChild2.get());
+      ASSERT_EQ(memoryArbitrationContext()->requestor, leafChild2.get());
     }
     ASSERT_TRUE(memoryArbitrationContext() == nullptr);
   });

velox/common/memory/tests/MemoryPoolTest.cpp

@@ -3662,7 +3662,7 @@ TEST_P(MemoryPoolTest, overuseUnderArbitration) {
   ASSERT_FALSE(child->maybeReserve(2 * kMaxSize));
   ASSERT_EQ(child->currentBytes(), 0);
   ASSERT_EQ(child->reservedBytes(), 0);
-  ScopedMemoryArbitrationContext scopedMemoryArbitration(*child);
+  ScopedMemoryArbitrationContext scopedMemoryArbitration(child.get());
   ASSERT_TRUE(underMemoryArbitration());
   ASSERT_TRUE(child->maybeReserve(2 * kMaxSize));
   ASSERT_EQ(child->currentBytes(), 0);

velox/common/memory/tests/MockSharedArbitratorTest.cpp

@@ -591,8 +591,30 @@ TEST_F(MockSharedArbitrationTest, arbitrationFailsTask) {
 }
 
 TEST_F(MockSharedArbitrationTest, shrinkMemory) {
-  std::vector<std::shared_ptr<MemoryPool>> pools;
-  ASSERT_THROW(arbitrator_->shrinkCapacity(pools, 128), VeloxException);
+  auto task1 = addTask(64 * MB);
+  auto op1 = addMemoryOp(task1);
+  auto task2 = addTask(64 * MB);
+  auto op2 = addMemoryOp(task2);
+
+  op1->allocate(64 * MB);
+  op1->freeAll();
+  auto bufOp11 = op1->allocate(32 * MB);
+  auto bufOp12 = op1->allocate(32 * MB);
+  op1->free(bufOp11);
+  ASSERT_EQ(op1->pool()->root()->capacity(), 64 * MB);
+  ASSERT_EQ(op1->pool()->root()->currentBytes(), 32 * MB);
+
+  op2->allocate(64 * MB);
+  op2->freeAll();
+  auto bufOp21 = op2->allocate(32 * MB);
+  auto bufOp22 = op2->allocate(32 * MB);
+  op2->free(bufOp21);
+  ASSERT_EQ(op2->pool()->root()->capacity(), 64 * MB);
+  ASSERT_EQ(op2->pool()->root()->currentBytes(), 32 * MB);
+
+  ASSERT_EQ(manager_->shrinkPools(kMaxMemory), 128 * MB);
+  ASSERT_EQ(op1->capacity(), 0);
+  ASSERT_EQ(op2->capacity(), 0);
 }
 
 TEST_F(MockSharedArbitrationTest, singlePoolGrowWithoutArbitration) {

velox/exec/SharedArbitrator.cpp

@@ -205,6 +205,22 @@ uint64_t SharedArbitrator::shrinkCapacity(
   return freedBytes;
 }
 
+uint64_t SharedArbitrator::shrinkCapacity(
+    const std::vector<std::shared_ptr<MemoryPool>>& pools,
+    uint64_t targetBytes) {
+  ScopedArbitration scopedArbitration(this);
+  targetBytes = std::max(memoryPoolTransferCapacity_, targetBytes);
+  std::vector<Candidate> candidates = getCandidateStats(pools);
+  auto freedBytes = reclaimFreeMemoryFromCandidates(candidates, targetBytes);
+  if (freedBytes >= targetBytes) {
+    return freedBytes;
+  }
+  freedBytes += reclaimUsedMemoryFromCandidates(
+      nullptr, candidates, targetBytes - freedBytes);
+  incrementFreeCapacity(freedBytes);
+  return freedBytes;
+}
+
 std::vector<SharedArbitrator::Candidate> SharedArbitrator::getCandidateStats(
     const std::vector<std::shared_ptr<MemoryPool>>& pools) {
   std::vector<SharedArbitrator::Candidate> candidates;
@@ -437,7 +453,8 @@ uint64_t SharedArbitrator::reclaimUsedMemoryFromCandidates(
         targetBytes - freedBytes, memoryPoolTransferCapacity_);
     VELOX_CHECK_GT(bytesToReclaim, 0);
     freedBytes += reclaim(candidate.pool, bytesToReclaim);
-    if ((freedBytes >= targetBytes) || requestor->aborted()) {
+    if ((freedBytes >= targetBytes) ||
+        (requestor != nullptr && requestor->aborted())) {
       break;
     }
   }
@@ -580,22 +597,39 @@ std::string SharedArbitrator::toStringLocked() const {
       statsLocked().toString());
 }
 
+SharedArbitrator::ScopedArbitration::ScopedArbitration(
+    SharedArbitrator* arbitrator)
+    : requestor_(nullptr),
+      arbitrator_(arbitrator),
+      startTime_(std::chrono::steady_clock::now()),
+      arbitrationCtx_(requestor_) {
+  VELOX_CHECK_NOT_NULL(arbitrator_);
+  arbitrator_->startArbitration("Wait for arbitration");
+}
+
 SharedArbitrator::ScopedArbitration::ScopedArbitration(
     MemoryPool* requestor,
     SharedArbitrator* arbitrator)
     : requestor_(requestor),
       arbitrator_(arbitrator),
       startTime_(std::chrono::steady_clock::now()),
-      arbitrationCtx_(*requestor_) {
+      arbitrationCtx_(requestor_) {
   VELOX_CHECK_NOT_NULL(arbitrator_);
-  arbitrator_->startArbitration(requestor);
-  if (arbitrator_->arbitrationStateCheckCb_ != nullptr) {
-    arbitrator_->arbitrationStateCheckCb_(*requestor);
+  requestor_->enterArbitration();
+  arbitrator_->startArbitration(fmt::format(
+      "Wait for arbitration, requestor: {}[{}]",
+      requestor_->name(),
+      requestor_->root()->name()));
+  if (arbitrator_->arbitrationStateCheckCb_ != nullptr &&
+      requestor_ != nullptr) {
+    arbitrator_->arbitrationStateCheckCb_(*requestor_);
   }
 }
 
 SharedArbitrator::ScopedArbitration::~ScopedArbitration() {
-  requestor_->leaveArbitration();
+  if (requestor_ != nullptr) {
+    requestor_->leaveArbitration();
+  }
   const auto arbitrationTimeUs =
       std::chrono::duration_cast<std::chrono::microseconds>(
           std::chrono::steady_clock::now() - startTime_)
@@ -609,18 +643,14 @@ SharedArbitrator::ScopedArbitration::~ScopedArbitration() {
   arbitrator_->finishArbitration();
 }
 
-void SharedArbitrator::startArbitration(MemoryPool* requestor) {
-  requestor->enterArbitration();
+void SharedArbitrator::startArbitration(const std::string& arbitrationContext) {
   ContinueFuture waitPromise{ContinueFuture::makeEmpty()};
   {
     std::lock_guard<std::mutex> l(mutex_);
     RECORD_METRIC_VALUE(kMetricArbitratorRequestsCount);
     ++numRequests_;
     if (running_) {
-      waitPromises_.emplace_back(fmt::format(
-          "Wait for arbitration, requestor: {}[{}]",
-          requestor->name(),
-          requestor->root()->name()));
+      waitPromises_.emplace_back(arbitrationContext);
       waitPromise = waitPromises_.back().getSemiFuture();
     } else {
       VELOX_CHECK(waitPromises_.empty());
@@ -629,7 +659,7 @@ void SharedArbitrator::startArbitration(MemoryPool* requestor) {
   }
 
   TestValue::adjust(
-      "facebook::velox::memory::SharedArbitrator::startArbitration", requestor);
+      "facebook::velox::memory::SharedArbitrator::startArbitration", nullptr);
 
   if (waitPromise.valid()) {
     uint64_t waitTimeUs{0};

velox/exec/SharedArbitrator.h

@@ -53,10 +53,8 @@ class SharedArbitrator : public memory::MemoryArbitrator {
   uint64_t shrinkCapacity(MemoryPool* pool, uint64_t freedBytes) final;
 
   uint64_t shrinkCapacity(
-      const std::vector<std::shared_ptr<MemoryPool>>& /*unused*/,
-      uint64_t /*unused*/) override final {
-    VELOX_NYI("shrinkCapacity is not supported by SharedArbitrator");
-  }
+      const std::vector<std::shared_ptr<MemoryPool>>& pools,
+      uint64_t targetBytes) override final;
 
   Stats stats() const final;
 
@@ -80,7 +78,14 @@ class SharedArbitrator : public memory::MemoryArbitrator {
 
   class ScopedArbitration {
    public:
-    ScopedArbitration(MemoryPool* requestor, SharedArbitrator* arbitrator);
+    // Used by arbitration request NOT initiated from memory pool, e.g. through
+    // shrinkPools() API.
+    explicit ScopedArbitration(SharedArbitrator* arbitrator);
+
+    // Used by arbitration request initiated from a memory pool.
+    explicit ScopedArbitration(
+        MemoryPool* requestor,
+        SharedArbitrator* arbitrator);
 
     ~ScopedArbitration();
 
@@ -129,7 +134,7 @@ class SharedArbitrator : public memory::MemoryArbitrator {
   // Invoked to start next memory arbitration request, and it will wait for the
   // serialized execution if there is a running or other waiting arbitration
   // requests.
-  void startArbitration(MemoryPool* requestor);
+  void startArbitration(const std::string& arbitrationContext);
 
   // Invoked by a finished memory arbitration request to kick off the next
   // arbitration request execution if there are any ones waiting.