perf(subscriber): take arrived rows without a task per row

CHANGELOG.md

@@ -20,6 +20,16 @@ there was nothing to have changed from. Everything above it is ordinary.
   `empty` from an older server. Found as an intermittent CI failure in the
   keyed-table example, and the same race was in `Stream.live` at open.
 
+- **`recv_many` and `batches` cost less per row than `recv`, as documented.**
+  Each row added to a batch started a task and gave the loop a turn to learn
+  whether it had arrived: 10.6–13.9 µs of client CPU a row, against `recv`'s
+  ~2.5. It now looks at `websockets`' queue of received frames instead, and
+  costs 1.3 µs a row against `recv`'s 1.45. Rows from the socket are no longer
+  read ahead between batches; a catch-up still reads one ahead.
+  - **`websockets` is capped below 18:** that queue is private API.
+    `tests/test_batches.py` checks its shape against the installed release, so
+    raising the cap is a matter of that test passing.
+
 ## 0.14.0 — 2026-10-04
 
 ### Added

docs/SPEC.md

@@ -491,7 +491,7 @@ Neither end can be run out of memory by the other, or by a stall:
 | broker | a replay | one batch at a time, into the bounded subscriber queue | — |
 | client | a publisher's unanswered sends | `publish(max_in_flight=)`, 64 frames | `submit` waits |
 | client | a catch-up | one batch at a time | — |
-| client | a `recv_many` batch, and the read it leaves waiting | `limit` rows, plus one | — |
+| client | a `recv_many` batch, and the read it leaves waiting | `limit` rows, plus one during a catch-up | — |
 | client | a snapshot's rows from the broker | `max_tail`, 1,000,000 rows | the snapshot is refused |
 | client | a live view's unpublished rows | `max_tail`, 1,000,000 rows | the view stops; its next query raises |
 

pyproject.toml

@@ -26,8 +26,10 @@ dependencies = [
     # floor because `websockets.asyncio` is what everything here imports: the
     # legacy `websockets.server`/`websockets.client` implementation is
     # deprecated, and the two have different connection objects, so a build
-    # that resolved 12 would import successfully and fail at `serve`.
-    "websockets>=14",
+    # that resolved 12 would import successfully and fail at `serve`. The cap
+    # is because `recv_many` reads a private queue (`_client._arrived`):
+    # raise it once `tests/test_batches.py` passes against the next major.
+    "websockets>=14,<18",
     # The durable tier. NOT optional, and that was a real decision: a
     # live-only broker needs nothing from it, and litelink pulls in duckdb,
     # pyarrow and pyiceberg, which is ~124 MB of wheel for a library whose own

src/streamcast/_client.py

@@ -290,8 +290,9 @@ def __init__(
         self._received: int | None = None
         # A read `recv_many` started and did not wait for: the next call's
         # first row. Kept rather than cancelled, because a cancelled read of
-        # the catch-up rows would end their generator (see `recv_many`).
-        self._ahead: asyncio.Task[Row] | None = None
+        # the catch-up rows would end their generator (see `recv_many`). Or a
+        # read that failed mid-batch, for the next call to raise.
+        self._ahead: asyncio.Future[Row] | None = None
 
     def __repr__(self) -> str:
         return (
@@ -406,12 +407,33 @@ async def recv_many(self, limit: int = BATCH) -> list[Row]:
 
         batch = [await self._take()]
         while len(batch) < limit:
-            # Whether a row is ready, asked the only way the public APIs
-            # allow: start a read and give it one turn of the loop. A ready
+            if self._prelude is None:
+                # From the socket, whether a row has arrived is a look at
+                # `websockets`' queue — no task and no turn of the loop, which
+                # cost the asking more per row than `recv` (measured at
+                # 16-36 µs a row against 12 on a 20k-row replay). The read
+                # cannot suspend, so a row that is not there is never started:
+                # nothing is held outside the queue between batches.
+                if not _arrived(self._live()):
+                    break
+
+                try:
+                    batch.append(await self._read())
+                except Exception as exc:  # noqa: BLE001 — raised by the next call
+                    failed: asyncio.Future[Row] = (
+                        asyncio.get_running_loop().create_future()
+                    )
+                    failed.set_exception(exc)
+                    self._ahead = failed
+                    break
+
+                continue
+
+            # From the catch-up generator, asked the only way its interface
+            # allows: start a read and give it one turn of the loop. A ready
             # row is taken; an unready read is KEPT as the next call's first
-            # row rather than cancelled, so it costs nothing, and a read of
-            # the catch-up generator — which a cancel would end — is never
-            # cancelled. One row at most is held outside `websockets`' queue.
+            # row rather than cancelled, since a cancel would end the
+            # generator. One row at most is held outside it.
             ahead = asyncio.ensure_future(self._read())
             await asyncio.sleep(0)
             if not ahead.done() or ahead.exception() is not None:
@@ -623,6 +645,22 @@ async def close(self, code: int = 1000, reason: str = "") -> None:
             await _close(self._connection, code, reason)
 
 
+def _arrived(connection: ClientConnection) -> bool:
+    """Whether a whole message is in `websockets`' queue, so `recv` returns it
+    without suspending.
+
+    **Private `websockets` API**: `Assembler.frames`, a `SimpleQueue` whose
+    `get` pops without awaiting when its deque is not empty — the same from 14
+    to 17, and the reason `pyproject.toml` caps the version.
+    `tests/test_batches.py` checks the shape against the installed release, so
+    an upgrade that moves it fails there rather than here. Only the head frame
+    is looked at: a fragmented message reads as not arrived, which ends the
+    batch early and loses nothing. The server never fragments a row.
+    """
+    frames = connection.recv_messages.frames.queue
+    return bool(frames) and frames[0].fin
+
+
 async def _close(connection: ClientConnection, code: int, reason: str) -> None:
     """Close, draining what is still in flight so the handshake can finish.
 

tests/test_batches.py

@@ -12,29 +12,43 @@
 
 import asyncio
 import collections
+from types import SimpleNamespace
 
 import pytest
 from websockets.exceptions import ConnectionClosed
 from websockets.frames import Close as CloseFrame
+from websockets.frames import Frame, Opcode
 
 import streamcast
-from streamcast._client import Subscription
+from streamcast._client import Subscription, _arrived
 from streamcast._errors import Close
 from streamcast._protocol import encode_projected, greeting, parse_greeting, refusal
 
+SCHEMA = {
+    "type": "object",
+    "properties": {"i": {"type": "integer"}},
+    "required": ["i"],
+}
+
 
 class FakeConnection:
-    """The two things `Subscription` asks of a connection: `recv` and `close`."""
+    """What `Subscription` asks of a connection: `recv`, `close`, and a look
+    at the queue of frames, at the same private path `websockets` keeps it.
+
+    `test_the_queue_it_looks_at_is_websockets` holds that path to the real one.
+    """
 
     def __init__(self) -> None:
-        self._frames: collections.deque[bytes] = collections.deque()
+        self._frames: collections.deque[Frame] = collections.deque()
+        self.recv_messages = SimpleNamespace(frames=SimpleNamespace(queue=self._frames))
         self._closed: ConnectionClosed | None = None
         self._arrived = asyncio.Event()
         self.reads = 0
 
     def arrive(self, *offsets: int) -> None:
         for offset in offsets:
-            self._frames.append(encode_projected(offset, 0, {"i": offset}))
+            data = encode_projected(offset, 0, {"i": offset})
+            self._frames.append(Frame(Opcode.TEXT, data))
 
         self._arrived.set()
 
@@ -51,7 +65,19 @@ async def recv(self, decode: bool | None = None) -> bytes:  # noqa: ARG002
             await self._arrived.wait()  # cancelling here takes nothing
 
         self.reads += 1
-        return self._frames.popleft()
+        # A fragmented message, reassembled as `websockets` does.
+        parts = [self._frames.popleft()]
+        while not parts[-1].fin:
+            parts.append(self._frames.popleft())
+
+        return b"".join(bytes(part.data) for part in parts)
+
+    def arrive_fragmented(self, offset: int) -> None:
+        """One message in two frames: whole once the second is here."""
+        data = encode_projected(offset, 0, {"i": offset})
+        self._frames.append(Frame(Opcode.TEXT, data[:5], fin=False))
+        self._frames.append(Frame(Opcode.CONT, data[5:]))
+        self._arrived.set()
 
     async def close(self, code: int = 1000, reason: str = "") -> None:  # noqa: ARG002
         self.end(code)
@@ -108,6 +134,18 @@ async def test_the_read_left_waiting_is_the_next_batchs_first_row(self):
         connection.arrive(4)
         assert (await sub.recv())[0] == 4  # `recv` takes it too
 
+    async def test_a_fragmented_message_ends_the_batch_and_loses_nothing(self):
+        """Only the head frame is looked at, so a message in fragments reads as
+        not arrived: the batch ends before it, and the next one starts with it."""
+        connection = FakeConnection()
+        connection.arrive(1)
+        connection.arrive_fragmented(2)
+        connection.arrive(3)
+        sub = subscription(connection)
+
+        assert offsets(await sub.recv_many()) == [1]
+        assert offsets(await sub.recv_many()) == [2, 3]
+
     def test_a_batch_holds_at_least_one_row(self):
         sub = subscription(FakeConnection())
         with pytest.raises(ValueError, match="limit=0"):
@@ -132,23 +170,38 @@ async def test_a_batch_is_saved_when_the_next_is_asked_for(self, tmp_path):
         await sub.recv_many()
         assert cursor.load() == 3  # the whole first batch, and no further
 
-    async def test_a_row_read_ahead_is_not_received(self, tmp_path):
+    async def test_from_the_socket_nothing_is_read_ahead(self, tmp_path):
+        """A row that has not arrived is never started: between batches,
+        everything not handed over is still in `websockets`' queue."""
         cursor = streamcast.Cursor(tmp_path / "cursor", every=0)
         connection = FakeConnection()
         connection.arrive(1)
         sub = subscription(connection, cursor)
 
         assert offsets(await sub.recv_many()) == [1]
         connection.arrive(2)
-        await asyncio.sleep(0.01)  # the waiting read takes row 2 off the socket
-        assert connection.reads == 2
+        await asyncio.sleep(0.01)
+        assert connection.reads == 1, "row 2 is still the socket's"
 
-        # But the caller never got it: it is not the offset, and a clean
-        # exit's commit does not save it.
         assert sub.offset == 1
         sub.commit()
         assert cursor.load() == 1
-        await sub.close()  # and closing with it still held is quiet
+        await sub.close()
+
+    async def test_a_batch_from_the_socket_starts_no_task(self, monkeypatch):
+        """What it costs: rows already arrived are taken without a task or a
+        turn of the loop, which had made a batch dearer per row than `recv`."""
+        connection = FakeConnection()
+        connection.arrive(*range(1, 101))
+        sub = subscription(connection)
+
+        def refused(*_: object, **__: object) -> None:
+            msg = "a task per row"
+            raise AssertionError(msg)
+
+        monkeypatch.setattr(asyncio, "ensure_future", refused)
+        monkeypatch.setattr(asyncio, "sleep", refused)
+        assert offsets(await sub.recv_many(limit=500)) == list(range(1, 101))
 
     async def test_a_failure_mid_batch_hands_over_the_rows_before_it(self):
         connection = FakeConnection()
@@ -203,13 +256,29 @@ async def test_a_handler_that_raises_leaves_its_batch_unsaved(self, tmp_path):
 
 
 class TestAgainstAServer:
+    async def test_the_queue_it_looks_at_is_websockets(self, tmp_path, serve):
+        """`_arrived` reads a private queue: this is what says the installed
+        `websockets` still keeps it where it did, and pops it the same way."""
+        stream = streamcast.Stream.new("t", root=tmp_path, schema=SCHEMA)
+        try:
+            async with serve(stream, maintain=False) as uri:
+                async with streamcast.connect(uri) as sub:
+                    assert not _arrived(sub.connection)
+                    await stream.send_many([{"i": i} for i in range(3)])
+                    for _ in range(100):
+                        if len(sub.connection.recv_messages.frames) == 3:
+                            break
+
+                        await asyncio.sleep(0.01)
+
+                    assert _arrived(sub.connection)
+                    assert offsets(await sub.recv_many()) == [1, 2, 3]
+                    assert not _arrived(sub.connection)
+        finally:
+            await stream.aclose()
+
     async def test_every_row_once_in_order_within_the_limit(self, tmp_path, serve):
-        schema = {
-            "type": "object",
-            "properties": {"i": {"type": "integer"}},
-            "required": ["i"],
-        }
-        stream = streamcast.Stream.new("t", root=tmp_path, schema=schema)
+        stream = streamcast.Stream.new("t", root=tmp_path, schema=SCHEMA)
         await stream.send_many([{"i": i} for i in range(300)])
         try:
             async with serve(stream, maintain=False) as uri: