From c89d8d6346f08c0bec13356641a210fbd67fec66 Mon Sep 17 00:00:00 2001
From: Jialiang Tan <jtan6@meta.com>
Date: Thu, 31 Oct 2024 10:20:53 -0700
Subject: [PATCH] Fix flaky test
 MockSharedArbitrationTest.arbitrateBySelfMemoryReclaim (#11397)

Summary:
Pull Request resolved: https://github.com/facebookincubator/velox/pull/11397

The test originally intended to test local arbitration. But after refactoring, this test is by default relying on global arbitration for reclaiming. Because global arbitration runs in a different thread, there is a race condition between the actual reclaim and the reclaimed bytes check in the test. Making the test disable global arbitration solves the flakiness.

Reviewed By: xiaoxmeng

Differential Revision: D65259718

fbshipit-source-id: 86a2937c620d1ac58bd9b76c9948f9d9b0aaa0d5
---
 .../memory/tests/MockSharedArbitratorTest.cpp | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/velox/common/memory/tests/MockSharedArbitratorTest.cpp b/velox/common/memory/tests/MockSharedArbitratorTest.cpp
index 67eb1ffa4d91..c30d6fab22e6 100644
--- a/velox/common/memory/tests/MockSharedArbitratorTest.cpp
+++ b/velox/common/memory/tests/MockSharedArbitratorTest.cpp
@@ -3283,14 +3283,26 @@ DEBUG_ONLY_TEST_F(MockSharedArbitrationTest, reclaimWithNoCandidate) {
 }
 
 TEST_F(MockSharedArbitrationTest, arbitrateBySelfMemoryReclaim) {
-  const std::vector<bool> isLeafReclaimables = {true, false};
-  for (const auto isLeafReclaimable : isLeafReclaimables) {
+  for (const auto isLeafReclaimable : {true, false}) {
     SCOPED_TRACE(fmt::format("isLeafReclaimable {}", isLeafReclaimable));
     const uint64_t memCapacity = 128 * MB;
     const uint64_t reservedCapacity = 8 * MB;
     const uint64_t poolReservedCapacity = 4 * MB;
     setupMemory(
-        memCapacity, reservedCapacity, reservedCapacity, poolReservedCapacity);
+        memCapacity,
+        reservedCapacity,
+        reservedCapacity,
+        poolReservedCapacity,
+        0,
+        0,
+        0,
+        0,
+        0,
+        0,
+        0,
+        kMemoryReclaimThreadsHwMultiplier,
+        nullptr,
+        false);
     std::shared_ptr<MockTask> task = addTask(kMemoryCapacity);
     auto* memOp = addMemoryOp(task, isLeafReclaimable);
     const int allocateSize = 8 * MB;
@@ -3305,7 +3317,6 @@ TEST_F(MockSharedArbitrationTest, arbitrateBySelfMemoryReclaim) {
           memOp->allocate(memCapacity), "Exceeded memory pool cap");
       ASSERT_EQ(oldNumRequests + 1, arbitrator_->stats().numRequests);
       ASSERT_EQ(arbitrator_->stats().numFailures, 1);
-      continue;
     } else {
       memOp->allocate(memCapacity / 2);
       ASSERT_EQ(oldNumRequests + 1, arbitrator_->stats().numRequests);