fix bug that ExchangeReceiver is not cancelled if exception happens i…

…n union/agg block input stream (#4285) (#4290) ref #4229
pingcap · Apr 22, 2022 · de32b2c · de32b2c
1 parent 839a659
commit de32b2c
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 12 deletions.
diff --git a/dbms/src/DataStreams/ParallelAggregatingBlockInputStream.cpp b/dbms/src/DataStreams/ParallelAggregatingBlockInputStream.cpp
@@ -183,7 +183,8 @@ void ParallelAggregatingBlockInputStream::Handler::onException(std::exception_pt
     parent.exceptions[thread_num] = exception;
     /// can not cancel parent inputStream or the exception might be lost
     if (!parent.executed)
-        parent.processor.cancel(false);
+        /// kill the processor so ExchangeReceiver will be closed
+        parent.processor.cancel(true);
 }
 
 

diff --git a/dbms/src/DataStreams/UnionBlockInputStream.h b/dbms/src/DataStreams/UnionBlockInputStream.h
@@ -26,11 +26,11 @@ struct OutputData<StreamUnionMode::Basic>
     Block block;
     std::exception_ptr exception;
 
-    OutputData() {}
+    OutputData() = default;
     explicit OutputData(Block & block_)
         : block(block_)
     {}
-    explicit OutputData(std::exception_ptr & exception_)
+    explicit OutputData(const std::exception_ptr & exception_)
         : exception(exception_)
     {}
 };
@@ -43,12 +43,12 @@ struct OutputData<StreamUnionMode::ExtraInfo>
     BlockExtraInfo extra_info;
     std::exception_ptr exception;
 
-    OutputData() {}
+    OutputData() = default;
     OutputData(Block & block_, BlockExtraInfo & extra_info_)
         : block(block_)
         , extra_info(extra_info_)
     {}
-    explicit OutputData(std::exception_ptr & exception_)
+    explicit OutputData(const std::exception_ptr & exception_)
         : exception(exception_)
     {}
 };
@@ -260,6 +260,23 @@ class UnionBlockInputStream final : public IProfilingBlockInputStream
       *  otherwise ParallelInputsProcessor can be blocked during insertion into the queue.
       */
     OutputQueue output_queue;
+    std::mutex mu;
+    bool meet_exception = false;
+
+    void handleException(const std::exception_ptr & exception)
+    {
+        std::unique_lock lock(mu);
+        if (meet_exception)
+            return;
+        meet_exception = true;
+        /// The order of the rows matters. If it is changed, then the situation is possible,
+        /// when before exception, an empty block (end of data) will be put into the queue,
+        /// and the exception is lost.
+        output_queue.emplace(exception);
+        /// can not cancel itself or the exception might be lost
+        /// kill the processor so ExchangeReceiver will be closed
+        processor.cancel(true);
+    }
 
     struct Handler
     {
@@ -290,13 +307,7 @@ class UnionBlockInputStream final : public IProfilingBlockInputStream
 
         void onException(std::exception_ptr & exception, size_t /*thread_num*/)
         {
-            /// The order of the rows matters. If it is changed, then the situation is possible,
-            /// when before exception, an empty block (end of data) will be put into the queue,
-            /// and the exception is lost.
-
-            parent.output_queue.emplace(exception);
-            /// can not cancel parent inputStream or the exception might be lost
-            parent.processor.cancel(false); /// Does not throw exceptions.
+            parent.handleException(exception);
         }
 
         String getName() const