diff --git a/velox/common/memory/MemoryArbitrator.cpp b/velox/common/memory/MemoryArbitrator.cpp index cda854de711d..1d664d1a0ac1 100644 --- a/velox/common/memory/MemoryArbitrator.cpp +++ b/velox/common/memory/MemoryArbitrator.cpp @@ -457,18 +457,4 @@ MemoryArbitrationContext* memoryArbitrationContext() { bool underMemoryArbitration() { return memoryArbitrationContext() != nullptr; } - -void testingRunArbitration(uint64_t targetBytes, MemoryManager* manager) { - if (manager == nullptr) { - manager = memory::memoryManager(); - } - manager->shrinkPools(targetBytes); -} - -void testingRunArbitration(MemoryPool* pool, uint64_t targetBytes) { - pool->enterArbitration(); - static_cast(pool)->testingManager()->shrinkPools( - targetBytes); - pool->leaveArbitration(); -} } // namespace facebook::velox::memory diff --git a/velox/common/memory/MemoryArbitrator.h b/velox/common/memory/MemoryArbitrator.h index 0f675a6dec37..20371f6bdc26 100644 --- a/velox/common/memory/MemoryArbitrator.h +++ b/velox/common/memory/MemoryArbitrator.h @@ -383,18 +383,4 @@ MemoryArbitrationContext* memoryArbitrationContext(); /// Returns true if the running thread is under memory arbitration or not. bool underMemoryArbitration(); - -/// The function triggers memory arbitration by shrinking memory pools from -/// 'manager' by invoking shrinkPools API. If 'manager' is not set, then it -/// shrinks from the process wide memory manager. If 'targetBytes' is zero, then -/// reclaims all the memory from 'manager' if possible. -class MemoryManager; -void testingRunArbitration( - uint64_t targetBytes = 0, - MemoryManager* manager = nullptr); - -/// The function triggers memory arbitration by shrinking memory pools from -/// 'manager' of 'pool' by invoking its shrinkPools API. If 'targetBytes' is -/// zero, then reclaims all the memory from 'manager' if possible. -void testingRunArbitration(MemoryPool* pool, uint64_t targetBytes = 0); } // namespace facebook::velox::memory diff --git a/velox/common/memory/MemoryPool.h b/velox/common/memory/MemoryPool.h index e9a55ad15126..841abcc5d9a8 100644 --- a/velox/common/memory/MemoryPool.h +++ b/velox/common/memory/MemoryPool.h @@ -679,10 +679,6 @@ class MemoryPoolImpl : public MemoryPool { void testingSetCapacity(int64_t bytes); - MemoryManager* testingManager() const { - return manager_; - } - MemoryAllocator* testingAllocator() const { return allocator_; } diff --git a/velox/core/QueryCtx.h b/velox/core/QueryCtx.h index e7a011e2eced..dfb33f1df1bb 100644 --- a/velox/core/QueryCtx.h +++ b/velox/core/QueryCtx.h @@ -142,10 +142,8 @@ class QueryCtx { void initPool(const std::string& queryId) { if (pool_ == nullptr) { - pool_ = memory::memoryManager()->addRootPool( - QueryCtx::generatePoolName(queryId), - memory::kMaxMemory, - memory::MemoryReclaimer::create()); + pool_ = memory::deprecatedDefaultMemoryManager().addRootPool( + QueryCtx::generatePoolName(queryId)); } } diff --git a/velox/exec/HashBuild.cpp b/velox/exec/HashBuild.cpp index 3730ce300373..2e3965e17742 100644 --- a/velox/exec/HashBuild.cpp +++ b/velox/exec/HashBuild.cpp @@ -452,7 +452,6 @@ bool HashBuild::ensureInputFits(RowVectorPtr& input) { bool HashBuild::reserveMemory(const RowVectorPtr& input) { VELOX_CHECK(spillEnabled()); - Operator::ReclaimableSectionGuard guard(this); numSpillRows_ = 0; numSpillBytes_ = 0; @@ -469,10 +468,9 @@ bool HashBuild::reserveMemory(const RowVectorPtr& input) { if (numRows != 0) { // Test-only spill path. if (testingTriggerSpill()) { - memory::testingRunArbitration(pool()); - // NOTE: the memory arbitration should have triggered spilling on this - // hash build operator so we return true to indicate have enough memory. - return true; + numSpillRows_ = std::max(1, numRows / 10); + numSpillBytes_ = numSpillRows_ * outOfLineBytesPerRow; + return false; } // We check usage from the parent pool to take peers' allocations into @@ -524,8 +522,11 @@ bool HashBuild::reserveMemory(const RowVectorPtr& input) { incrementBytes * 2, currentUsage * spillConfig_->spillableReservationGrowthPct / 100); - if (pool()->maybeReserve(targetIncrementBytes)) { - return true; + { + Operator::ReclaimableSectionGuard guard(this); + if (pool()->maybeReserve(targetIncrementBytes)) { + return true; + } } LOG(WARNING) << "Failed to reserve " << succinctBytes(targetIncrementBytes) diff --git a/velox/exec/tests/HashJoinTest.cpp b/velox/exec/tests/HashJoinTest.cpp index 794528d7b447..43ea05ecfdd1 100644 --- a/velox/exec/tests/HashJoinTest.cpp +++ b/velox/exec/tests/HashJoinTest.cpp @@ -24,6 +24,7 @@ #include "velox/exec/HashBuild.h" #include "velox/exec/HashJoinBridge.h" #include "velox/exec/PlanNodeStats.h" +#include "velox/exec/TableScan.h" #include "velox/exec/tests/utils/ArbitratorTestUtil.h" #include "velox/exec/tests/utils/AssertQueryBuilder.h" #include "velox/exec/tests/utils/Cursor.h" @@ -733,11 +734,6 @@ class HashJoinBuilder { class HashJoinTest : public HiveConnectorTestBase { protected: - static void SetUpTestCase() { - FLAGS_velox_testing_enable_arbitration = true; - HiveConnectorTestBase::SetUpTestCase(); - } - HashJoinTest() : HashJoinTest(TestParam(1)) {} explicit HashJoinTest(const TestParam& param) diff --git a/velox/exec/tests/utils/ArbitratorTestUtil.cpp b/velox/exec/tests/utils/ArbitratorTestUtil.cpp index 58c5d0af4b25..e72b0743c778 100644 --- a/velox/exec/tests/utils/ArbitratorTestUtil.cpp +++ b/velox/exec/tests/utils/ArbitratorTestUtil.cpp @@ -348,4 +348,21 @@ QueryTestResult runWriteTask( } return result; } + +void testingRunArbitration( + memory::MemoryPool* pool, + uint64_t targetBytes, + memory::MemoryManager* manager) { + if (manager == nullptr) { + manager = memory::memoryManager(); + } + if (pool != nullptr) { + pool->enterArbitration(); + manager->shrinkPools(targetBytes); + pool->leaveArbitration(); + } else { + manager->shrinkPools(targetBytes); + } +} + } // namespace facebook::velox::exec::test diff --git a/velox/exec/tests/utils/ArbitratorTestUtil.h b/velox/exec/tests/utils/ArbitratorTestUtil.h index 3b6610fd12f5..ddbddd232dc4 100644 --- a/velox/exec/tests/utils/ArbitratorTestUtil.h +++ b/velox/exec/tests/utils/ArbitratorTestUtil.h @@ -179,4 +179,15 @@ QueryTestResult runWriteTask( const std::string& kHiveConnectorId, bool enableSpilling, const RowVectorPtr& expectedResult = nullptr); + +/// The function triggers memory arbitration by shrinking memory pools from +/// 'manager' by invoking shrinkPools API. If 'manager' is not set, then it +/// shrinks from the process wide memory manager. If 'pool' is provided, the +/// function puts 'pool' in arbitration state before the arbitration to ease +/// test use. If 'targetBytes' is zero, then reclaims all the memory from +/// 'manager' if possible. +void testingRunArbitration( + memory::MemoryPool* pool = nullptr, + uint64_t targetBytes = 0, + memory::MemoryManager* manager = nullptr); } // namespace facebook::velox::exec::test diff --git a/velox/exec/tests/utils/OperatorTestBase.cpp b/velox/exec/tests/utils/OperatorTestBase.cpp index ac19629621de..4871d58ca325 100644 --- a/velox/exec/tests/utils/OperatorTestBase.cpp +++ b/velox/exec/tests/utils/OperatorTestBase.cpp @@ -34,10 +34,6 @@ DECLARE_bool(velox_memory_leak_check_enabled); DECLARE_bool(velox_enable_memory_usage_track_in_default_memory_pool); -DEFINE_bool( - velox_testing_enable_arbitration, - false, - "Enable to turn on arbitration for tests by default"); using namespace facebook::velox::common::testutil; @@ -62,12 +58,6 @@ void OperatorTestBase::SetUpTestCase() { exec::SharedArbitrator::registerFactory(); memory::MemoryManagerOptions options; options.allocatorCapacity = 8L << 30; - if (FLAGS_velox_testing_enable_arbitration) { - options.arbitratorCapacity = 6L << 30; - options.arbitratorKind = "SHARED"; - options.checkUsageLeak = true; - options.arbitrationStateCheckCb = memoryArbitrationStateCheck; - } memory::MemoryManager::testingSetInstance(options); asyncDataCache_ = cache::AsyncDataCache::create(memoryManager()->allocator()); cache::AsyncDataCache::setInstance(asyncDataCache_.get()); diff --git a/velox/exec/tests/utils/OperatorTestBase.h b/velox/exec/tests/utils/OperatorTestBase.h index a10b316c173f..54eb0d6b97eb 100644 --- a/velox/exec/tests/utils/OperatorTestBase.h +++ b/velox/exec/tests/utils/OperatorTestBase.h @@ -28,8 +28,6 @@ #include "velox/vector/tests/utils/VectorMaker.h" #include "velox/vector/tests/utils/VectorTestBase.h" -DECLARE_bool(velox_testing_enable_arbitration); - namespace facebook::velox::exec::test { class OperatorTestBase : public testing::Test, public velox::test::VectorTestBase { diff --git a/velox/expression/tests/FuzzerRunner.cpp b/velox/expression/tests/FuzzerRunner.cpp index 76d9aff62e46..aef6b786f8da 100644 --- a/velox/expression/tests/FuzzerRunner.cpp +++ b/velox/expression/tests/FuzzerRunner.cpp @@ -210,7 +210,6 @@ int FuzzerRunner::run( void FuzzerRunner::runFromGtest( size_t seed, const std::unordered_set& skipFunctions) { - memory::MemoryManager::testingSetInstance({}); auto signatures = facebook::velox::getFunctionSignatures(); ExpressionFuzzerVerifier( signatures, seed, getExpressionFuzzerVerifierOptions(skipFunctions))