fix(asyncio): preserve post-delimiter bytes in _read_until

The asyncio backend's _read_until previously called streamreader.read(1024)
in a loop and, when the delimiter appeared mid-chunk, kept only the bytes
up to and including the delimiter, silently discarding any trailing bytes
from the same chunk. The original author flagged the limitation in an
inline comment ("for production use, consider asyncio.StreamReader
replacement").

Under load this surfaced as ReqlDriverError "Connection interrupted during
handshake". The V1.0 SCRAM-SHA-256 handshake sends the version banner and
the first JSON message in one initial write, then expects to read two
null-terminated server replies back-to-back without an intermediate write
(state 1 returns ""). When many connections handshake concurrently, the
typical pattern when warming up a connection pool, both server replies
can land in a single TCP frame. The legacy code returned the first reply
and dropped the second; the next _read_until call then blocked on EOF or
raised the misleading "interrupted" error.

The fix delegates to asyncio.StreamReader.readuntil, which leaves trailing
bytes in the reader's internal buffer for the next call. The EOF contract
"return whatever was read" is preserved by mapping
asyncio.IncompleteReadError to its .partial bytes.

Tests in tests/test_asyncio_read_until.py are rewritten to drive a real
asyncio.StreamReader (via feed_data / feed_eof) so they exercise the
buffering path that production hits, and two regression tests are added
that fail on the legacy implementation:

* test_post_delimiter_bytes_preserved_for_next_call directly pins the
  IsardVDI bug scenario (two pipelined handshake replies in one frame).
* test_three_pipelined_messages_in_one_frame extends the contract to more
  than two messages so future protocol changes do not regress.

Author: vinolas

rethinkdb/asyncio_net/net_asyncio.py

@@ -71,28 +71,42 @@ def __setitem__(self, key, value):
 
 
 async def _read_until(streamreader, delimiter):
-    """Optimized implementation of reading until a delimiter"""
-    buffer = bytearray()
-
-    # Read in chunks for better performance
-    chunk_size = 1024
-
-    while True:
-        chunk = await streamreader.read(chunk_size)
-        if not chunk:
-            break  # EOF
-
-        delimiter_pos = chunk.find(delimiter)
-        if delimiter_pos >= 0:
-            # Found delimiter, add up to and including delimiter
-            buffer.extend(chunk[: delimiter_pos + len(delimiter)])
-            # Put back remaining data (if any) - not directly possible with asyncio StreamReader
-            # This is a limitation; for production use, consider asyncio.StreamReader replacement
-            break
-        else:
-            buffer.extend(chunk)
-
-    return bytes(buffer)
+    """Read from an :class:`asyncio.StreamReader` up to and including
+    ``delimiter``.
+
+    Defers to :meth:`asyncio.StreamReader.readuntil` so any bytes that
+    arrive **after** the delimiter in the same TCP frame stay in the
+    reader's internal buffer and are available to the next
+    ``_read_until`` call. The previous implementation called
+    ``streamreader.read(1024)`` and silently discarded any post-delimiter
+    bytes; the only hint that this was lossy was an inline comment
+    saying "for production use, consider asyncio.StreamReader
+    replacement". In practice it manifested as
+    ``ReqlDriverError: Connection interrupted during handshake`` when a
+    consumer (typically a connection pool warming up) opened many
+    asyncio connections concurrently. The V1.0 handshake's two-step
+    pipelined JSON exchange can deliver both messages in the same TCP
+    frame under load, after which the original chunked read consumed
+    the first message and dropped the second on the floor — the next
+    read would then see EOF or block forever waiting for bytes that
+    had already been delivered.
+
+    Uses stdlib :meth:`StreamReader.readuntil`, which buffers correctly:
+    only the delimited prefix is consumed; trailing bytes remain
+    available for subsequent reads.
+
+    Behaviour preserved from the legacy implementation:
+
+    * Returns whatever bytes were read (including a trailing delimiter
+      if one was found) as a :class:`bytes` instance.
+    * On EOF before the delimiter is encountered, returns the partial
+      bytes read so far rather than raising — the caller's protocol
+      decoder is in charge of reporting "incomplete frame" errors.
+    """
+    try:
+        return await streamreader.readuntil(delimiter)
+    except asyncio.IncompleteReadError as err:
+        return err.partial
 
 
 def reusable_waiter(loop, timeout):

tests/test_asyncio_read_until.py

@@ -1,12 +1,15 @@
-"""Edge-case tests for the asyncio backend's chunked _read_until() helper.
-
-The fork rewrote _read_until to do 1 KB chunked reads instead of upstream's
-byte-by-byte loop. The optimisation is correct on a normal network but has a
-documented edge case: if the delimiter is the last byte of one TCP frame
-and the next bytes belong to a different message, the chunked read could
-overshoot. RethinkDB's handshake protocol uses null-terminated framing
-where the delimiter only appears at message ends, so in practice this is
-safe — but we want a regression net for it.
+"""Edge-case tests for the asyncio backend's ``_read_until`` helper.
+
+The asyncio handshake reads null-terminated JSON messages with this helper.
+The legacy implementation did its own chunked reads and silently discarded
+any bytes that arrived after the delimiter in the same TCP frame; under
+connection-pool warm-up the V1.0 handshake's pipelined replies could land
+in one frame, the post-delimiter half got dropped, and the next read
+surfaced as ``ReqlDriverError: Connection interrupted during handshake``.
+
+These tests pin the corrected behaviour: the helper now defers to
+``StreamReader.readuntil`` so trailing bytes remain in the reader's
+internal buffer for the next call.
 """
 
 import asyncio
@@ -16,93 +19,111 @@
 from rethinkdb.asyncio_net.net_asyncio import _read_until
 
 
-class FakeStreamReader:
-    """StreamReader stand-in that hands out pre-staged byte chunks.
-
-    Each call to ``read(n)`` returns the next staged chunk verbatim,
-    regardless of ``n`` — the implementation should still find the
-    delimiter and stop, even if a chunk is split across what would be
-    multiple TCP frames.
+def _make_reader(*chunks: bytes, eof: bool = True) -> asyncio.StreamReader:
+    """Build a fully-buffered :class:`asyncio.StreamReader` from the given
+    chunks. Equivalent to feeding ``data_received`` callbacks in the
+    transport layer — but with the data already buffered, the reader
+    behaves exactly like a real socket that received the same frames.
     """
-
-    def __init__(self, chunks):
-        self._chunks = list(chunks)
-        self._read_calls = 0
-
-    async def read(self, n):
-        self._read_calls += 1
-        if not self._chunks:
-            return b""
-        return self._chunks.pop(0)
+    reader = asyncio.StreamReader()
+    for chunk in chunks:
+        reader.feed_data(chunk)
+    if eof:
+        reader.feed_eof()
+    return reader
 
 
 @pytest.mark.unit
 class TestReadUntil:
     @pytest.mark.asyncio
     async def test_delimiter_in_single_chunk(self):
-        reader = FakeStreamReader([b"hello\0world", b""])
+        reader = _make_reader(b"hello\0world")
         result = await _read_until(reader, b"\0")
         assert result == b"hello\0"
-        # Should stop at the delimiter; only one read call needed.
-        assert reader._read_calls == 1
 
     @pytest.mark.asyncio
     async def test_delimiter_split_across_chunks(self):
-        # The handshake response is delivered in two halves, with the
-        # delimiter at the start of the second chunk.
-        reader = FakeStreamReader([b'{"success":true}', b"\0", b""])
+        # Handshake response delivered in two halves with the delimiter
+        # at the start of the second chunk.
+        reader = _make_reader(b'{"success":true}', b"\0")
         result = await _read_until(reader, b"\0")
         assert result == b'{"success":true}\0'
 
     @pytest.mark.asyncio
     async def test_delimiter_at_frame_boundary(self):
-        # Delimiter is the last byte of the chunk — common case for short
-        # ack messages on a slow network.
-        reader = FakeStreamReader([b"ack\0", b""])
+        reader = _make_reader(b"ack\0")
         result = await _read_until(reader, b"\0")
         assert result == b"ack\0"
 
     @pytest.mark.asyncio
     async def test_eof_before_delimiter(self):
-        # Connection died mid-message; we get back what we read, no infinite loop.
-        reader = FakeStreamReader([b"partial", b""])
+        # Connection died mid-message — return what we have, no exception.
+        reader = _make_reader(b"partial")
         result = await _read_until(reader, b"\0")
         assert result == b"partial"
 
     @pytest.mark.asyncio
     async def test_immediate_eof(self):
-        reader = FakeStreamReader([b""])
+        reader = _make_reader()
         result = await _read_until(reader, b"\0")
         assert result == b""
 
     @pytest.mark.asyncio
     async def test_multibyte_delimiter(self):
-        reader = FakeStreamReader([b"prefix\r\nsuffix", b""])
+        reader = _make_reader(b"prefix\r\nsuffix\r\n")
         result = await _read_until(reader, b"\r\n")
         assert result == b"prefix\r\n"
 
     @pytest.mark.asyncio
     async def test_handshake_split_with_concurrent_yields(self):
-        """Simulate slow network: each chunk arrives after an event-loop tick.
+        """Bytes arriving after event-loop ticks must be reassembled."""
+        reader = asyncio.StreamReader()
+
+        async def feeder() -> None:
+            msg = (
+                b'{"success":true,"min_protocol_version":0,'
+                b'"max_protocol_version":1}\0'
+            )
+            for i in range(0, len(msg), 12):
+                # Yield to the loop between feeds — simulates separate
+                # ``data_received`` callbacks on the transport.
+                await asyncio.sleep(0)
+                reader.feed_data(msg[i : i + 12])
+            reader.feed_eof()
+
+        feed_task = asyncio.create_task(feeder())
+        result = await _read_until(reader, b"\0")
+        await feed_task
+        assert result == (
+            b'{"success":true,"min_protocol_version":0,' b'"max_protocol_version":1}\0'
+        )
 
-        This is the scenario the fork's IsardVDI evaluation flagged — handshake
-        bytes split across multiple data_received callbacks. The chunked read
-        must correctly accumulate them and find the terminating null.
+    @pytest.mark.asyncio
+    async def test_post_delimiter_bytes_preserved_for_next_call(self):
+        """Regression for the IsardVDI handshake-race bug.
+
+        When two null-terminated messages arrive in one TCP frame (the
+        V1.0 handshake's pipelined replies under pool warm-up), the
+        bytes after the first delimiter MUST remain in the stream for
+        the next ``_read_until`` call. The previous chunked
+        implementation silently discarded them and the second
+        handshake step would block on EOF.
         """
+        reader = _make_reader(b"first-message\0second-message\0")
 
-        class SlowReader:
-            def __init__(self, chunks):
-                self._chunks = list(chunks)
+        first = await _read_until(reader, b"\0")
+        second = await _read_until(reader, b"\0")
 
-            async def read(self, n):
-                if not self._chunks:
-                    return b""
-                # Yield to the event loop between chunks so other tasks run.
-                await asyncio.sleep(0)
-                return self._chunks.pop(0)
+        assert first == b"first-message\0"
+        assert second == b"second-message\0"
 
-        # Mimic a 60-byte JSON handshake response delivered in 5 frames.
-        msg = b'{"success":true,"min_protocol_version":0,"max_protocol_version":1}\0'
-        reader = SlowReader([msg[i : i + 12] for i in range(0, len(msg), 12)])
-        result = await _read_until(reader, b"\0")
-        assert result == msg
+    @pytest.mark.asyncio
+    async def test_three_pipelined_messages_in_one_frame(self):
+        """Stretch case: SCRAM-SHA-256 doesn't pipeline three replies,
+        but the helper's contract is "consume exactly one delimited
+        message per call" so verify it scales beyond two."""
+        reader = _make_reader(b"a\0bb\0ccc\0")
+
+        assert (await _read_until(reader, b"\0")) == b"a\0"
+        assert (await _read_until(reader, b"\0")) == b"bb\0"
+        assert (await _read_until(reader, b"\0")) == b"ccc\0"