diff --git a/velox/common/memory/MemoryArbitrator.cpp b/velox/common/memory/MemoryArbitrator.cpp index a98ecd395788..49725f0228ae 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_NE(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> leafPools; + + const uint32_t kNumChildren = 10; + const uint64_t kPoolMemoryBytes = 1024; + std::vector priorityOrder; + priorityOrder.reserve(kNumChildren); + ReclaimCallback reclaimerCb = [&](MemoryPool* pool) { + auto* mockReclaimer = dynamic_cast(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