diff --git a/velox/common/memory/tests/MockSharedArbitratorTest.cpp b/velox/common/memory/tests/MockSharedArbitratorTest.cpp index d9d3b68f68d6..ce90cca420f5 100644 --- a/velox/common/memory/tests/MockSharedArbitratorTest.cpp +++ b/velox/common/memory/tests/MockSharedArbitratorTest.cpp @@ -486,7 +486,9 @@ class MockSharedArbitrationTest : public testing::Test { {std::string(ExtraConfig::kGlobalArbitrationEnabled), folly::to(globalArtbitrationEnabled)}, {std::string(ExtraConfig::kGlobalArbitrationWithoutSpill), - folly::to(globalArbitrationWithoutSpill)}}; + folly::to(globalArbitrationWithoutSpill)}, + {std::string(ExtraConfig::kGlobalArbitrationAbortTimeRatio), + folly::to(globalArbitrationAbortTimeRatio)}}; options.arbitrationStateCheckCb = std::move(arbitrationStateCheckCb); options.checkUsageLeak = true; manager_ = std::make_unique(options); @@ -858,6 +860,26 @@ TEST_F(MockSharedArbitrationTest, asyncArbitrationWork) { // Test different kinds of arbitraton failures. TEST_F(MockSharedArbitrationTest, arbitrationFailures) { + // Set the globalArbitrationAbortTimeRatio to be very small so that the query + // can be aborted sooner and the test would not timeout. + setupMemory( + kMemoryCapacity, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + kMemoryReclaimThreadsHwMultiplier, + nullptr, + true, + 5 * 60 * 1'000'000'000UL, + false, + 0.005); // Local arbitration failure with exceeded capacity limit. { auto task = addTask(64 * MB); @@ -3027,8 +3049,27 @@ TEST_F(MockSharedArbitrationTest, minReclaimBytes) { for (const auto& testData : testSettings) { SCOPED_TRACE(testData.debugString()); - // Make simple settings to focus shrink capacity logic testing. - setupMemory(memoryCapacity, 0, 0, 0, 0, 0, 0, 0, testData.minReclaimBytes); + // Make simple settings to focus shrink capacity logic testing. Set the + // globalArbitrationAbortTimeRatio to be very small so that the query + // can be aborted sooner and the test would not timeout. + setupMemory( + memoryCapacity, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + testData.minReclaimBytes, + 0, + 0, + kMemoryReclaimThreadsHwMultiplier, + nullptr, + true, + 5 * 60 * 1'000'000'000UL, + false, + 0.005); std::vector taskContainers; for (const auto& testTask : testData.testTasks) { auto task = addTask(); @@ -3903,8 +3944,26 @@ TEST_F(MockSharedArbitrationTest, arbitrationFailure) { for (const auto& testData : testSettings) { SCOPED_TRACE(testData.debugString()); - - setupMemory(maxCapacity, 0, initialCapacity, 0); + // Set the globalArbitrationAbortTimeRatio to be very small so that the + // query can be aborted sooner and the test would not timeout. + setupMemory( + maxCapacity, + 0, + initialCapacity, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + kMemoryReclaimThreadsHwMultiplier, + nullptr, + true, + 5 * 60 * 1'000'000'000UL, + false, + 0.005); std::shared_ptr requestorTask = addTask(); MockMemoryOperator* requestorOp = addMemoryOp(requestorTask, false); requestorOp->allocate(testData.requestorCapacity);