From ba2c06a1801029752e5e089673d74764c3f592da Mon Sep 17 00:00:00 2001
From: xiaoxmeng <xiaoxmeng@fb.com>
Date: Thu, 11 Apr 2024 15:58:54 -0700
Subject: [PATCH] Wait for all the tasks to be deleted before reset memory
 system

---
 velox/exec/tests/HashJoinTest.cpp           | 5 -----
 velox/exec/tests/MultiFragmentTest.cpp      | 2 --
 velox/exec/tests/OperatorUtilsTest.cpp      | 1 -
 velox/exec/tests/TableWriteTest.cpp         | 5 -----
 velox/exec/tests/TaskTest.cpp               | 5 -----
 velox/exec/tests/ValuesTest.cpp             | 1 -
 velox/exec/tests/utils/OperatorTestBase.cpp | 1 +
 7 files changed, 1 insertion(+), 19 deletions(-)

diff --git a/velox/exec/tests/HashJoinTest.cpp b/velox/exec/tests/HashJoinTest.cpp
index b859f86351b6..84f9c6370cc1 100644
--- 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
index 3e8f4c03517b..110729da8850 100644
--- 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
index 24ce90dd7480..c0abaa5f2226 100644
--- 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
index 23fc416e08ab..5e688f793fda 100644
--- 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
index 3bf0a1d487a6..7963c9ad2a16 100644
--- 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
index f9d79872af45..368fb17bfb52 100644
--- 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
index 7ad53cc8889a..9c59dd09c853 100644
--- 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();