Fix ReadBuffer canceled assertion in TCPHandler idle loop

nerve-bot · claude · nerve-bot · commit 95e1ecdc73b0 · 2026-03-24T04:29:58.000Z
TCPHandler::runImpl() idle loop called in-&gt;eof() on a ReadBuffer that
had been canceled by a previous query's pipeline callback. In debug and
sanitizer builds, ReadBuffer::next() has chassert(!isCanceled()) which
aborts the server process. This was observed as recurring "Logical error:
ReadBuffer is canceled" crashes in stress tests (amd_tsan, amd_msan,
amd_debug) — STID 2508-2913, 2508-25af.

Add ReadBuffer::eofOrCanceled() which checks isCanceled() before eof(),
preventing the assertion. TCPHandler idle loop now uses eofOrCanceled().

Also add explicit rethrowExceptionIfHas() after cancel() in
PullingAsyncPipelineExecutor::pull() as a defensive measure for the
is_execution_finished path.

Co-Authored-By: Claude Opus 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/src/IO/ReadBuffer.h b/src/IO/ReadBuffer.h
@@ -81,6 +81,17 @@ class ReadBuffer : public BufferBase
         return !hasPendingData() && !next();
     }
 
+    /// Safe EOF check that handles canceled buffers. Returns true if the buffer
+    /// is canceled OR at EOF. Use this instead of eof() when the buffer might
+    /// have been canceled by a failed read (e.g., a broken TCP connection
+    /// detected inside a pipeline callback). Calling eof() directly on a
+    /// canceled buffer triggers chassert(!isCanceled()) in debug/sanitizer
+    /// builds — see ReadBuffer::next().
+    bool ALWAYS_INLINE eofOrCanceled()
+    {
+        return isCanceled() || eof();
+    }
+
     void ignore()
     {
         if (!eof())
diff --git a/src/IO/tests/gtest_read_buffer_cancel.cpp b/src/IO/tests/gtest_read_buffer_cancel.cpp
@@ -0,0 +1,74 @@
+#include <gtest/gtest.h>
+
+#include <string>
+#include <IO/ReadBuffer.h>
+#include <IO/ReadBufferFromString.h>
+
+/// Regression test for the "ReadBuffer is canceled" crash (STID 2508-2913).
+///
+/// Bug: TCPHandler::runImpl() idle loop called in->eof() on a ReadBuffer that
+/// had been canceled by a previous query's pipeline callback. In debug/sanitizer
+/// builds, ReadBuffer::next() has chassert(!isCanceled()) which aborts the
+/// server process.
+///
+/// Fix: Added ReadBuffer::eofOrCanceled() which checks isCanceled() BEFORE
+/// calling eof(). TCPHandler now uses eofOrCanceled() in the idle loop.
+
+using namespace DB;
+
+/// This test calls the PRODUCTION method ReadBuffer::eofOrCanceled() on a
+/// canceled buffer. Without the isCanceled() check inside eofOrCanceled(),
+/// eof() would be called on the canceled buffer, triggering
+/// chassert(!isCanceled()) → abort → test FAILS (crash).
+///
+/// With the fix (isCanceled() short-circuits before eof()), the test PASSES.
+TEST(ReadBufferCancelTest, EofOrCanceledOnCanceledBuffer)
+{
+    std::string data = "some data in the buffer";
+    ReadBufferFromString buf(data);
+
+    /// Consume all pending data so hasPendingData() returns false.
+    /// This simulates a TCP ReadBuffer where all received bytes have been
+    /// processed and eof() would need to call next() to check for more.
+    buf.position() = buf.buffer().end();
+
+    /// Simulate what happens when a callback's read from the TCP buffer
+    /// fails: ReadBuffer::next() catches the exception from nextImpl()
+    /// and calls cancel(), setting canceled = true.
+    buf.cancel();
+
+    /// This is the exact method called by TCPHandler::runImpl() idle loop.
+    /// If eofOrCanceled() doesn't check isCanceled() first, this crashes
+    /// with "ReadBuffer is canceled" in debug/sanitizer builds.
+    EXPECT_TRUE(buf.eofOrCanceled());
+}
+
+/// Verify eofOrCanceled() behaves like eof() on a non-canceled buffer.
+TEST(ReadBufferCancelTest, EofOrCanceledOnNormalBuffer)
+{
+    std::string data = "test data";
+    ReadBufferFromString buf(data);
+
+    /// Buffer has pending data — not at EOF
+    EXPECT_FALSE(buf.eofOrCanceled());
+
+    /// Consume all data
+    buf.position() = buf.buffer().end();
+
+    /// Now at EOF (no more data, not canceled)
+    EXPECT_TRUE(buf.eofOrCanceled());
+}
+
+/// Verify that cancel() is persistent and isCanceled() reflects the state.
+TEST(ReadBufferCancelTest, CancelIsPersistent)
+{
+    std::string data = "test";
+    ReadBufferFromString buf(data);
+
+    EXPECT_FALSE(buf.isCanceled());
+    buf.cancel();
+    EXPECT_TRUE(buf.isCanceled());
+
+    /// Cancel is permanent — by design, a canceled TCP connection cannot be reused
+    EXPECT_TRUE(buf.isCanceled());
+}
diff --git a/src/Processors/Executors/PullingAsyncPipelineExecutor.cpp b/src/Processors/Executors/PullingAsyncPipelineExecutor.cpp
@@ -115,8 +115,13 @@ bool PullingAsyncPipelineExecutor::pull(Chunk & chunk, uint64_t milliseconds)
     {
         /// If lazy format is finished, we don't cancel pipeline but wait for main thread to be finished.
         data->is_finished = true;
-        /// Wait thread and rethrow exception if any.
+        /// Wait for the pipeline thread to finish.
         cancel();
+        /// Rethrow any exception from the pipeline thread. Without this check,
+        /// exceptions from pipeline callbacks (e.g. ClusterFunctionReadTaskCallback
+        /// canceling the TCP ReadBuffer) would be silently lost when pull() returns
+        /// false, causing the caller to believe the query succeeded.
+        data->rethrowExceptionIfHas();
         return false;
     }
 
diff --git a/src/Server/TCPHandler.cpp b/src/Server/TCPHandler.cpp
@@ -492,9 +492,14 @@ void TCPHandler::runImpl()
             }
 
             /// If we need to shut down, or client disconnects.
-            if (!tcp_server.isOpen() || server.isCancelled() || in->eof())
+            /// Use eofOrCanceled() instead of eof() as a safety net: if a previous
+            /// query's pipeline callback canceled the ReadBuffer but the exception
+            /// was not propagated, we must close the connection rather than hitting
+            /// chassert(!isCanceled()) inside eof().
+            if (!tcp_server.isOpen() || server.isCancelled() || in->eofOrCanceled())
             {
-                LOG_TEST(log, "Closing connection (open: {}, cancelled: {}, eof: {})", tcp_server.isOpen(), server.isCancelled(), in->eof());
+                LOG_TEST(log, "Closing connection (open: {}, cancelled: {}, eof_or_canceled: {})",
+                    tcp_server.isOpen(), server.isCancelled(), in->eofOrCanceled());
                 return;
             }
         }

Original file line number	Diff line number	Diff line change
`@@ -115,8 +115,13 @@ bool PullingAsyncPipelineExecutor::pull(Chunk & chunk, uint64_t milliseconds)`
`115`	`115`	`{`
`116`	`116`	`/// If lazy format is finished, we don't cancel pipeline but wait for main thread to be finished.`
`117`	`117`	`data->is_finished = true;`
`118`		`- /// Wait thread and rethrow exception if any.`
	`118`	`+ /// Wait for the pipeline thread to finish.`
`119`	`119`	`cancel();`
	`120`	`+ /// Rethrow any exception from the pipeline thread. Without this check,`
	`121`	`+ /// exceptions from pipeline callbacks (e.g. ClusterFunctionReadTaskCallback`
	`122`	`+ /// canceling the TCP ReadBuffer) would be silently lost when pull() returns`
	`123`	`+ /// false, causing the caller to believe the query succeeded.`
	`124`	`+ data->rethrowExceptionIfHas();`
`120`	`125`	`return false;`
`121`	`126`	`}`
`122`	`127`
Original file line number	Diff line number	Diff line change
`@@ -492,9 +492,14 @@ void TCPHandler::runImpl()`
`492`	`492`	`}`
`493`	`493`
`494`	`494`	`/// If we need to shut down, or client disconnects.`
`495`		`- if (!tcp_server.isOpen() \|\| server.isCancelled() \|\| in->eof())`
	`495`	`+ /// Use eofOrCanceled() instead of eof() as a safety net: if a previous`
	`496`	`+ /// query's pipeline callback canceled the ReadBuffer but the exception`
	`497`	`+ /// was not propagated, we must close the connection rather than hitting`
	`498`	`+ /// chassert(!isCanceled()) inside eof().`
	`499`	`+ if (!tcp_server.isOpen() \|\| server.isCancelled() \|\| in->eofOrCanceled())`
`496`	`500`	`{`
`497`		`- LOG_TEST(log, "Closing connection (open: {}, cancelled: {}, eof: {})", tcp_server.isOpen(), server.isCancelled(), in->eof());`
	`501`	`+ LOG_TEST(log, "Closing connection (open: {}, cancelled: {}, eof_or_canceled: {})",`
	`502`	`+ tcp_server.isOpen(), server.isCancelled(), in->eofOrCanceled());`
`498`	`503`	`return;`
`499`	`504`	`}`
`500`	`505`	`}`