Add shrinkPool API to MemoryManager/MemoryArbitrator

velox/common/memory/Memory.cpp

@@ -181,6 +181,11 @@ bool MemoryManager::growPool(MemoryPool* pool, uint64_t incrementBytes) {
   return arbitrator_->growMemory(pool, candidates, incrementBytes);
 }
 
+uint64_t MemoryManager::shrinkPools(uint64_t targetBytes) {
+  VELOX_CHECK_NOT_NULL(arbitrator_);
+  return arbitrator_->shrinkMemory(getAlivePools(), targetBytes);
+}
+
 void MemoryManager::dropPool(MemoryPool* pool) {
   VELOX_CHECK_NOT_NULL(pool);
   folly::SharedMutex::WriteHolder guard{mutex_};

velox/common/memory/Memory.h

@@ -157,6 +157,11 @@ class MemoryManager {
   /// 'incrementBytes'. The function returns true on success, otherwise false.
   bool growPool(MemoryPool* pool, uint64_t incrementBytes);
 
+  /// Invoked to shrink currently alive pools by 'targetBytes', by reducing this
+  /// memory pool's capacity without actually freeing any used memory. Returns
+  /// the actual shrank memory 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

@@ -140,6 +140,13 @@ class MemoryArbitrator {
       const std::vector<std::shared_ptr<MemoryPool>>& candidatePools,
       uint64_t targetBytes) = 0;
 
+  /// Invoked by the memory manager to shrink memory for a given vector of
+  /// memory pools. The shrank memory will directly be given back to the
+  /// arbitrator. Returns the bytes shrank.
+  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

@@ -95,6 +95,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/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;