From e9bb6c18df1144e5581b1dc2ef71b210cae97f54 Mon Sep 17 00:00:00 2001
From: Jialiang Tan <jacob.jialiang.tan@gmail.com>
Date: Sat, 21 Dec 2024 00:14:23 -0800
Subject: [PATCH] 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: https://github.com/facebookincubator/velox/pull/11925

Reviewed By: xiaoxmeng

Differential Revision: D67554750

Pulled By: tanjialiang

fbshipit-source-id: 7e669daebfaa7eeaf14dec6c7ad35e852d1b5924
---
 velox/common/memory/MemoryArbitrator.cpp |  7 +++--
 velox/exec/tests/MemoryReclaimerTest.cpp | 34 ++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/velox/common/memory/MemoryArbitrator.cpp b/velox/common/memory/MemoryArbitrator.cpp
index a98ecd395788..74e1af1731e1 100644
--- a/velox/common/memory/MemoryArbitrator.cpp
+++ b/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) {
diff --git a/velox/exec/tests/MemoryReclaimerTest.cpp b/velox/exec/tests/MemoryReclaimerTest.cpp
index 2e12003f4bc8..adaef7302697 100644
--- a/velox/exec/tests/MemoryReclaimerTest.cpp
+++ b/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