fix bug that SharedQueryBlockInputStream may loss block randomly (#2759…

…) (#2761)
pingcap · Sep 1, 2021 · 550e6bc · 550e6bc
1 parent 884f5fb
commit 550e6bc
Showing 1 changed file with 8 additions and 15 deletions.
diff --git a/dbms/src/DataStreams/SharedQueryBlockInputStream.h b/dbms/src/DataStreams/SharedQueryBlockInputStream.h
@@ -1,13 +1,12 @@
 #pragma once
 
-#include <thread>
-
 #include <Common/ConcurrentBoundedQueue.h>
 #include <common/logger_useful.h>
 #include <Common/typeid_cast.h>
-
 #include <DataStreams/IProfilingBlockInputStream.h>
 
+#include <thread>
+
 namespace DB
 {
 
@@ -37,15 +36,9 @@ class SharedQueryBlockInputStream : public IProfilingBlockInputStream
         }
     }
 
-    String getName() const override
-    {
-        return "SharedQuery";
-    }
+    String getName() const override { return "SharedQuery"; }
 
-    Block getHeader() const override
-    {
-        return children.back()->getHeader();
-    }
+    Block getHeader() const override { return children.back()->getHeader(); }
 
     void readPrefix() override
     {
@@ -102,7 +95,7 @@ class SharedQueryBlockInputStream : public IProfilingBlockInputStream
             in->readPrefix();
             while (!isCancelled())
             {
-                Block block;
+                Block block = in->read();
                 do
                 {
                     if (isCancelled() || read_suffixed)
@@ -111,7 +104,7 @@ class SharedQueryBlockInputStream : public IProfilingBlockInputStream
                         queue.tryEmplace(0);
                         break;
                     }
-                } while (!queue.tryPush(block = in->read(), try_action_millisecionds));
+                } while (!queue.tryPush(block, try_action_millisecionds));
 
                 if (!block)
                     break;
@@ -141,11 +134,11 @@ class SharedQueryBlockInputStream : public IProfilingBlockInputStream
     bool read_suffixed = false;
 
     std::thread thread;
-    std::mutex  mutex;
+    std::mutex mutex;
 
     std::string exception_msg;
 
     Logger * log;
     BlockInputStreamPtr in;
 };
-}
+} // namespace DB