make SharedQueryInputStream compatible with UnionBlockInputStream (#2656

) (#2662)
pingcap · Sep 16, 2021 · a1debcb · a1debcb
1 parent c56ae2c
commit a1debcb
Showing 1 changed file with 16 additions and 6 deletions.
diff --git a/dbms/src/DataStreams/SharedQueryBlockInputStream.h b/dbms/src/DataStreams/SharedQueryBlockInputStream.h
@@ -69,8 +69,8 @@ class SharedQueryBlockInputStream : public IProfilingBlockInputStream
 
         if (thread.joinable())
             thread.join();
-        if (exception)
-            std::rethrow_exception(exception);
+        if (!exception_msg.empty())
+            throw Exception(exception_msg);
     }
 
 protected:
@@ -84,8 +84,10 @@ class SharedQueryBlockInputStream : public IProfilingBlockInputStream
         Block block;
         do
         {
-            if (exception)
-                std::rethrow_exception(exception);
+            if (!exception_msg.empty())
+            {
+                throw Exception(exception_msg);
+            }
             if (isCancelled() || read_suffixed)
                 return {};
         } while (!queue.tryPop(block, try_action_millisecionds));
@@ -116,9 +118,17 @@ class SharedQueryBlockInputStream : public IProfilingBlockInputStream
             }
             in->readSuffix();
         }
+        catch (Exception & e)
+        {
+            exception_msg = e.message();
+        }
+        catch (std::exception & e)
+        {
+            exception_msg = e.what();
+        }
         catch (...)
         {
-            exception = std::current_exception();
+            exception_msg = "other error";
         }
     }
 
@@ -133,7 +143,7 @@ class SharedQueryBlockInputStream : public IProfilingBlockInputStream
     std::thread thread;
     std::mutex  mutex;
 
-    std::exception_ptr exception;
+    std::string exception_msg;
 
     Logger * log;
     BlockInputStreamPtr in;