Wait for all the tasks to be deleted before reset memory system (face…

…bookincubator#9458) Summary: Pull Request resolved: facebookincubator#9458 Reviewed By: tanjialiang Differential Revision: D56041868 Pulled By: xiaoxmeng
xiaoxmeng · Apr 11, 2024 · 6d43fa2 · 6d43fa2
1 parent 0643556
commit 6d43fa2
Show file tree

Hide file tree

Showing 8 changed files with 1 addition and 20 deletions.
diff --git a/velox/common/memory/tests/MemoryCapExceededTest.cpp b/velox/common/memory/tests/MemoryCapExceededTest.cpp
@@ -37,7 +37,6 @@ class MemoryCapExceededTest : public OperatorTestBase,
   }
 
   void TearDown() override {
-    waitForAllTasksToBeDeleted();
     OperatorTestBase::TearDown();
     FLAGS_velox_suppress_memory_capacity_exceeding_error_message = false;
   }

diff --git a/velox/exec/tests/HashJoinTest.cpp b/velox/exec/tests/HashJoinTest.cpp
@@ -792,11 +792,6 @@ class HashJoinTest : public HiveConnectorTestBase {
         .allowLazyVector = false};
   }
 
-  void TearDown() override {
-    waitForAllTasksToBeDeleted();
-    HiveConnectorTestBase::TearDown();
-  }
-
   // Make splits with each plan node having a number of source files.
   SplitInput makeSpiltInput(
       const std::vector<core::PlanNodeId>& nodeIds,

diff --git a/velox/exec/tests/MultiFragmentTest.cpp b/velox/exec/tests/MultiFragmentTest.cpp
@@ -47,8 +47,6 @@ class MultiFragmentTest : public HiveConnectorTestBase {
   }
 
   void TearDown() override {
-    waitForAllTasksToBeDeleted();
-
     // There might be lingering exchange source on executor even after all tasks
     // are deleted. This can cause memory leak because exchange source holds
     // reference to memory pool. We need to make sure they are properly cleaned.

diff --git a/velox/exec/tests/OperatorUtilsTest.cpp b/velox/exec/tests/OperatorUtilsTest.cpp
@@ -31,7 +31,6 @@ class OperatorUtilsTest : public OperatorTestBase {
     driverCtx_.reset();
     driver_.reset();
     task_.reset();
-    waitForAllTasksToBeDeleted();
     OperatorTestBase::TearDown();
   }
 

diff --git a/velox/exec/tests/TableWriteTest.cpp b/velox/exec/tests/TableWriteTest.cpp
@@ -248,11 +248,6 @@ class TableWriteTest : public HiveConnectorTestBase {
     HiveConnectorTestBase::SetUp();
   }
 
-  void TearDown() override {
-    waitForAllTasksToBeDeleted();
-    HiveConnectorTestBase::TearDown();
-  }
-
   std::shared_ptr<Task> assertQueryWithWriterConfigs(
       const core::PlanNodePtr& plan,
       std::vector<std::shared_ptr<TempFilePath>> filePaths,

diff --git a/velox/exec/tests/TaskTest.cpp b/velox/exec/tests/TaskTest.cpp
@@ -460,11 +460,6 @@ class TestBadMemoryTranslator : public exec::Operator::PlanNodeTranslator {
 
 class TaskTest : public HiveConnectorTestBase {
  protected:
-  void TearDown() override {
-    waitForAllTasksToBeDeleted();
-    HiveConnectorTestBase::TearDown();
-  }
-
   static std::pair<std::shared_ptr<exec::Task>, std::vector<RowVectorPtr>>
   executeSingleThreaded(
       core::PlanFragment plan,

diff --git a/velox/exec/tests/ValuesTest.cpp b/velox/exec/tests/ValuesTest.cpp
@@ -26,7 +26,6 @@ namespace facebook::velox::exec::test {
 class ValuesTest : public OperatorTestBase {
  protected:
   void TearDown() override {
-    waitForAllTasksToBeDeleted();
     input_.reset();
     input2_.reset();
     OperatorTestBase::TearDown();

diff --git a/velox/exec/tests/utils/OperatorTestBase.cpp b/velox/exec/tests/utils/OperatorTestBase.cpp
@@ -103,6 +103,7 @@ void OperatorTestBase::SetUp() {
 }
 
 void OperatorTestBase::TearDown() {
+  waitForAllTasksToBeDeleted();
   pool_.reset();
   rootPool_.reset();
   resetMemory();