Update memory reclaimer code (#430)

Author: zhztheplayer

velox/common/memory/Memory.cpp

@@ -149,14 +149,6 @@ 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);
@@ -175,7 +167,7 @@ void MemoryManager::dropPool(MemoryPool* pool) {
     VELOX_FAIL("The dropped memory pool {} not found", pool->name());
   }
   pools_.erase(it);
-  arbitrator_->releaseMemory(pool, 0);
+  arbitrator_->releaseMemory(pool);
 }
 
 MemoryPool& MemoryManager::deprecatedSharedLeafPool() {

velox/common/memory/Memory.h

@@ -160,10 +160,6 @@ 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

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

velox/common/memory/MemoryArbitrator.h

@@ -124,10 +124,9 @@ 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 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 return back all the reserved memory
+  /// capacity of a destroying memory pool.
+  virtual void releaseMemory(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
@@ -153,7 +152,6 @@ 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

@@ -759,16 +759,6 @@ 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

@@ -366,12 +366,6 @@ 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
@@ -638,8 +632,6 @@ class MemoryPoolImpl : public MemoryPool {
   uint64_t reclaim(uint64_t targetBytes, memory::MemoryReclaimer::Stats& stats)
       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/tests/MemoryArbitratorTest.cpp

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

velox/common/memory/tests/MemoryManagerTest.cpp

@@ -127,7 +127,7 @@ class FakeTestArbitrator : public MemoryArbitrator {
     VELOX_NYI();
   }
 
-  uint64_t releaseMemory(MemoryPool* pool, uint64_t bytes) override {
+  void releaseMemory(MemoryPool* pool) override {
     VELOX_NYI();
   }
 
@@ -211,6 +211,9 @@ 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

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

velox/dwio/dwrf/test/WriterFlushTest.cpp

@@ -206,10 +206,6 @@ 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__);
   }