From 0a6a01822653b97bca02d2e2fb7b51ac4996c621 Mon Sep 17 00:00:00 2001 From: Jialiang Tan Date: Thu, 21 Mar 2024 17:17:24 -0700 Subject: [PATCH] Enable arbitrator in OperatorTestBase by default (#9141) Summary: Enable arbitrator by default for all test suites that extend from OperatorTestBase. Pull Request resolved: https://github.com/facebookincubator/velox/pull/9141 Reviewed By: xiaoxmeng Differential Revision: D55044372 Pulled By: tanjialiang fbshipit-source-id: 6b2f1b752dd91ba7f976935696536bd02924b405 --- .../memory/tests/MemoryCapExceededTest.cpp | 8 ++-- .../hive/tests/HiveDataSinkTest.cpp | 39 +++++++++---------- velox/exec/tests/AggregationTest.cpp | 6 --- velox/exec/tests/GroupedExecutionTest.cpp | 1 - velox/exec/tests/HashJoinTest.cpp | 10 ----- velox/exec/tests/RowNumberTest.cpp | 10 ----- velox/exec/tests/TaskTest.cpp | 10 ----- velox/exec/tests/utils/OperatorTestBase.cpp | 14 ++----- velox/exec/tests/utils/OperatorTestBase.h | 2 - .../tests/utils/AggregationTestBase.cpp | 10 ----- .../tests/utils/AggregationTestBase.h | 2 - 11 files changed, 27 insertions(+), 85 deletions(-) diff --git a/velox/common/memory/tests/MemoryCapExceededTest.cpp b/velox/common/memory/tests/MemoryCapExceededTest.cpp index 7157a5e3aacbf..6960d751857d6 100644 --- a/velox/common/memory/tests/MemoryCapExceededTest.cpp +++ b/velox/common/memory/tests/MemoryCapExceededTest.cpp @@ -96,8 +96,8 @@ TEST_P(MemoryCapExceededTest, singleDriver) { .orderBy({"c0"}, false) .planNode(); auto queryCtx = std::make_shared(executor_.get()); - queryCtx->testingOverrideMemoryPool( - memory::memoryManager()->addRootPool(queryCtx->queryId(), kMaxBytes)); + queryCtx->testingOverrideMemoryPool(memory::memoryManager()->addRootPool( + queryCtx->queryId(), kMaxBytes, exec::MemoryReclaimer::create())); CursorParameters params; params.planNode = plan; params.queryCtx = queryCtx; @@ -154,8 +154,8 @@ TEST_P(MemoryCapExceededTest, multipleDrivers) { .singleAggregation({"c0"}, {"sum(c1)"}) .planNode(); auto queryCtx = std::make_shared(executor_.get()); - queryCtx->testingOverrideMemoryPool( - memory::memoryManager()->addRootPool(queryCtx->queryId(), kMaxBytes)); + queryCtx->testingOverrideMemoryPool(memory::memoryManager()->addRootPool( + queryCtx->queryId(), kMaxBytes, exec::MemoryReclaimer::create())); const int32_t numDrivers = 10; CursorParameters params; diff --git a/velox/connectors/hive/tests/HiveDataSinkTest.cpp b/velox/connectors/hive/tests/HiveDataSinkTest.cpp index 34528f18a249e..7695c4e42993b 100644 --- a/velox/connectors/hive/tests/HiveDataSinkTest.cpp +++ b/velox/connectors/hive/tests/HiveDataSinkTest.cpp @@ -565,8 +565,8 @@ TEST_F(HiveDataSinkTest, abort) { } TEST_F(HiveDataSinkTest, memoryReclaim) { - const int numBatches = 20; - auto vectors = createVectors(500, 20); + const int numBatches = 200; + auto vectors = createVectors(500, 200); struct { dwio::common::FileFormat format; @@ -587,7 +587,7 @@ TEST_F(HiveDataSinkTest, memoryReclaim) { expectedWriterReclaimed); } } testSettings[] = { - //{dwio::common::FileFormat::DWRF, true, true, 1 << 30, true, true}, + // {dwio::common::FileFormat::DWRF, true, true, 1 << 30, true, true}, {dwio::common::FileFormat::DWRF, true, true, 1, true, true}, {dwio::common::FileFormat::DWRF, true, false, 1 << 30, false, false}, {dwio::common::FileFormat::DWRF, true, false, 1, false, false}, @@ -670,30 +670,29 @@ TEST_F(HiveDataSinkTest, memoryReclaim) { for (int i = 0; i < numBatches; ++i) { dataSink->appendData(vectors[i]); } - memory::MemoryReclaimer::Stats stats; + memory::MemoryArbitrator::Stats oldStats = + memory::memoryManager()->arbitrator()->stats(); uint64_t reclaimableBytes{0}; if (testData.expectedWriterReclaimed) { reclaimableBytes = root_->reclaimableBytes().value(); ASSERT_GT(reclaimableBytes, 0); - ASSERT_GT(root_->reclaim(256L << 20, 0, stats), 0); - ASSERT_GT(stats.reclaimExecTimeUs, 0); - ASSERT_GT(stats.reclaimedBytes, 0); + memory::testingRunArbitration(); + memory::MemoryArbitrator::Stats curStats = + memory::memoryManager()->arbitrator()->stats(); + ASSERT_GT(curStats.reclaimTimeUs - oldStats.reclaimTimeUs, 0); + ASSERT_GT(curStats.numReclaimedBytes - oldStats.numReclaimedBytes, 0); // We expect dwrf writer set numNonReclaimableAttempts counter. - ASSERT_LE(stats.numNonReclaimableAttempts, 1); + ASSERT_LE( + curStats.numNonReclaimableAttempts - + oldStats.numNonReclaimableAttempts, + 1); } else { ASSERT_FALSE(root_->reclaimableBytes().has_value()); - ASSERT_EQ(root_->reclaim(256L << 20, 0, stats), 0); - ASSERT_EQ(stats.reclaimExecTimeUs, 0); - ASSERT_EQ(stats.reclaimedBytes, 0); - if (testData.expectedWriterReclaimEnabled) { - if (testData.sortWriter) { - ASSERT_GE(stats.numNonReclaimableAttempts, 1); - } else { - ASSERT_EQ(stats.numNonReclaimableAttempts, 1); - } - } else { - ASSERT_EQ(stats.numNonReclaimableAttempts, 0); - } + memory::testingRunArbitration(); + memory::MemoryArbitrator::Stats curStats = + memory::memoryManager()->arbitrator()->stats(); + ASSERT_EQ(curStats.reclaimTimeUs - oldStats.reclaimTimeUs, 0); + ASSERT_EQ(curStats.numReclaimedBytes - oldStats.numReclaimedBytes, 0); } const auto partitions = dataSink->close(); if (testData.sortWriter && testData.expectedWriterReclaimed) { diff --git a/velox/exec/tests/AggregationTest.cpp b/velox/exec/tests/AggregationTest.cpp index 0bafdd8b62ea9..4e919bcfa3d42 100644 --- a/velox/exec/tests/AggregationTest.cpp +++ b/velox/exec/tests/AggregationTest.cpp @@ -138,16 +138,10 @@ void checkSpillStats(PlanNodeStats& stats, bool expectedSpill) { class AggregationTest : public OperatorTestBase { protected: static void SetUpTestCase() { - FLAGS_velox_testing_enable_arbitration = true; OperatorTestBase::SetUpTestCase(); TestValue::enable(); } - static void TearDownTestCase() { - FLAGS_velox_testing_enable_arbitration = false; - OperatorTestBase::TearDownTestCase(); - } - void SetUp() override { OperatorTestBase::SetUp(); filesystems::registerLocalFileSystem(); diff --git a/velox/exec/tests/GroupedExecutionTest.cpp b/velox/exec/tests/GroupedExecutionTest.cpp index 06cd5c2ddbe12..ed4c8c2647429 100644 --- a/velox/exec/tests/GroupedExecutionTest.cpp +++ b/velox/exec/tests/GroupedExecutionTest.cpp @@ -34,7 +34,6 @@ class GroupedExecutionTest : public virtual HiveConnectorTestBase { } static void SetUpTestCase() { - FLAGS_velox_testing_enable_arbitration = true; HiveConnectorTestBase::SetUpTestCase(); } diff --git a/velox/exec/tests/HashJoinTest.cpp b/velox/exec/tests/HashJoinTest.cpp index eec0053650d30..abca47ad7f22b 100644 --- a/velox/exec/tests/HashJoinTest.cpp +++ b/velox/exec/tests/HashJoinTest.cpp @@ -749,16 +749,6 @@ class HashJoinTest : public HiveConnectorTestBase { explicit HashJoinTest(const TestParam& param) : numDrivers_(param.numDrivers) {} - static void SetUpTestCase() { - FLAGS_velox_testing_enable_arbitration = true; - OperatorTestBase::SetUpTestCase(); - } - - static void TearDownTestCase() { - FLAGS_velox_testing_enable_arbitration = false; - OperatorTestBase::TearDownTestCase(); - } - void SetUp() override { HiveConnectorTestBase::SetUp(); diff --git a/velox/exec/tests/RowNumberTest.cpp b/velox/exec/tests/RowNumberTest.cpp index 6ec773d2c4774..39c1c8d90b207 100644 --- a/velox/exec/tests/RowNumberTest.cpp +++ b/velox/exec/tests/RowNumberTest.cpp @@ -26,16 +26,6 @@ namespace facebook::velox::exec::test { class RowNumberTest : public OperatorTestBase { protected: - static void SetUpTestCase() { - FLAGS_velox_testing_enable_arbitration = true; - OperatorTestBase::SetUpTestCase(); - } - - static void TearDownTestCase() { - FLAGS_velox_testing_enable_arbitration = false; - OperatorTestBase::TearDownTestCase(); - } - RowNumberTest() { filesystems::registerLocalFileSystem(); rowType_ = ROW( diff --git a/velox/exec/tests/TaskTest.cpp b/velox/exec/tests/TaskTest.cpp index 80357b3a2cd84..8c411f45421a9 100644 --- a/velox/exec/tests/TaskTest.cpp +++ b/velox/exec/tests/TaskTest.cpp @@ -457,16 +457,6 @@ class TestBadMemoryTranslator : public exec::Operator::PlanNodeTranslator { } // namespace class TaskTest : public HiveConnectorTestBase { protected: - static void SetUpTestCase() { - FLAGS_velox_testing_enable_arbitration = true; - OperatorTestBase::SetUpTestCase(); - } - - static void TearDownTestCase() { - FLAGS_velox_testing_enable_arbitration = false; - OperatorTestBase::TearDownTestCase(); - } - static std::pair, std::vector> executeSingleThreaded( core::PlanFragment plan, diff --git a/velox/exec/tests/utils/OperatorTestBase.cpp b/velox/exec/tests/utils/OperatorTestBase.cpp index cf09e11e394a8..b0f9210b4f769 100644 --- a/velox/exec/tests/utils/OperatorTestBase.cpp +++ b/velox/exec/tests/utils/OperatorTestBase.cpp @@ -32,10 +32,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; using namespace facebook::velox::memory; @@ -61,12 +57,10 @@ void OperatorTestBase::SetUpTestCase() { memory::SharedArbitrator::registerFactory(); 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; - } + options.arbitratorCapacity = 6L << 30; + options.arbitratorKind = "SHARED"; + options.checkUsageLeak = true; + options.arbitrationStateCheckCb = memoryArbitrationStateCheck; memory::MemoryManager::testingSetInstance(options); asyncDataCache_ = cache::AsyncDataCache::create(memory::memoryManager()->allocator()); diff --git a/velox/exec/tests/utils/OperatorTestBase.h b/velox/exec/tests/utils/OperatorTestBase.h index a10b316c173f8..54eb0d6b97eb4 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/functions/lib/aggregates/tests/utils/AggregationTestBase.cpp b/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.cpp index 49a5691b1b0f8..b0ef8ca0aa8c4 100644 --- a/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.cpp +++ b/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.cpp @@ -79,16 +79,6 @@ void AggregationTestBase::TearDown() { OperatorTestBase::TearDown(); } -void AggregationTestBase::SetUpTestCase() { - FLAGS_velox_testing_enable_arbitration = true; - OperatorTestBase::SetUpTestCase(); -} - -void AggregationTestBase::TearDownTestCase() { - FLAGS_velox_testing_enable_arbitration = false; - OperatorTestBase::TearDownTestCase(); -} - void AggregationTestBase::testAggregations( const std::vector& data, const std::vector& groupingKeys, diff --git a/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.h b/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.h index 3e22078c32919..0b0968f5d0094 100644 --- a/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.h +++ b/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.h @@ -29,8 +29,6 @@ class AggregationTestBase : public exec::test::OperatorTestBase { protected: void SetUp() override; void TearDown() override; - static void SetUpTestCase(); - static void TearDownTestCase(); std::vector makeVectors(const RowTypePtr& rowType, vector_size_t size, int numVectors);