From cc48ac9a61586fe2767f9ace43e59fe1d7a7b060 Mon Sep 17 00:00:00 2001
From: Jialiang Tan <jacob.jialiang.tan@gmail.com>
Date: Mon, 11 Nov 2024 10:45:03 -0800
Subject: [PATCH] Fix MockSharedArbitrationTest.ensureMemoryPoolMaxCapacity
 (#11470)

Summary:
global arbitration now runs in a separate thread. The stats updating is async to the main testing thread. So the check might be flaky in this test. Making sure global arbitration finishes before checking stats eliminates the flakiness.

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

Reviewed By: xiaoxmeng

Differential Revision: D65613023

Pulled By: tanjialiang

fbshipit-source-id: 4a41fffc524b014424931f605a68f13f3375389a
---
 velox/common/memory/SharedArbitrator.cpp               | 1 -
 velox/common/memory/tests/MockSharedArbitratorTest.cpp | 3 +++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/velox/common/memory/SharedArbitrator.cpp b/velox/common/memory/SharedArbitrator.cpp
index 2dc13af45901..6ca1601826f0 100644
--- a/velox/common/memory/SharedArbitrator.cpp
+++ b/velox/common/memory/SharedArbitrator.cpp
@@ -1384,7 +1384,6 @@ SharedArbitrator::GlobalArbitrationSection::GlobalArbitrationSection(
 SharedArbitrator::GlobalArbitrationSection::~GlobalArbitrationSection() {
   VELOX_CHECK(arbitrator_->globalArbitrationRunning_);
   arbitrator_->globalArbitrationRunning_ = false;
-  ;
 }
 
 std::string SharedArbitrator::kind() const {
diff --git a/velox/common/memory/tests/MockSharedArbitratorTest.cpp b/velox/common/memory/tests/MockSharedArbitratorTest.cpp
index b506c5c39876..1a36784bc887 100644
--- a/velox/common/memory/tests/MockSharedArbitratorTest.cpp
+++ b/velox/common/memory/tests/MockSharedArbitratorTest.cpp
@@ -2287,6 +2287,7 @@ TEST_F(MockSharedArbitrationTest, ensureMemoryPoolMaxCapacity) {
        memCapacity - memCapacity / 2,
        false,
        false}};
+
   for (const auto& testData : testSettings) {
     SCOPED_TRACE(testData.debugString());
     setupMemory(
@@ -2322,6 +2323,8 @@ TEST_F(MockSharedArbitrationTest, ensureMemoryPoolMaxCapacity) {
     } else if (testData.hasOtherTask) {
       ASSERT_EQ(otherOp->reclaimer()->stats().numReclaims, 0);
     }
+    test::SharedArbitratorTestHelper arbitratorHelper(arbitrator_);
+    arbitratorHelper.waitForGlobalArbitrationToFinish();
     if (testData.expectedSuccess &&
         (((testData.allocatedBytes + testData.requestBytes) >
           testData.poolMaxCapacity) ||