From acd57170b6d98206def1ef02b74c467e3ba18061 Mon Sep 17 00:00:00 2001
From: Satadru Pan <satadru@meta.com>
Date: Wed, 9 Oct 2024 00:14:23 -0700
Subject: [PATCH] De-flake DriverTest.pause (#11202)

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

This diff fixes a rare race between test shutdown and Task shutdowm. Even though Test shutdown waits for Task constructor to finish, Test shutdown was not waiting for all member variables of Task to be cleaned up: for example task.queryCtx_ holds memory pool and if task.queryCtx_ is not freed before test shutdown, we get test failure. This diff fixes queryCtx_ destruction by reseting the pointer in ~Task.

Reviewed By: xiaoxmeng

Differential Revision: D64026457

fbshipit-source-id: f0c350530b7f1c8cbfc051df3e12ad4d0b95e364
---
 velox/exec/Task.cpp |  5 +++--
 velox/exec/Task.h   | 10 +++++++---
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/velox/exec/Task.cpp b/velox/exec/Task.cpp
index c7efac498dc9..c8acf566e086 100644
--- a/velox/exec/Task.cpp
+++ b/velox/exec/Task.cpp
@@ -288,11 +288,11 @@ Task::Task(
     std::function<void(std::exception_ptr)> onError)
     : uuid_{makeUuid()},
       taskId_(taskId),
-      planFragment_(std::move(planFragment)),
       destination_(destination),
+      mode_(mode),
       queryCtx_(std::move(queryCtx)),
+      planFragment_(std::move(planFragment)),
       traceConfig_(maybeMakeTraceConfig()),
-      mode_(mode),
       consumerSupplier_(std::move(consumerSupplier)),
       onError_(std::move(onError)),
       splitsStates_(buildSplitStates(planFragment_.planNode)),
@@ -347,6 +347,7 @@ Task::~Task() {
   CLEAR(childPools_.clear());
   CLEAR(pool_.reset());
   CLEAR(planFragment_ = core::PlanFragment());
+  CLEAR(queryCtx_.reset());
   clearStage = "exiting ~Task()";
 
   // Ful-fill the task deletion promises at the end.
diff --git a/velox/exec/Task.h b/velox/exec/Task.h
index 11bb852f56b6..b3e159624904 100644
--- a/velox/exec/Task.h
+++ b/velox/exec/Task.h
@@ -992,15 +992,19 @@ class Task : public std::enable_shared_from_this<Task> {
   // Application specific task ID specified at construction time. May not be
   // unique or universally unique.
   const std::string taskId_;
-  core::PlanFragment planFragment_;
+
   const int destination_;
-  const std::shared_ptr<core::QueryCtx> queryCtx_;
-  const std::optional<trace::QueryTraceConfig> traceConfig_;
 
   // The execution mode of the task. It is enforced that a task can only be
   // executed in a single mode throughout its lifetime
   const ExecutionMode mode_;
 
+  std::shared_ptr<core::QueryCtx> queryCtx_;
+
+  core::PlanFragment planFragment_;
+
+  const std::optional<trace::QueryTraceConfig> traceConfig_;
+
   // Hook in the system wide task list.
   TaskListEntry taskListEntry_;