Add shrinkPool API to MemoryManager/MemoryArbitrator (#6102)

Summary: Pull Request resolved: #6102

Reviewed By: xiaoxmeng

Differential Revision: D48362262

Pulled By: tanjialiang

fbshipit-source-id: 5ac3d1deef75cb6a745a4d8644d2f7fc398d5469

velox/common/memory/Memory.cpp

@@ -175,10 +175,12 @@ bool MemoryManager::growPool(MemoryPool* pool, uint64_t incrementBytes) {
   if (arbitrator_ == nullptr) {
     return false;
   }
-  // Holds a shared reference to each alive memory pool in 'pools_' to keep
-  // their aliveness during the memory arbitration process.
-  std::vector<std::shared_ptr<MemoryPool>> candidates = getAlivePools();
-  return arbitrator_->growMemory(pool, candidates, incrementBytes);
+  return arbitrator_->growMemory(pool, getAlivePools(), incrementBytes);
+}
+
+uint64_t MemoryManager::shrinkPools(uint64_t targetBytes) {
+  VELOX_CHECK_NOT_NULL(arbitrator_);
+  return arbitrator_->shrinkMemory(getAlivePools(), targetBytes);
 }
 
 void MemoryManager::dropPool(MemoryPool* pool) {

velox/common/memory/Memory.h

@@ -157,6 +157,10 @@ class MemoryManager {
   /// 'incrementBytes'. The function returns true on success, otherwise false.
   bool growPool(MemoryPool* pool, uint64_t incrementBytes);
 
+  /// Invoked to shrink alive pools to free 'targetBytes' capacity. The function
+  /// returns the actual freed memory capacity in bytes.
+  uint64_t shrinkPools(uint64_t targetBytes);
+
   /// Default unmanaged leaf pool with no threadsafe stats support. Libraries
   /// using this method can get a pool that is shared with other threads. The
   /// goal is to minimize lock contention while supporting such use cases.

velox/common/memory/MemoryArbitrator.h

@@ -141,6 +141,13 @@ class MemoryArbitrator {
       const std::vector<std::shared_ptr<MemoryPool>>& candidatePools,
       uint64_t targetBytes) = 0;
 
+  /// Invoked by the memory manager to shrink memory from a given list of memory
+  /// pools. The freed memory capacity is given back to the arbitrator. The
+  /// function returns the actual freed memory capacity in bytes.
+  virtual uint64_t shrinkMemory(
+      const std::vector<std::shared_ptr<MemoryPool>>& pools,
+      uint64_t targetBytes) = 0;
+
   /// The internal execution stats of the memory arbitrator.
   struct Stats {
     /// The number of arbitration requests.

velox/common/memory/SharedArbitrator.h

@@ -50,6 +50,12 @@ class SharedArbitrator : public MemoryArbitrator {
       const std::vector<std::shared_ptr<MemoryPool>>& candidatePools,
       uint64_t targetBytes) final;
 
+  uint64_t shrinkMemory(
+      const std::vector<std::shared_ptr<MemoryPool>>& /*unused*/,
+      uint64_t /*unused*/) override final {
+    VELOX_NYI("shrinkMemory is not supported by SharedArbitrator");
+  }
+
   Stats stats() const final;
 
   std::string kind() override;

velox/common/memory/tests/MemoryArbitratorTest.cpp

@@ -81,10 +81,12 @@ class FakeTestArbitrator : public MemoryArbitrator {
              .retryArbitrationFailure = config.retryArbitrationFailure}) {}
 
   void reserveMemory(MemoryPool* pool, uint64_t bytes) override {
-    VELOX_NYI()
+    VELOX_NYI();
   }
 
-  void releaseMemory(MemoryPool* pool) override{VELOX_NYI()}
+  void releaseMemory(MemoryPool* pool) override {
+    VELOX_NYI();
+  }
 
   std::string kind() override {
     return "USER";
@@ -93,12 +95,22 @@ class FakeTestArbitrator : public MemoryArbitrator {
   bool growMemory(
       MemoryPool* pool,
       const std::vector<std::shared_ptr<MemoryPool>>& candidatePools,
-      uint64_t targetBytes) override{VELOX_NYI()}
+      uint64_t targetBytes) override {
+    VELOX_NYI();
+  }
 
-  Stats stats() const override{VELOX_NYI()}
+  uint64_t shrinkMemory(
+      const std::vector<std::shared_ptr<MemoryPool>>& pools,
+      uint64_t targetBytes) override {
+    VELOX_NYI();
+  }
+
+  Stats stats() const override {
+    VELOX_NYI();
+  }
 
   std::string toString() const override {
-    VELOX_NYI()
+    VELOX_NYI();
   }
 };
 } // namespace

velox/common/memory/tests/MemoryManagerTest.cpp

@@ -113,6 +113,10 @@ class FakeTestArbitrator : public MemoryArbitrator {
       const std::vector<std::shared_ptr<MemoryPool>>& candidatePools,
       uint64_t targetBytes) override{VELOX_NYI()}
 
+  uint64_t shrinkMemory(
+      const std::vector<std::shared_ptr<MemoryPool>>& pools,
+      uint64_t targetBytes) override{VELOX_NYI()}
+
   Stats stats() const override{VELOX_NYI()}
 
   std::string toString() const override{VELOX_NYI()}

velox/common/memory/tests/MockSharedArbitratorTest.cpp

@@ -535,6 +535,11 @@ TEST_F(MockSharedArbitrationTest, arbitrationFailsTask) {
   growOp->freeAll();
 }
 
+TEST_F(MockSharedArbitrationTest, shrinkMemory) {
+  std::vector<std::shared_ptr<MemoryPool>> pools;
+  ASSERT_THROW(arbitrator_->shrinkMemory(pools, 128), VeloxException);
+}
+
 TEST_F(MockSharedArbitrationTest, singlePoolGrowWithoutArbitration) {
   auto* memOp = addMemoryOp();
   const int allocateSize = 1 * MB;