Backport #95466 to 26.1: Avoid possible crash for distributed queries in case of cancellation

robot-clickhouse · robot-clickhouse · commit 2cb868ba843b · 2026-01-30T10:22:30.000Z
diff --git a/src/Client/Connection.cpp b/src/Client/Connection.cpp
@@ -84,6 +84,7 @@ namespace ErrorCodes
     extern const int SUPPORT_IS_DISABLED;
     extern const int BAD_ARGUMENTS;
     extern const int EMPTY_DATA_PASSED;
+    extern const int LOGICAL_ERROR;
 }
 
 Connection::~Connection()
@@ -1326,6 +1327,11 @@ Packet Connection::receivePacket()
 {
     try
     {
+        /// We are trying to send something to already disconnected connection,
+        /// this means that we continue using Connection after exception.
+        if (!in)
+            throw Exception(ErrorCodes::LOGICAL_ERROR, "Connection to {} is terminated", getDescription());
+
         Packet res;
 
         /// Have we already read packet type?
diff --git a/src/Processors/Sources/RemoteSource.cpp b/src/Processors/Sources/RemoteSource.cpp
@@ -249,6 +249,8 @@ void RemoteSource::onCancel() noexcept
 
 void RemoteSource::onUpdatePorts()
 {
+    if (isCancelled())
+        return;
     if (getPort().isFinished())
         query_executor->finish();
 }
diff --git a/src/QueryPipeline/RemoteQueryExecutor.cpp b/src/QueryPipeline/RemoteQueryExecutor.cpp
@@ -785,7 +785,7 @@ void RemoteQueryExecutor::finish()
       * - received an unknown packet from one replica;
       * then you do not need to read anything.
       */
-    if (!isQueryPending() || hasThrownException())
+    if (!isQueryPending() || hasThrownException() || was_cancelled)
         return;
 
     /// To make sure finish is only called once

Original file line number	Diff line number	Diff line change
`@@ -84,6 +84,7 @@ namespace ErrorCodes`
`84`	`84`	`extern const int SUPPORT_IS_DISABLED;`
`85`	`85`	`extern const int BAD_ARGUMENTS;`
`86`	`86`	`extern const int EMPTY_DATA_PASSED;`
	`87`	`+ extern const int LOGICAL_ERROR;`
`87`	`88`	`}`
`88`	`89`
`89`	`90`	`Connection::~Connection()`
`@@ -1326,6 +1327,11 @@ Packet Connection::receivePacket()`
`1326`	`1327`	`{`
`1327`	`1328`	`try`
`1328`	`1329`	`{`
	`1330`	`+ /// We are trying to send something to already disconnected connection,`
	`1331`	`+ /// this means that we continue using Connection after exception.`
	`1332`	`+ if (!in)`
	`1333`	`+ throw Exception(ErrorCodes::LOGICAL_ERROR, "Connection to {} is terminated", getDescription());`
	`1334`	`+`
`1329`	`1335`	`Packet res;`
`1330`	`1336`
`1331`	`1337`	`/// Have we already read packet type?`
Original file line number	Diff line number	Diff line change
`@@ -249,6 +249,8 @@ void RemoteSource::onCancel() noexcept`
`249`	`249`
`250`	`250`	`void RemoteSource::onUpdatePorts()`
`251`	`251`	`{`
	`252`	`+ if (isCancelled())`
	`253`	`+ return;`
`252`	`254`	`if (getPort().isFinished())`
`253`	`255`	`query_executor->finish();`
`254`	`256`	`}`