diff --git a/CHANGELOG.md b/CHANGELOG.md index e5a8207..ae201d4 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/SPEC.md b/docs/SPEC.md index 69339e0..b8ca94a 100644 --- a/docs/SPEC.md +++ b/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 | diff --git a/pyproject.toml b/pyproject.toml index 73d9c23..5ad4aed 100644 --- a/pyproject.toml +++ b/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 diff --git a/src/streamcast/_client.py b/src/streamcast/_client.py index 180340d..fb66b6a 100644 --- a/src/streamcast/_client.py +++ b/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. diff --git a/tests/test_batches.py b/tests/test_batches.py index 50e1b9f..d188b1c 100644 --- a/tests/test_batches.py +++ b/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,7 +170,9 @@ 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) @@ -140,15 +180,28 @@ async def test_a_row_read_ahead_is_not_received(self, tmp_path): 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: diff --git a/uv.lock b/uv.lock index 02cad09..5ecae1e 100644 --- a/uv.lock +++ b/uv.lock @@ -2006,7 +2006,7 @@ requires-dist = [ { name = "litelink", specifier = ">=0.10.1,<0.11" }, { name = "msgspec", specifier = ">=0.18" }, { name = "starlette", marker = "extra == 'asgi'", specifier = ">=0.40" }, - { name = "websockets", specifier = ">=14" }, + { name = "websockets", specifier = ">=14,<18" }, ] provides-extras = ["asgi"]