diff --git a/velox/common/memory/SharedArbitrator.cpp b/velox/common/memory/SharedArbitrator.cpp index bd98f4e9f4d..d16334776a1 100644 --- a/velox/common/memory/SharedArbitrator.cpp +++ b/velox/common/memory/SharedArbitrator.cpp @@ -243,6 +243,12 @@ bool SharedArbitrator::ExtraConfig::globalArbitrationWithoutSpill( kDefaultGlobalArbitrationWithoutSpill); } +double SharedArbitrator::ExtraConfig::abortPriorityWeightPct( + const std::unordered_map& configs) { + return getConfig( + configs, kAbortPriorityWeightPct, kDefaultAbortPriorityWeightPct); +} + double SharedArbitrator::ExtraConfig::globalArbitrationAbortTimeRatio( const std::unordered_map& configs) { return getConfig( @@ -279,6 +285,8 @@ SharedArbitrator::SharedArbitrator(const Config& config) ExtraConfig::globalArbitrationAbortTimeRatio(config.extraConfigs)), globalArbitrationWithoutSpill_( ExtraConfig::globalArbitrationWithoutSpill(config.extraConfigs)), + abortPriorityWeightPct_( + ExtraConfig::abortPriorityWeightPct(config.extraConfigs)), freeReservedCapacity_(reservedCapacity_), freeNonReservedCapacity_(capacity_ - freeReservedCapacity_) { VELOX_CHECK_EQ(kind_, config.kind); @@ -641,46 +649,39 @@ SharedArbitrator::sortAndGroupSpillCandidates( return candidateGroups; } +namespace { +// Abort badness, higher = better victim. Uses currentCapacity (abort frees the +// whole pool). A large priorityWeight degenerates to the legacy priority-first +// order. +__int128 abortBadnessScore( + const ArbitrationCandidate& candidate, + uint64_t priorityWeight) { + const auto* reclaimer = candidate.participant->pool()->reclaimer(); + const int64_t priority = reclaimer == nullptr + ? std::numeric_limits::max() + : reclaimer->priority(); + return static_cast<__int128>(candidate.currentCapacity) + + static_cast<__int128>(priority) * static_cast<__int128>(priorityWeight); +} +} // namespace + // static -std::vector> -SharedArbitrator::sortAndGroupAbortCandidates( - std::vector&& candidates) { +std::vector SharedArbitrator::sortAbortCandidates( + std::vector&& candidates, + uint64_t priorityWeight) { std::sort( candidates.begin(), candidates.end(), - [](const ArbitrationCandidate& lhs, const ArbitrationCandidate& rhs) { - const auto* lhsReclaimer = lhs.participant->pool()->reclaimer(); - const auto* rhsReclaimer = rhs.participant->pool()->reclaimer(); - if (lhsReclaimer == nullptr || rhsReclaimer == nullptr) { - // Participants without reclaimer are treated as low priority, putting - // them in front. - return (lhsReclaimer == nullptr) > (rhsReclaimer == nullptr); + [priorityWeight]( + const ArbitrationCandidate& lhs, const ArbitrationCandidate& rhs) { + const auto lhsScore = abortBadnessScore(lhs, priorityWeight); + const auto rhsScore = abortBadnessScore(rhs, priorityWeight); + if (lhsScore != rhsScore) { + return lhsScore > rhsScore; } - return lhsReclaimer->priority() > rhsReclaimer->priority(); + return lhs.participant->id() > rhs.participant->id(); }); - - std::vector> candidateGroups; - std::optional prevPriority; - for (auto i = 0; i < candidates.size(); ++i) { - const auto* curReclaimer = candidates[i].participant->pool()->reclaimer(); - const auto curPriority = curReclaimer == nullptr - ? std::nullopt - : std::optional(curReclaimer->priority()); - if (i == 0) { - prevPriority = curPriority; - candidateGroups.emplace_back( - std::vector{std::move(candidates[i])}); - continue; - } - if (curPriority != prevPriority) { - prevPriority = curPriority; - candidateGroups.emplace_back( - std::vector{std::move(candidates[i])}); - } else { - candidateGroups.back().push_back(std::move(candidates[i])); - } - } - return candidateGroups; + return std::move(candidates); } std::optional SharedArbitrator::findAbortCandidate( @@ -698,34 +699,15 @@ std::optional SharedArbitrator::findAbortCandidate( return std::nullopt; } - auto candidateGroups = sortAndGroupAbortCandidates(std::move(candidates)); + const uint64_t priorityWeightBytes = + static_cast(abortPriorityWeightPct_ * capacity_); + auto sorted = sortAbortCandidates(std::move(candidates), priorityWeightBytes); - for (auto& candidateGroup : candidateGroups) { - for (uint64_t capacityLimit : globalArbitrationAbortCapacityLimits_) { - int32_t candidateIdx{-1}; - for (int32_t i = 0; i < candidateGroup.size(); ++i) { - if (candidateGroup[i].participant->aborted()) { - continue; - } - if (candidateGroup[i].currentCapacity < capacityLimit || - candidateGroup[i].currentCapacity == 0) { - continue; - } - if (candidateIdx == -1) { - candidateIdx = i; - continue; - } - // With the same capacity size bucket, we favor the old participant to - // not to be killed, to let long running query proceed first. - if (candidateGroup[candidateIdx].participant->id() < - candidateGroup[i].participant->id()) { - candidateIdx = i; - } - } - if (candidateIdx != -1) { - return candidateGroup[candidateIdx]; - } + for (auto& candidate : sorted) { + if (candidate.participant->aborted() || candidate.currentCapacity == 0) { + continue; } + return candidate; } if (!force) { @@ -733,23 +715,8 @@ std::optional SharedArbitrator::findAbortCandidate( return std::nullopt; } - // Can't find an eligible abort candidate and then return the youngest - // candidate (which has the largest participant id) in the lowest priority - // bucket. - VELOX_CHECK(!candidateGroups.empty() && !candidateGroups[0].empty()); - int32_t candidateIdx{0}; - for (auto i = 0; i < candidateGroups[0].size(); ++i) { - if (candidateGroups[0][i].participant->id() > - candidateGroups[0][candidateIdx].participant->id()) { - candidateIdx = i; - } - } - - VELOX_MEM_LOG(WARNING) - << "Can't find an eligible abort victim and force to abort the youngest " - "participant " - << candidateGroups[0][candidateIdx].participant->name(); - return candidateGroups[0][candidateIdx]; + VELOX_CHECK(!sorted.empty()); + return sorted.front(); } void SharedArbitrator::updateArbitrationRequestStats() { diff --git a/velox/common/memory/SharedArbitrator.h b/velox/common/memory/SharedArbitrator.h index e9ab03e6b5b..95ea9a87385 100644 --- a/velox/common/memory/SharedArbitrator.h +++ b/velox/common/memory/SharedArbitrator.h @@ -242,6 +242,16 @@ class SharedArbitrator : public memory::MemoryArbitrator { static bool globalArbitrationWithoutSpill( const std::unordered_map& configs); + /// What one unit of priority difference is worth in the abort badness + /// score, as a fraction of total capacity (so sensitivity is consistent + /// across deployment sizes). Higher makes priority dominate; the default + /// 1.0 keeps the priority-dominant ordering. + static constexpr std::string_view kAbortPriorityWeightPct{ + "abort-priority-weight-pct"}; + static constexpr double kDefaultAbortPriorityWeightPct{1.0}; + static double abortPriorityWeightPct( + const std::unordered_map& configs); + /// If true, do sanity check on the arbitrator state on destruction. /// /// TODO: deprecate this flag after all the existing memory leak use cases @@ -506,11 +516,12 @@ class SharedArbitrator : public memory::MemoryArbitrator { std::vector> sortAndGroupSpillCandidates( std::vector&& candidates); - // Sorts 'candidates' based on participant's reclaimer priority in descending - // order, putting lower priority ones (with higher priority value) first, and - // high priority ones (with lower priority value) later. - static std::vector> - sortAndGroupAbortCandidates(std::vector&& candidates); + /// Sorts 'candidates' by abort badness score (priority blended with capacity) + /// in descending order. Higher score = better victim. Tie-break: younger + /// participant (larger id) first. + static std::vector sortAbortCandidates( + std::vector&& candidates, + uint64_t priorityWeight); // Finds the participant victim to abort to free used memory based on the // participant's memory capacity and age. The function returns std::nullopt if @@ -639,6 +650,7 @@ class SharedArbitrator : public memory::MemoryArbitrator { const uint32_t globalArbitrationMemoryReclaimPct_; const double globalArbitrationAbortTimeRatio_; const bool globalArbitrationWithoutSpill_; + const double abortPriorityWeightPct_; // The executor used to reclaim memory from multiple participants in parallel // at the background for global arbitration or external memory reclamation. diff --git a/velox/common/memory/tests/CMakeLists.txt b/velox/common/memory/tests/CMakeLists.txt index d797e39e5ed..f2e9a81794b 100644 --- a/velox/common/memory/tests/CMakeLists.txt +++ b/velox/common/memory/tests/CMakeLists.txt @@ -26,6 +26,7 @@ add_executable( CustomMemoryRegistrationTest.cpp HashStringAllocatorTest.cpp MemoryAllocatorTest.cpp + MemoryArbitrationAbortScoringTest.cpp MemoryArbitratorTest.cpp MemoryCapExceededTest.cpp MemoryManagerTest.cpp diff --git a/velox/common/memory/tests/MemoryArbitrationAbortScoringTest.cpp b/velox/common/memory/tests/MemoryArbitrationAbortScoringTest.cpp new file mode 100644 index 00000000000..8c36c1cd986 --- /dev/null +++ b/velox/common/memory/tests/MemoryArbitrationAbortScoringTest.cpp @@ -0,0 +1,387 @@ +/* + * Copyright (c) Facebook, Inc. and its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/common/memory/MallocAllocator.h" +#include "velox/common/memory/Memory.h" +#include "velox/common/memory/SharedArbitrator.h" +#include "velox/core/QueryCtx.h" + +namespace facebook::velox::memory { +namespace { + +class MemoryArbitrationAbortScoringTest : public testing::Test { + protected: + static void SetUpTestCase() { + SharedArbitrator::registerFactory(); + } + + static void TearDownTestCase() { + SharedArbitrator::unregisterFactory(); + } + + std::unique_ptr makeManager( + int64_t capacity, + double priorityWeightPct) { + MemoryManager::Options options; + options.allocatorCapacity = capacity; + options.arbitratorCapacity = capacity; + options.arbitratorKind = "SHARED"; + options.checkUsageLeak = true; + using ExtraConfig = SharedArbitrator::ExtraConfig; + options.extraArbitratorConfigs = { + {std::string(ExtraConfig::kMemoryPoolInitialCapacity), "0B"}, + {std::string(ExtraConfig::kMemoryPoolReservedCapacity), "0B"}, + {std::string(ExtraConfig::kReservedCapacity), "0B"}, + {std::string(ExtraConfig::kGlobalArbitrationEnabled), "true"}, + {std::string(ExtraConfig::kGlobalArbitrationWithoutSpill), "true"}, + {std::string(ExtraConfig::kMemoryPoolAbortCapacityLimit), + folly::to(capacity) + "B"}, + {std::string(ExtraConfig::kAbortPriorityWeightPct), + folly::to(priorityWeightPct)}}; + return std::make_unique(options); + } + + std::shared_ptr makeQueryCtx( + MemoryManager* manager, + int64_t capacity, + int32_t priority, + const std::string& id) { + std::unordered_map config{ + {core::QueryConfig::kQueryMemoryReclaimerPriority, + folly::to(priority)}}; + return core::QueryCtx::Builder() + .executor(executor_.get()) + .pool(manager->addRootPool(id, capacity)) + .queryConfig(core::QueryConfig{std::move(config)}) + .queryId(id) + .build(); + } + + std::shared_ptr executor_{ + std::make_shared(4)}; +}; + +TEST_F(MemoryArbitrationAbortScoringTest, hugeHigherPriorityAbortedFirst) { + const int64_t capacity = 512 << 20; + auto manager = makeManager(capacity, 0.0625); + auto* arbitrator = static_cast(manager->arbitrator()); + + auto hugeQuery = + makeQueryCtx(manager.get(), capacity, /*priority=*/1, "huge"); + auto hugeLeaf = hugeQuery->pool()->addLeafChild("huge.leaf"); + void* hugeBuf = hugeLeaf->allocate(448 << 20); + + auto tinyQuery = + makeQueryCtx(manager.get(), capacity, /*priority=*/2, "tiny"); + auto tinyLeaf = tinyQuery->pool()->addLeafChild("tiny.leaf"); + void* tinyBuf = tinyLeaf->allocate(64 << 20); + + arbitrator->shrinkCapacity(64 << 20, /*allowSpill=*/false, /*force=*/true); + + EXPECT_TRUE(hugeQuery->pool()->aborted()); + EXPECT_FALSE(tinyQuery->pool()->aborted()); + + hugeLeaf->free(hugeBuf, 448 << 20); + tinyLeaf->free(tinyBuf, 64 << 20); +} + +TEST_F(MemoryArbitrationAbortScoringTest, highWeightAbortsLowestPriority) { + const int64_t capacity = 512 << 20; + auto manager = makeManager(capacity, 1.0); + auto* arbitrator = static_cast(manager->arbitrator()); + + auto hugeQuery = + makeQueryCtx(manager.get(), capacity, /*priority=*/1, "huge"); + auto hugeLeaf = hugeQuery->pool()->addLeafChild("huge.leaf"); + void* hugeBuf = hugeLeaf->allocate(448 << 20); + + auto tinyQuery = + makeQueryCtx(manager.get(), capacity, /*priority=*/2, "tiny"); + auto tinyLeaf = tinyQuery->pool()->addLeafChild("tiny.leaf"); + void* tinyBuf = tinyLeaf->allocate(64 << 20); + + arbitrator->shrinkCapacity(64 << 20, /*allowSpill=*/false, /*force=*/true); + + EXPECT_TRUE(tinyQuery->pool()->aborted()); + EXPECT_FALSE(hugeQuery->pool()->aborted()); + + hugeLeaf->free(hugeBuf, 448 << 20); + tinyLeaf->free(tinyBuf, 64 << 20); +} + +TEST_F(MemoryArbitrationAbortScoringTest, largeWeightDegeneratesToPriority) { + const int64_t capacity = 512 << 20; + auto manager = makeManager(capacity, 1.0); + auto* arbitrator = static_cast(manager->arbitrator()); + + auto hugeQuery = + makeQueryCtx(manager.get(), capacity, /*priority=*/1, "huge"); + auto hugeLeaf = hugeQuery->pool()->addLeafChild("huge.leaf"); + void* hugeBuf = hugeLeaf->allocate(448 << 20); + + auto tinyQuery = + makeQueryCtx(manager.get(), capacity, /*priority=*/2, "tiny"); + auto tinyLeaf = tinyQuery->pool()->addLeafChild("tiny.leaf"); + void* tinyBuf = tinyLeaf->allocate(64 << 20); + + arbitrator->shrinkCapacity(64 << 20, /*allowSpill=*/false, /*force=*/true); + + EXPECT_TRUE(tinyQuery->pool()->aborted()); + EXPECT_FALSE(hugeQuery->pool()->aborted()); + + hugeLeaf->free(hugeBuf, 448 << 20); + tinyLeaf->free(tinyBuf, 64 << 20); +} + +TEST_F(MemoryArbitrationAbortScoringTest, equalPriorityAbortsLargest) { + const int64_t capacity = 512 << 20; + auto manager = makeManager(capacity, 0.0625); + auto* arbitrator = static_cast(manager->arbitrator()); + + auto bigQuery = makeQueryCtx(manager.get(), capacity, /*priority=*/1, "big"); + auto bigLeaf = bigQuery->pool()->addLeafChild("big.leaf"); + void* bigBuf = bigLeaf->allocate(320 << 20); + + auto smallQuery = + makeQueryCtx(manager.get(), capacity, /*priority=*/1, "small"); + auto smallLeaf = smallQuery->pool()->addLeafChild("small.leaf"); + void* smallBuf = smallLeaf->allocate(192 << 20); + + arbitrator->shrinkCapacity(64 << 20, /*allowSpill=*/false, /*force=*/true); + + EXPECT_TRUE(bigQuery->pool()->aborted()); + EXPECT_FALSE(smallQuery->pool()->aborted()); + + bigLeaf->free(bigBuf, 320 << 20); + smallLeaf->free(smallBuf, 192 << 20); +} + +TEST_F(MemoryArbitrationAbortScoringTest, multiQueryScoreOrder) { + const int64_t capacity = 512 << 20; + auto manager = makeManager(capacity, 0.0625); + auto* arbitrator = static_cast(manager->arbitrator()); + + auto midPrioHuge = + makeQueryCtx(manager.get(), capacity, /*priority=*/1, "midPrioHuge"); + auto hugeLeaf = midPrioHuge->pool()->addLeafChild("huge.leaf"); + void* hugeBuf = hugeLeaf->allocate(300 << 20); + + auto lowPrioMid = + makeQueryCtx(manager.get(), capacity, /*priority=*/2, "lowPrioMid"); + auto midLeaf = lowPrioMid->pool()->addLeafChild("mid.leaf"); + void* midBuf = midLeaf->allocate(160 << 20); + + auto highPrioTiny = + makeQueryCtx(manager.get(), capacity, /*priority=*/0, "highPrioTiny"); + auto tinyLeaf = highPrioTiny->pool()->addLeafChild("tiny.leaf"); + void* tinyBuf = tinyLeaf->allocate(40 << 20); + + arbitrator->shrinkCapacity(32 << 20, /*allowSpill=*/false, /*force=*/true); + EXPECT_TRUE(midPrioHuge->pool()->aborted()); + EXPECT_FALSE(lowPrioMid->pool()->aborted()); + EXPECT_FALSE(highPrioTiny->pool()->aborted()); + + arbitrator->shrinkCapacity(32 << 20, /*allowSpill=*/false, /*force=*/true); + EXPECT_TRUE(lowPrioMid->pool()->aborted()); + EXPECT_FALSE(highPrioTiny->pool()->aborted()); + + hugeLeaf->free(hugeBuf, 300 << 20); + midLeaf->free(midBuf, 160 << 20); + tinyLeaf->free(tinyBuf, 40 << 20); +} + +TEST_F(MemoryArbitrationAbortScoringTest, equalScorePrefersAbortingYounger) { + const int64_t capacity = 512 << 20; + auto manager = makeManager(capacity, 0.0625); + auto* arbitrator = static_cast(manager->arbitrator()); + + auto older = makeQueryCtx(manager.get(), capacity, /*priority=*/1, "older"); + auto olderLeaf = older->pool()->addLeafChild("older.leaf"); + void* olderBuf = olderLeaf->allocate(200 << 20); + + auto younger = + makeQueryCtx(manager.get(), capacity, /*priority=*/1, "younger"); + auto youngerLeaf = younger->pool()->addLeafChild("younger.leaf"); + void* youngerBuf = youngerLeaf->allocate(200 << 20); + + arbitrator->shrinkCapacity(32 << 20, /*allowSpill=*/false, /*force=*/true); + + EXPECT_TRUE(younger->pool()->aborted()); + EXPECT_FALSE(older->pool()->aborted()); + + olderLeaf->free(olderBuf, 200 << 20); + youngerLeaf->free(youngerBuf, 200 << 20); +} + +TEST_F(MemoryArbitrationAbortScoringTest, smallPriorityGapNotAbsoluteShield) { + const int64_t capacity = 1LL << 30; + + auto run = [&](double weightPct) { + auto manager = makeManager(capacity, weightPct); + auto* arbitrator = static_cast(manager->arbitrator()); + + auto bigEtl = + makeQueryCtx(manager.get(), capacity, /*priority=*/1, "big_etl"); + auto bigLeaf = bigEtl->pool()->addLeafChild("big_etl.leaf"); + void* bigBuf = bigLeaf->allocate(600 << 20); + + auto q1 = makeQueryCtx(manager.get(), capacity, /*priority=*/2, "q1"); + auto q1Leaf = q1->pool()->addLeafChild("q1.leaf"); + void* q1Buf = q1Leaf->allocate(100 << 20); + + auto q2 = makeQueryCtx(manager.get(), capacity, /*priority=*/2, "q2"); + auto q2Leaf = q2->pool()->addLeafChild("q2.leaf"); + void* q2Buf = q2Leaf->allocate(80 << 20); + + auto q3 = makeQueryCtx(manager.get(), capacity, /*priority=*/2, "q3"); + auto q3Leaf = q3->pool()->addLeafChild("q3.leaf"); + void* q3Buf = q3Leaf->allocate(70 << 20); + + auto q4 = makeQueryCtx(manager.get(), capacity, /*priority=*/2, "q4"); + auto q4Leaf = q4->pool()->addLeafChild("q4.leaf"); + void* q4Buf = q4Leaf->allocate(50 << 20); + + auto q5 = makeQueryCtx(manager.get(), capacity, /*priority=*/2, "q5"); + auto q5Leaf = q5->pool()->addLeafChild("q5.leaf"); + void* q5Buf = q5Leaf->allocate(40 << 20); + + arbitrator->shrinkCapacity(300 << 20, /*allowSpill=*/false, /*force=*/true); + + bool bigAborted = bigEtl->pool()->aborted(); + bool allSmallSurvived = !q1->pool()->aborted() && !q2->pool()->aborted() && + !q3->pool()->aborted() && !q4->pool()->aborted() && + !q5->pool()->aborted(); + + bigLeaf->free(bigBuf, 600 << 20); + q1Leaf->free(q1Buf, 100 << 20); + q2Leaf->free(q2Buf, 80 << 20); + q3Leaf->free(q3Buf, 70 << 20); + q4Leaf->free(q4Buf, 50 << 20); + q5Leaf->free(q5Buf, 40 << 20); + + return std::make_pair(bigAborted, allSmallSurvived); + }; + + auto [highWeightBigAborted, highWeightSmallSurvived] = run(1.0); + EXPECT_FALSE(highWeightBigAborted); + EXPECT_FALSE(highWeightSmallSurvived); + + auto [lowWeightBigAborted, lowWeightSmallSurvived] = run(0.0625); + EXPECT_TRUE(lowWeightBigAborted); + EXPECT_TRUE(lowWeightSmallSurvived); +} + +TEST_F(MemoryArbitrationAbortScoringTest, wastedAbortsCascade) { + const int64_t capacity = 1LL << 30; + + auto runWithWeight = [&](double weightPct) { + auto manager = makeManager(capacity, weightPct); + auto* arbitrator = static_cast(manager->arbitrator()); + + struct QState { + std::shared_ptr ctx; + std::shared_ptr leaf; + void* buf; + int64_t size; + }; + std::vector queries; + + auto addQuery = [&](const std::string& id, int32_t priority, int64_t mb) { + auto ctx = makeQueryCtx(manager.get(), capacity, priority, id); + auto leaf = ctx->pool()->addLeafChild(id + ".leaf"); + void* buf = leaf->allocate(mb << 20); + queries.push_back({ctx, leaf, buf, mb << 20}); + }; + + addQuery("big_etl", 1, 600); + addQuery("q1", 2, 100); + addQuery("q2", 2, 80); + addQuery("q3", 2, 70); + addQuery("q4", 2, 50); + addQuery("q5", 2, 40); + + auto statsBefore = arbitrator->stats(); + + uint64_t totalReclaimed = 0; + const uint64_t target = 300UL << 20; + while (totalReclaimed < target) { + uint64_t reclaimed = + arbitrator->shrinkCapacity(target - totalReclaimed, false, true); + if (reclaimed == 0) { + break; + } + totalReclaimed += reclaimed; + } + + auto statsAfter = arbitrator->stats(); + uint64_t numAborted = statsAfter.numAborted - statsBefore.numAborted; + + uint64_t surviving = 0; + for (auto& q : queries) { + if (!q.ctx->pool()->aborted()) { + surviving++; + } + q.leaf->free(q.buf, q.size); + } + + return std::make_pair(numAborted, surviving); + }; + + auto [highWeightAborts, highWeightSurviving] = runWithWeight(1.0); + auto [lowWeightAborts, lowWeightSurviving] = runWithWeight(0.0625); + + EXPECT_GT(highWeightAborts, lowWeightAborts); + EXPECT_LT(highWeightSurviving, lowWeightSurviving); + + EXPECT_EQ(lowWeightAborts, 1); + EXPECT_EQ(lowWeightSurviving, 5); +} + +TEST_F( + MemoryArbitrationAbortScoringTest, + highWeightPreservesPriorityDominance) { + const int64_t capacity = 1LL << 30; + auto manager = makeManager(capacity, 1.0); + auto* arbitrator = static_cast(manager->arbitrator()); + + auto bigEtl = + makeQueryCtx(manager.get(), capacity, /*priority=*/1, "big_etl"); + auto bigLeaf = bigEtl->pool()->addLeafChild("big_etl.leaf"); + void* bigBuf = bigLeaf->allocate(600 << 20); + + auto q1 = makeQueryCtx(manager.get(), capacity, /*priority=*/2, "q1"); + auto q1Leaf = q1->pool()->addLeafChild("q1.leaf"); + void* q1Buf = q1Leaf->allocate(100 << 20); + + auto q2 = makeQueryCtx(manager.get(), capacity, /*priority=*/2, "q2"); + auto q2Leaf = q2->pool()->addLeafChild("q2.leaf"); + void* q2Buf = q2Leaf->allocate(80 << 20); + + arbitrator->shrinkCapacity(64 << 20, /*allowSpill=*/false, /*force=*/true); + + EXPECT_FALSE(bigEtl->pool()->aborted()); + + bigLeaf->free(bigBuf, 600 << 20); + q1Leaf->free(q1Buf, 100 << 20); + q2Leaf->free(q2Buf, 80 << 20); +} + +} // namespace +} // namespace facebook::velox::memory diff --git a/velox/docs/develop/memory.rst b/velox/docs/develop/memory.rst index 94437cb90c9..c540d213df8 100644 --- a/velox/docs/develop/memory.rst +++ b/velox/docs/develop/memory.rst @@ -623,11 +623,16 @@ The end-to-end memory arbitration process in *SharedArbitrator* works as follows reclaimed memory to the requestor pool by increasing its memory capacity (*MemoryPool::grow*). If not, the memory arbitrator has to call *SharedArbitrator::handleOOM* to send the memory pool abort - (*MemoryPool::abort*) request to the candidate memory pool with the largest - capacity as victim to free up memory to let the other running queries - with enough memory proceed. The memory pool abort fails the query - execution and waits for its completion to release all the held memory - resources. + (*MemoryPool::abort*) request to a victim memory pool to free up memory to + let the other running queries with enough memory proceed. The victim is + chosen by a badness score that blends reclaimer priority with capacity: + ``score = currentCapacity + priority * (abort-priority-weight-pct * + arbitratorCapacity)``. A higher score means better victim. The + *abort-priority-weight-pct* config (default 1.0) controls how strongly + priority dominates: a large value keeps priority dominant, a small value + lets capacity matter. When scores are equal, the younger participant is + aborted first. The memory pool abort fails the query execution and waits + for its completion to release all the held memory resources. f. If the victim query pool is the requestor pool itself, then memory arbitration fails. Otherwise, go back to step-6-c to retry the memory