Return try-catch for unexpected EOF

Author: ianton-ru

src/QueryPipeline/RemoteQueryExecutor.cpp

@@ -730,6 +730,15 @@ RemoteQueryExecutor::ReadResult RemoteQueryExecutor::processPacket(Packet packet
             break;
 
         case Protocol::Server::ConnectionLost:
+            if (extension && extension->task_iterator && extension->task_iterator->supportRerunTask() && extension->replica_info)
+            {
+                if (!replica_has_processed_data.contains(extension->replica_info->number_of_current_replica))
+                {
+                    finished = true;
+                    extension->task_iterator->rescheduleTasksFromReplica(extension->replica_info->number_of_current_replica);
+                    return ReadResult(Block{});
+                }
+            }
             packet.exception->rethrow();
             break;
 
@@ -1000,6 +1009,11 @@ void RemoteQueryExecutor::setProfileInfoCallback(ProfileInfoCallback callback)
     profile_info_callback = std::move(callback);
 }
 
+bool RemoteQueryExecutor::skipUnavailableShards() const
+{
+    return context->getSettingsRef()[Setting::skip_unavailable_shards];
+}
+
 bool RemoteQueryExecutor::needToSkipUnavailableShard() const
 {
     return context->getSettingsRef()[Setting::skip_unavailable_shards] && (0 == connections->size());

src/QueryPipeline/RemoteQueryExecutor.h

@@ -233,6 +233,8 @@ class RemoteQueryExecutor
 
     IConnections & getConnections() { return *connections; }
 
+    bool skipUnavailableShards() const;
+
     bool needToSkipUnavailableShard() const;
 
     bool isReplicaUnavailable() const { return extension && extension->parallel_reading_coordinator && connections->size() == 0; }

src/QueryPipeline/RemoteQueryExecutorReadContext.cpp

@@ -57,22 +57,42 @@ void RemoteQueryExecutorReadContext::Task::run(AsyncCallback async_callback, Sus
     if (read_context.executor.needToSkipUnavailableShard())
         return;
 
-    while (true)
+    try
     {
-        read_context.has_read_packet_part = PacketPart::None;
+        while (true)
+        {
+            read_context.has_read_packet_part = PacketPart::None;
+
+            if (read_context.read_packet_type_separately)
+            {
+                read_context.packet.type = read_context.executor.getConnections().receivePacketTypeUnlocked(async_callback);
+                read_context.has_read_packet_part = PacketPart::Type;
+                suspend_callback();
+            }
+            read_context.packet = read_context.executor.getConnections().receivePacketUnlocked(async_callback);
+            read_context.has_read_packet_part = PacketPart::Body;
+            if (read_context.packet.type == Protocol::Server::Data)
+                read_context.has_data_packets = true;
 
-        if (read_context.read_packet_type_separately)
+            suspend_callback();
+        }
+    }
+    catch (const Exception & e)
+    {
+        /// If cluster node unxepectedly shutted down (kill/segfault/power off/etc.) socket just closes.
+        /// If initiator did not process any data packets before, this fact can be ignored.
+        /// Unprocessed tasks will be executed on other nodes.
+        if (e.code() == ErrorCodes::ATTEMPT_TO_READ_AFTER_EOF
+            && !read_context.has_data_packets.load()
+            && read_context.executor.skipUnavailableShards())
         {
-            read_context.packet.type = read_context.executor.getConnections().receivePacketTypeUnlocked(async_callback);
-            read_context.has_read_packet_part = PacketPart::Type;
+            read_context.packet.type = Protocol::Server::ConnectionLost;
+            read_context.packet.exception = std::make_unique<Exception>(getCurrentExceptionMessageAndPattern(true), getCurrentExceptionCode());
+            read_context.has_read_packet_part = PacketPart::Body;
             suspend_callback();
         }
-        read_context.packet = read_context.executor.getConnections().receivePacketUnlocked(async_callback);
-        read_context.has_read_packet_part = PacketPart::Body;
-        if (read_context.packet.type == Protocol::Server::Data)
-            read_context.has_data_packets = true;
-
-        suspend_callback();
+        else
+            throw;
     }
 }