From 76362793e21f4d571b05a9d2c83deb85fab6c602 Mon Sep 17 00:00:00 2001
From: SeaRise <hhssearise@foxmail.com>
Date: Wed, 22 Mar 2023 10:37:12 +0800
Subject: [PATCH] fix

---
 dbms/src/Flash/Coprocessor/DAGRequest.cpp | 10 +++++-----
 dbms/src/Flash/Mpp/MPPTask.cpp            |  2 +-
 dbms/src/Flash/Mpp/MPPTask.h              |  2 ++
 3 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/dbms/src/Flash/Coprocessor/DAGRequest.cpp b/dbms/src/Flash/Coprocessor/DAGRequest.cpp
index 8f25e7ad45e..a1a6440dd9e 100644
--- a/dbms/src/Flash/Coprocessor/DAGRequest.cpp
+++ b/dbms/src/Flash/Coprocessor/DAGRequest.cpp
@@ -87,7 +87,7 @@ DAGRequest::DAGRequest(tipb::DAGRequest * dag_request_)
 
 void DAGRequest::checkOrSetExecutorId()
 {
-    if (isTreeBased())
+    if (is_tree_based)
     {
         // check duplicate executor_id for tree based request.
         std::unordered_set<String> ids;
@@ -125,13 +125,13 @@ const tipb::Executor & DAGRequest::rootExecutor() const
 
 void DAGRequest::traverse(std::function<bool(const tipb::Executor &)> && func) const
 {
-    if likely (dag_request)
-        traverseExecutors(dag_request, std::move(func));
+    check(dag_request, "dagrequest cannot be null");
+    traverseExecutors(dag_request, std::move(func));
 }
 
 void DAGRequest::traverseReverse(std::function<void(const tipb::Executor &)> && func) const
 {
-    if likely (dag_request)
-        traverseExecutorsReverse(dag_request, std::move(func));
+    check(dag_request, "dagrequest cannot be null");
+    traverseExecutorsReverse(dag_request, std::move(func));
 }
 } // namespace DB
diff --git a/dbms/src/Flash/Mpp/MPPTask.cpp b/dbms/src/Flash/Mpp/MPPTask.cpp
index babe5ac19ba..5752159f618 100644
--- a/dbms/src/Flash/Mpp/MPPTask.cpp
+++ b/dbms/src/Flash/Mpp/MPPTask.cpp
@@ -265,7 +265,7 @@ void MPPTask::unregisterTask()
 
 void MPPTask::prepare(const mpp::DispatchTaskRequest & task_request)
 {
-    auto dag_req = getDAGRequestFromStringWithRetry(task_request.encoded_plan());
+    dag_req = getDAGRequestFromStringWithRetry(task_request.encoded_plan());
     TMTContext & tmt_context = context->getTMTContext();
     /// MPP task will only use key ranges in mpp::DispatchTaskRequest::regions/mpp::DispatchTaskRequest::table_regions.
     /// The ones defined in tipb::TableScan will never be used and can be removed later.
diff --git a/dbms/src/Flash/Mpp/MPPTask.h b/dbms/src/Flash/Mpp/MPPTask.h
index 044b5ab1608..302b64fff87 100644
--- a/dbms/src/Flash/Mpp/MPPTask.h
+++ b/dbms/src/Flash/Mpp/MPPTask.h
@@ -109,6 +109,8 @@ class MPPTask : public std::enable_shared_from_this<MPPTask>
 
     void initExchangeReceivers();
 
+    // To make sure dag_req is not destroyed before the mpp task ends.
+    tipb::DAGRequest dag_req;
     mpp::TaskMeta meta;
     MPPTaskId id;