Free the unused free capacity when reclaim (#5034)

Summary:
Add to reclaim unused capacity in SharedArbitrator::reclaim since
query is still running so add a step to free unused capacity before
the actual memory reclamation

Pull Request resolved: #5034

Reviewed By: amitkdutta

Differential Revision: D46132850

Pulled By: xiaoxmeng

fbshipit-source-id: 9572442d0fd3cff314f98e7e89fe29b0cb2a7ff4

velox/common/memory/SharedArbitrator.cpp

@@ -215,16 +215,19 @@ uint64_t SharedArbitrator::reclaimUsedMemoryFromCandidates(
       break;
     }
   }
-  numReclaimedBytes_ += freedBytes;
   return freedBytes;
 }
 
 uint64_t SharedArbitrator::reclaim(
     MemoryPool* pool,
     uint64_t targetBytes) noexcept {
   const uint64_t oldCapacity = pool->capacity();
+  uint64_t freedBytes{0};
   try {
-    pool->reclaim(targetBytes);
+    freedBytes = pool->shrink(targetBytes);
+    if (freedBytes < targetBytes) {
+      pool->reclaim(targetBytes - freedBytes);
+    }
   } catch (const std::exception& e) {
     VELOX_MEM_LOG(ERROR) << "Failed to reclaim from memory pool "
                          << pool->name() << ", aborting it!";
@@ -235,7 +238,10 @@ uint64_t SharedArbitrator::reclaim(
   }
   const uint64_t newCapacity = pool->capacity();
   VELOX_CHECK_GE(oldCapacity, newCapacity);
-  return oldCapacity - newCapacity;
+  const uint64_t reclaimedbytes = oldCapacity - newCapacity;
+  numShrunkBytes_ += freedBytes;
+  numReclaimedBytes_ += reclaimedbytes - freedBytes;
+  return reclaimedbytes;
 }
 
 void SharedArbitrator::abort(MemoryPool* pool) {

velox/common/memory/tests/MockSharedArbitratorTest.cpp

@@ -1323,6 +1323,55 @@ TEST_F(MockSharedArbitrationTest, memoryPoolAbortThrow) {
   ASSERT_EQ(arbitrator_->stats().numAborted, 1);
 }
 
+DEBUG_ONLY_TEST_F(
+    MockSharedArbitrationTest,
+    freeUnusedCapacityWhenReclaimMemoryPool) {
+  const int allocationSize = kMemoryCapacity / 4;
+  const int reclaimedMemoryCapacity = kMemoryCapacity;
+  std::shared_ptr<MockQuery> reclaimedQuery = addQuery(kMemoryCapacity);
+  MockMemoryOperator* reclaimedQueryOp = addMemoryOp(reclaimedQuery);
+  // The buffer to free later.
+  void* bufferToFree = reclaimedQueryOp->allocate(allocationSize);
+  reclaimedQueryOp->allocate(kMemoryCapacity - allocationSize);
+
+  std::shared_ptr<MockQuery> arbitrationQuery = addQuery(kMemoryCapacity);
+  MockMemoryOperator* arbitrationQueryOp = addMemoryOp(arbitrationQuery);
+
+  folly::EventCount reclaimWait;
+  auto reclaimWaitKey = reclaimWait.prepareWait();
+  folly::EventCount reclaimBlock;
+  auto reclaimBlockKey = reclaimBlock.prepareWait();
+  SCOPED_TESTVALUE_SET(
+      "facebook::velox::memory::SharedArbitrator::sortCandidatesByReclaimableMemory",
+      std::function<void(const MemoryPool*)>(([&](const MemoryPool* /*unsed*/) {
+        reclaimWait.notify();
+        reclaimBlock.wait(reclaimBlockKey);
+      })));
+
+  const auto oldStats = arbitrator_->stats();
+
+  std::thread allocThread([&]() {
+    // Allocate to trigger arbitration.
+    arbitrationQueryOp->allocate(allocationSize);
+  });
+
+  reclaimWait.wait(reclaimWaitKey);
+  reclaimedQueryOp->free(bufferToFree);
+  reclaimBlock.notify();
+
+  allocThread.join();
+
+  const auto stats = arbitrator_->stats();
+  ASSERT_EQ(stats.numFailures, 0);
+  ASSERT_EQ(stats.numAborted, 0);
+  ASSERT_EQ(stats.numRequests, oldStats.numRequests + 1);
+  // We count the freed capacity in reclaimed bytes.
+  ASSERT_EQ(stats.numShrunkBytes, oldStats.numShrunkBytes + allocationSize);
+  ASSERT_EQ(stats.numReclaimedBytes, 0);
+  ASSERT_EQ(reclaimedQueryOp->capacity(), kMemoryCapacity - allocationSize);
+  ASSERT_EQ(arbitrationQueryOp->capacity(), allocationSize);
+}
+
 DEBUG_ONLY_TEST_F(
     MockSharedArbitrationTest,
     raceBetweenInitialReservationAndArbitration) {