From 74558a4e27d14ff7dc6ae21b5e7dc10cf3f48d06 Mon Sep 17 00:00:00 2001
From: Alexey Milovidov <milovidov@yandex-team.ru>
Date: Sat, 24 Oct 2020 06:41:47 +0300
Subject: [PATCH] Better diagnostics when client has dropped connection

---
 src/Server/TCPHandler.cpp | 26 +++++++++++++++++++++++---
 src/Server/TCPHandler.h   |  1 +
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/src/Server/TCPHandler.cpp b/src/Server/TCPHandler.cpp
index 406e29ba4ab..a37f88f9306 100644
--- a/src/Server/TCPHandler.cpp
+++ b/src/Server/TCPHandler.cpp
@@ -277,6 +277,9 @@ void TCPHandler::runImpl()
             /// Do it before sending end of stream, to have a chance to show log message in client.
             query_scope->logPeakMemoryUsage();
 
+            if (state.is_connection_closed)
+                break;
+
             sendLogs();
             sendEndOfStream();
 
@@ -444,7 +447,11 @@ bool TCPHandler::readDataNext(const size_t & poll_interval, const int & receive_
 
     /// If client disconnected.
     if (in->eof())
+    {
+        LOG_INFO(log, "Client has dropped the connection, cancel the query.");
+        state.is_connection_closed = true;
         return false;
+    }
 
     /// We accept and process data. And if they are over, then we leave.
     if (!receivePacket())
@@ -477,9 +484,8 @@ void TCPHandler::readData(const Settings & connection_settings)
     std::tie(poll_interval, receive_timeout) = getReadTimeouts(connection_settings);
     sendLogs();
 
-    while (true)
-        if (!readDataNext(poll_interval, receive_timeout))
-            return;
+    while (readDataNext(poll_interval, receive_timeout))
+        ;
 }
 
 
@@ -567,6 +573,9 @@ void TCPHandler::processOrdinaryQuery()
             sendProgress();
         }
 
+        if (state.is_connection_closed)
+            return;
+
         sendData({});
     }
 
@@ -632,6 +641,9 @@ void TCPHandler::processOrdinaryQueryWithProcessors()
             sendLogs();
         }
 
+        if (state.is_connection_closed)
+            return;
+
         sendData({});
     }
 
@@ -1179,6 +1191,14 @@ bool TCPHandler::isQueryCancelled()
     /// During request execution the only packet that can come from the client is stopping the query.
     if (static_cast<ReadBufferFromPocoSocket &>(*in).poll(0))
     {
+        if (in->eof())
+        {
+            LOG_INFO(log, "Client has dropped the connection, cancel the query.");
+            state.is_cancelled = true;
+            state.is_connection_closed = true;
+            return true;
+        }
+
         UInt64 packet_type = 0;
         readVarUInt(packet_type, *in);
 
diff --git a/src/Server/TCPHandler.h b/src/Server/TCPHandler.h
index 12149d9a66f..2f2bf35e59e 100644
--- a/src/Server/TCPHandler.h
+++ b/src/Server/TCPHandler.h
@@ -57,6 +57,7 @@ struct QueryState
 
     /// Is request cancelled
     bool is_cancelled = false;
+    bool is_connection_closed = false;
     /// empty or not
     bool is_empty = true;
     /// Data was sent.