Fix MockSharedArbitrationTest.ensureMemoryPoolMaxCapacity

tanjialiang · Nov 7, 2024 · dca6292 · dca6292
1 parent 8c69b8b
commit dca6292
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 1 deletion.
diff --git 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/SharedArbitrator.h b/velox/common/memory/SharedArbitrator.h
@@ -258,6 +258,10 @@ class SharedArbitrator : public memory::MemoryArbitrator {
 
   std::string toString() const final;
 
+  bool testingGlobalArbitrationRunning() const {
+    return globalArbitrationRunning_;
+  }
+
   /// Operator level runtime stats reported for an arbitration operation
   /// execution.
   static inline const std::string kMemoryArbitrationWallNanos{

diff --git a/velox/common/memory/tests/MockSharedArbitratorTest.cpp b/velox/common/memory/tests/MockSharedArbitratorTest.cpp
@@ -2341,6 +2341,9 @@ TEST_F(MockSharedArbitrationTest, ensureMemoryPoolMaxCapacity) {
     } else if (testData.hasOtherTask) {
       ASSERT_EQ(otherOp->reclaimer()->stats().numReclaims, 0);
     }
+    while (arbitrator_->testingGlobalArbitrationRunning()) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(200));
+    }
     if (testData.expectedSuccess &&
         (((testData.allocatedBytes + testData.requestBytes) >
           testData.poolMaxCapacity) ||