make mpp_fail test stable (pingcap#4300)

close pingcap#4299
ti-chi-bot · Jun 15, 2022 · 90dcd38 · 90dcd38
1 parent a518c82
commit 90dcd38
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/dbms/src/DataStreams/ParallelAggregatingBlockInputStream.cpp b/dbms/src/DataStreams/ParallelAggregatingBlockInputStream.cpp
@@ -181,6 +181,9 @@ void ParallelAggregatingBlockInputStream::Handler::onFinish()
 void ParallelAggregatingBlockInputStream::Handler::onException(std::exception_ptr & exception, size_t thread_num)
 {
     parent.exceptions[thread_num] = exception;
+    Int32 old_value = -1;
+    parent.first_exception_index.compare_exchange_strong(old_value, static_cast<Int32>(thread_num), std::memory_order_seq_cst, std::memory_order_relaxed);
+
     /// can not cancel parent inputStream or the exception might be lost
     if (!parent.executed)
         parent.processor.cancel(false);
@@ -205,7 +208,8 @@ void ParallelAggregatingBlockInputStream::execute()
     processor.process();
     processor.wait();
 
-    rethrowFirstException(exceptions);
+    if (first_exception_index != -1)
+        std::rethrow_exception(exceptions[first_exception_index]);
 
     if (isCancelledOrThrowIfKilled())
         return;

diff --git a/dbms/src/DataStreams/ParallelAggregatingBlockInputStream.h b/dbms/src/DataStreams/ParallelAggregatingBlockInputStream.h
@@ -83,6 +83,7 @@ class ParallelAggregatingBlockInputStream : public IProfilingBlockInputStream
 
     ManyAggregatedDataVariants many_data;
     Exceptions exceptions;
+    std::atomic<Int32> first_exception_index{-1};
 
     struct ThreadData
     {