avoid tiflash crash when query is killed (#3434) (#3451)

close #3401
pingcap · Feb 23, 2022 · 9b47191 · 9b47191
1 parent 2eae1bc
commit 9b47191
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 7 deletions.
diff --git a/dbms/src/DataStreams/ParallelAggregatingBlockInputStream.cpp b/dbms/src/DataStreams/ParallelAggregatingBlockInputStream.cpp
@@ -181,7 +181,9 @@ void ParallelAggregatingBlockInputStream::Handler::onFinish()
 void ParallelAggregatingBlockInputStream::Handler::onException(std::exception_ptr & exception, size_t thread_num)
 {
     parent.exceptions[thread_num] = exception;
-    parent.cancel(false);
+    /// can not cancel parent inputStream or the exception might be lost
+    if (!parent.executed)
+        parent.processor.cancel(false);
 }
 
 

diff --git a/dbms/src/DataStreams/UnionBlockInputStream.h b/dbms/src/DataStreams/UnionBlockInputStream.h
@@ -298,7 +298,8 @@ class UnionBlockInputStream final : public IProfilingBlockInputStream
             ///  and the exception is lost.
 
             parent.output_queue.push(exception);
-            parent.cancel(false); /// Does not throw exceptions.
+            /// can not cancel parent inputStream or the exception might be lost
+            parent.processor.cancel(false); /// Does not throw exceptions.
         }
 
         String getName() const

diff --git a/dbms/src/Flash/Mpp/MPPTask.cpp b/dbms/src/Flash/Mpp/MPPTask.cpp
@@ -58,7 +58,8 @@ MPPTask::~MPPTask()
 {
     /// MPPTask maybe destructed by different thread, set the query memory_tracker
     /// to current_memory_tracker in the destructor
-    current_memory_tracker = memory_tracker;
+    if (current_memory_tracker != memory_tracker)
+        current_memory_tracker = memory_tracker;
     closeAllTunnels("");
     LOG_DEBUG(log, "finish MPPTask: " << id.toString());
 }

diff --git a/dbms/src/Flash/Mpp/MPPTask.h b/dbms/src/Flash/Mpp/MPPTask.h
@@ -82,17 +82,18 @@ class MPPTask : public std::enable_shared_from_this<MPPTask>
 
     bool switchStatus(TaskStatus from, TaskStatus to);
 
-    Context context;
-
     RegionInfoMap local_regions;
     RegionInfoList remote_regions;
 
     std::unique_ptr<tipb::DAGRequest> dag_req;
-    std::unique_ptr<DAGContext> dag_context;
 
+    Context context;
     /// store io in MPPTask to keep the life cycle of memory_tracker for the current query
-    /// BlockIO contains some information stored in Context and DAGContext, so need deconstruct it before Context and DAGContext
+    /// BlockIO contains some information stored in Context, so need deconstruct it before Context
     BlockIO io;
+    /// The inputStreams should be released in the destructor of BlockIO, since DAGContext contains
+    /// some reference to inputStreams, so it need to be destructed before BlockIO
+    std::unique_ptr<DAGContext> dag_context;
     MemoryTracker * memory_tracker = nullptr;
 
     MPPTaskId id;