From 894ba7d1f3c1c99e1beab9d0985b2d03cb2b66f9 Mon Sep 17 00:00:00 2001
From: Jialiang Tan <jacob.jialiang.tan@gmail.com>
Date: Tue, 5 Nov 2024 16:36:45 -0800
Subject: [PATCH] Fix test MockSharedArbitratorTest.arbitrationFailure

---
 .../memory/tests/MockSharedArbitratorTest.cpp | 46 +++++++++++++++++--
 1 file changed, 43 insertions(+), 3 deletions(-)

diff --git a/velox/common/memory/tests/MockSharedArbitratorTest.cpp b/velox/common/memory/tests/MockSharedArbitratorTest.cpp
index d9d3b68f68d6..7c99b5ceffa1 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<std::string>(globalArtbitrationEnabled)},
         {std::string(ExtraConfig::kGlobalArbitrationWithoutSpill),
-         folly::to<std::string>(globalArbitrationWithoutSpill)}};
+         folly::to<std::string>(globalArbitrationWithoutSpill)},
+        {std::string(ExtraConfig::kGlobalArbitrationAbortTimeRatio),
+         folly::to<std::string>(globalArbitrationAbortTimeRatio)}};
     options.arbitrationStateCheckCb = std::move(arbitrationStateCheckCb);
     options.checkUsageLeak = true;
     manager_ = std::make_unique<MemoryManager>(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);
@@ -3903,8 +3925,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<MockTask> requestorTask = addTask();
     MockMemoryOperator* requestorOp = addMemoryOp(requestorTask, false);
     requestorOp->allocate(testData.requestorCapacity);