fix: Fix 0 reclaimable bytes handling in priority reclaiming (#11925)

Summary:
Current logic skips processing memory reclaim on next priority batch if the current batch has reclaimable bytes of 0. This might cause reclaim not happen in some cases. Making it traverse all candidates fix the bug.

Pull Request resolved: #11925

Reviewed By: xiaoxmeng

Differential Revision: D67554750

Pulled By: tanjialiang

fbshipit-source-id: 7e669daebfaa7eeaf14dec6c7ad35e852d1b5924

velox/common/memory/MemoryArbitrator.cpp

@@ -257,7 +257,8 @@ uint64_t MemoryReclaimer::reclaim(
       auto child = entry.second.lock();
       if (child != nullptr) {
         const auto reclaimableBytesOpt = child->reclaimableBytes();
-        if (!reclaimableBytesOpt.has_value()) {
+        if (!reclaimableBytesOpt.has_value() ||
+            reclaimableBytesOpt.value() == 0) {
           nonReclaimableCandidates.push_back(Candidate{std::move(child), 0});
           continue;
         }
@@ -282,9 +283,7 @@ uint64_t MemoryReclaimer::reclaim(
 
   uint64_t reclaimedBytes{0};
   for (const auto& candidate : candidates) {
-    if (candidate.reclaimableBytes == 0) {
-      break;
-    }
+    VELOX_CHECK_GT(candidate.reclaimableBytes, 0);
     const auto bytes = candidate.pool->reclaim(targetBytes, maxWaitMs, stats);
     reclaimedBytes += bytes;
     if (targetBytes != 0) {

velox/exec/tests/MemoryReclaimerTest.cpp

@@ -359,6 +359,40 @@ TEST_F(MemoryReclaimerTest, reclaimerPriorities) {
   }
 }
 
+TEST_F(MemoryReclaimerTest, reclaimerPrioritiesWithZeroUsage) {
+  // This test makes sure the reclaim continues to lower priority reclaimers
+  // even if the higher priority ones have nothing to reclaim.
+  auto rootPool = memory::memoryManager()->addRootPool(
+      "reclaimerPriorities", kMaxMemory, exec::MemoryReclaimer::create());
+  std::vector<std::shared_ptr<MemoryPool>> leafPools;
+
+  const uint32_t kNumChildren = 10;
+  const uint64_t kPoolMemoryBytes = 1024;
+  std::vector<int32_t> priorityOrder;
+  priorityOrder.reserve(kNumChildren);
+  ReclaimCallback reclaimerCb = [&](MemoryPool* pool) {
+    auto* mockReclaimer = dynamic_cast<MockMemoryReclaimer*>(pool->reclaimer());
+    ASSERT_TRUE(mockReclaimer != nullptr);
+    priorityOrder.push_back(mockReclaimer->priority());
+  };
+
+  for (uint32_t i = 0; i < kNumChildren; ++i) {
+    auto reclaimer = MockMemoryReclaimer::create(
+        kPoolMemoryBytes,
+        i % 5 == 0 ? kPoolMemoryBytes : 0,
+        reclaimerCb,
+        i / 5);
+    leafPools.push_back(rootPool->addLeafChild(
+        fmt::format("leaf-{}", i), true, std::move(reclaimer)));
+  }
+
+  memory::ScopedMemoryArbitrationContext context(rootPool.get());
+  memory::MemoryReclaimer::Stats stats;
+  rootPool->reclaim(2 * kPoolMemoryBytes, 0, stats);
+
+  ASSERT_EQ(priorityOrder.size(), 2);
+}
+
 TEST_F(MemoryReclaimerTest, multiLevelReclaimerPriorities) {
   // Following tree structure is built with priorities
   //  rootPool