Backport #92807 to 25.12: Fix crash when receiving from disconnected Connection

robot-clickhouse · robot-clickhouse · commit 817dca965ba1 · 2025-12-24T11:14:47.000Z
diff --git a/src/Processors/Sources/RemoteSource.cpp b/src/Processors/Sources/RemoteSource.cpp
@@ -93,8 +93,11 @@ ISource::Status RemoteSource::prepare()
     if (is_async_state)
         return Status::Async;
 
-    if (executor_finished)
+    if (query_executor->isFinished())
+    {
+        getPort().finish();
         return Status::Finished;
+    }
 
     Status status = ISource::prepare();
     /// To avoid resetting the connection (because of "unfinished" query) in the
@@ -126,7 +129,6 @@ void RemoteSource::work()
     if (need_drain)
     {
         query_executor->finish();
-        executor_finished = true;
         return;
     }
 
@@ -248,9 +250,7 @@ void RemoteSource::onCancel() noexcept
 void RemoteSource::onUpdatePorts()
 {
     if (getPort().isFinished())
-    {
         query_executor->finish();
-    }
 }
 
 
diff --git a/src/Processors/Sources/RemoteSource.h b/src/Processors/Sources/RemoteSource.h
@@ -48,7 +48,6 @@ class RemoteSource final : public ISource
 private:
     std::atomic_bool was_query_sent = false;
     bool need_drain = false;
-    bool executor_finished = false;
     bool add_aggregation_info = false;
     RemoteQueryExecutorPtr query_executor;
     RowsBeforeStepCounterPtr rows_before_limit;
diff --git a/src/QueryPipeline/RemoteQueryExecutor.cpp b/src/QueryPipeline/RemoteQueryExecutor.cpp
@@ -778,6 +778,9 @@ void RemoteQueryExecutor::finish()
 {
     LockAndBlocker guard(was_cancelled_mutex);
 
+    /// To make sure finish is only called once
+    SCOPE_EXIT({ finished = true; });
+
     /** If one of:
       * - nothing started to do;
       * - received all packets before EndOfStream;
diff --git a/src/QueryPipeline/RemoteQueryExecutor.h b/src/QueryPipeline/RemoteQueryExecutor.h
@@ -221,6 +221,8 @@ class RemoteQueryExecutor
     /// return true if parallel replica packet was processed
     bool processParallelReplicaPacketIfAny();
 
+    bool isFinished() const { return finished; }
+
 private:
     RemoteQueryExecutor(
         const String & query_,

Original file line number	Diff line number	Diff line change
`@@ -93,8 +93,11 @@ ISource::Status RemoteSource::prepare()`
`93`	`93`	`if (is_async_state)`
`94`	`94`	`return Status::Async;`
`95`	`95`
`96`		`- if (executor_finished)`
	`96`	`+ if (query_executor->isFinished())`
	`97`	`+ {`
	`98`	`+ getPort().finish();`
`97`	`99`	`return Status::Finished;`
	`100`	`+ }`
`98`	`101`
`99`	`102`	`Status status = ISource::prepare();`
`100`	`103`	`/// To avoid resetting the connection (because of "unfinished" query) in the`
`@@ -126,7 +129,6 @@ void RemoteSource::work()`
`126`	`129`	`if (need_drain)`
`127`	`130`	`{`
`128`	`131`	`query_executor->finish();`
`129`		`- executor_finished = true;`
`130`	`132`	`return;`
`131`	`133`	`}`
`132`	`134`
`@@ -248,9 +250,7 @@ void RemoteSource::onCancel() noexcept`
`248`	`250`	`void RemoteSource::onUpdatePorts()`
`249`	`251`	`{`
`250`	`252`	`if (getPort().isFinished())`
`251`		`- {`
`252`	`253`	`query_executor->finish();`
`253`		`- }`
`254`	`254`	`}`
`255`	`255`
`256`	`256`
Original file line number	Diff line number	Diff line change
`@@ -778,6 +778,9 @@ void RemoteQueryExecutor::finish()`
`778`	`778`	`{`
`779`	`779`	`LockAndBlocker guard(was_cancelled_mutex);`
`780`	`780`
	`781`	`+ /// To make sure finish is only called once`
	`782`	`+ SCOPE_EXIT({ finished = true; });`
	`783`	`+`
`781`	`784`	`/** If one of:`
`782`	`785`	`* - nothing started to do;`
`783`	`786`	`* - received all packets before EndOfStream;`