diff --git a/docs/source/api-reference.md b/docs/source/api-reference.md index 9ae8366..d45607d 100644 --- a/docs/source/api-reference.md +++ b/docs/source/api-reference.md @@ -22,6 +22,9 @@ .. autoclass:: BatchSubmitTicket() :members: +.. autoclass:: ReadSession() + :members: + .. autoclass:: Producer() :members: diff --git a/s2-specs b/s2-specs index 329de93..973e0d9 160000 --- a/s2-specs +++ b/s2-specs @@ -1 +1 @@ -Subproject commit 329de93f7b240a4daef9edbeb98ced0699aab7d0 +Subproject commit 973e0d92166ee1386b71d4e7ade8bd0fc2aaf4b8 diff --git a/src/s2_sdk/__init__.py b/src/s2_sdk/__init__.py index 0950512..fde2544 100644 --- a/src/s2_sdk/__init__.py +++ b/src/s2_sdk/__init__.py @@ -14,6 +14,7 @@ ) from s2_sdk._ops import S2, S2Basin, S2Stream from s2_sdk._producer import Producer, RecordSubmitTicket +from s2_sdk._read_session import ReadSession from s2_sdk._types import ( AccessTokenInfo, AccessTokenScope, @@ -98,6 +99,7 @@ "Timestamp", "TailOffset", "ReadBatch", + "ReadSession", "ReadLimit", "SequencedRecord", "Page", diff --git a/src/s2_sdk/_append_session.py b/src/s2_sdk/_append_session.py index 75750bd..5c7094f 100644 --- a/src/s2_sdk/_append_session.py +++ b/src/s2_sdk/_append_session.py @@ -28,6 +28,8 @@ class AppendSession: Supports pipelining multiple :class:`AppendInput`\\ s while preserving submission order. + Use it as an async context manager, or call :meth:`close` explicitly to close the session. + Caution: Returned by :meth:`S2Stream.append_session`. Do not instantiate directly. """ diff --git a/src/s2_sdk/_client.py b/src/s2_sdk/_client.py index b9d9241..892f0a2 100644 --- a/src/s2_sdk/_client.py +++ b/src/s2_sdk/_client.py @@ -5,7 +5,7 @@ import ssl import time from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, suppress from dataclasses import dataclass, field from importlib.metadata import version from typing import Any, NamedTuple @@ -423,10 +423,8 @@ async def close(self) -> None: self._closed = True if self._reaper_task is not None: self._reaper_task.cancel() - try: + with suppress(asyncio.CancelledError): await self._reaper_task - except asyncio.CancelledError: - pass for conns in self._hosts.values(): for pc in conns: await pc.close() diff --git a/src/s2_sdk/_ops.py b/src/s2_sdk/_ops.py index 6ddfc0e..b2700f9 100644 --- a/src/s2_sdk/_ops.py +++ b/src/s2_sdk/_ops.py @@ -2,7 +2,7 @@ import uuid from collections.abc import AsyncIterator from datetime import datetime -from typing import Any, AsyncIterable, Self +from typing import Any, Self from urllib.parse import quote import s2_sdk._generated.s2.v1.s2_pb2 as pb @@ -33,6 +33,7 @@ tail_from_json, ) from s2_sdk._producer import Producer +from s2_sdk._read_session import ReadSession from s2_sdk._retrier import Retrier, http_retry_on, is_safe_to_retry_unary from s2_sdk._s2s._read_session import run_read_session from s2_sdk._types import ( @@ -1172,7 +1173,7 @@ async def read( return batch @fallible - async def read_session( + def read_session( self, *, start: types.SeqNum | types.Timestamp | types.TailOffset, @@ -1181,7 +1182,7 @@ async def read_session( clamp_to_tail: bool = False, wait: int | None = None, ignore_command_records: bool = False, - ) -> AsyncIterable[types.ReadBatch]: + ) -> ReadSession: """Read batches of records from a stream continuously. Args: @@ -1196,9 +1197,8 @@ async def read_session( reached. ignore_command_records: Filter out command records from batches. - Yields: - :class:`ReadBatch` — each containing a batch of records and an - optional tail position. + Returns: + A :class:`ReadSession` that yields batches of records. Note: Sessions without bounds (no ``limit`` or ``until_timestamp``) default @@ -1207,19 +1207,20 @@ async def read_session( ``wait`` makes a bounded session wait up to that many seconds for new records before ending. """ - async for batch in run_read_session( - self._client, - self.name, - start, - limit, - until_timestamp, - clamp_to_tail, - wait, - ignore_command_records, - retry=self._retry, - encryption_key=self._encryption_key, - ): - yield batch + return ReadSession( + run_read_session( + self._client, + self.name, + start, + limit, + until_timestamp, + clamp_to_tail, + wait, + retry=self._retry, + encryption_key=self._encryption_key, + ), + ignore_command_records=ignore_command_records, + ) def _s2_request_token() -> str: diff --git a/src/s2_sdk/_producer.py b/src/s2_sdk/_producer.py index 6014ce5..986b6d6 100644 --- a/src/s2_sdk/_producer.py +++ b/src/s2_sdk/_producer.py @@ -33,6 +33,8 @@ class Producer: Handles batching into :class:`AppendInput` automatically and uses an append session internally. + Use it as an async context manager, or call :meth:`close` explicitly to close the producer. + Caution: Returned by :meth:`S2Stream.producer`. Do not instantiate directly. """ diff --git a/src/s2_sdk/_read_session.py b/src/s2_sdk/_read_session.py new file mode 100644 index 0000000..9396ae0 --- /dev/null +++ b/src/s2_sdk/_read_session.py @@ -0,0 +1,245 @@ +from __future__ import annotations + +import asyncio +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable +from contextlib import suppress +from dataclasses import dataclass, field +from typing import Self + +from s2_sdk._exceptions import S2ClientError, normalize_exception +from s2_sdk._types import ReadBatch, StreamPosition + + +class _ReadSessionRetrying: + pass + + +@dataclass(frozen=True, slots=True) +class _ReadSessionHeartbeat: + tail: StreamPosition + + +@dataclass(frozen=True, slots=True) +class _ReadSessionBatch: + batch: ReadBatch + + +_ReadSessionEvent = _ReadSessionRetrying | _ReadSessionHeartbeat | _ReadSessionBatch + + +@dataclass(slots=True) +class _ReadDelivery: + batch: ReadBatch + caught_up_tail: StreamPosition | None + delivered: asyncio.Event = field(default_factory=asyncio.Event) + + +class ReadSession(AsyncIterator[ReadBatch]): + """Async iterator that yields :class:`ReadBatch` values and tracks caught-up state. + + Use it as an async context manager, or call :meth:`close` explicitly to close the session. + + Caution: + Returned by :meth:`S2Stream.read_session`. Do not instantiate directly. + """ + + __slots__ = ( + "_caught_up_tail", + "_closed", + "_delivery_queue", + "_error", + "_ignore_command_records", + "_task", + "_events", + "_caught_up_futs", + ) + + def __init__( + self, + events: AsyncGenerator[_ReadSessionEvent, None], + *, + ignore_command_records: bool = False, + ) -> None: + self._events = events + self._ignore_command_records = ignore_command_records + self._delivery_queue: asyncio.Queue[_ReadDelivery | BaseException | None] = ( + asyncio.Queue(maxsize=1) + ) + self._task: asyncio.Task[None] | None = None + self._closed = False + self._error: BaseException | None = None + self._caught_up_tail: StreamPosition | None = None + self._caught_up_futs: set[asyncio.Future[StreamPosition]] = set() + + def is_caught_up(self) -> bool: + """Whether this session has yielded all records with sequence numbers less than + the sequence number of the last observed tail position. + + The stream's tail may have advanced since this session last observed it. + Use :meth:`S2Stream.check_tail` if you need the current tail. + """ + return self._caught_up_tail is not None + + def caught_up(self) -> Awaitable[StreamPosition]: + """Get an awaitable that resolves to a tail position once this session is caught up to it. + + See :meth:`is_caught_up` for the semantics of caught up. + + This awaitable only signals that the session is caught up. To yield batches, continue + iterating over the session. + """ + loop = asyncio.get_running_loop() + caught_up_fut: asyncio.Future[StreamPosition] = loop.create_future() + if self._caught_up_tail is not None: + caught_up_fut.set_result(_copy_position(self._caught_up_tail)) + elif self._closed: + caught_up_fut.set_exception( + self._error if self._error is not None else _session_closed_error() + ) + else: + self._caught_up_futs.add(caught_up_fut) + caught_up_fut.add_done_callback(self._caught_up_futs.discard) + self._ensure_started() + return caught_up_fut + + async def close(self) -> None: + """Close the session.""" + if self._closed: + return + if self._task is None: + self._finish() + return + current = asyncio.current_task() + cancellation_count = current.cancelling() if current is not None else 0 + self._task.cancel() + with suppress(asyncio.CancelledError): + await self._task + if not self._closed: + self._finish() + if current is not None and current.cancelling() > cancellation_count: + raise asyncio.CancelledError + + def __aiter__(self) -> Self: + self._ensure_started() + return self + + async def __anext__(self) -> ReadBatch: + session = self.__aiter__() + try: + delivery = await session._delivery_queue.get() + except asyncio.CancelledError: + await self.close() + raise + + if delivery is None: + self._delivery_queue.put_nowait(None) + raise StopAsyncIteration + if isinstance(delivery, BaseException): + self._delivery_queue.put_nowait(None) + raise delivery + + if delivery.caught_up_tail is not None: + self._mark_caught_up(delivery.caught_up_tail) + delivery.delivered.set() + return delivery.batch + + async def __aenter__(self) -> Self: + return self.__aiter__() + + async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool: + await self.close() + return False + + def _ensure_started(self) -> None: + if self._task is None and not self._closed: + self._task = asyncio.get_running_loop().create_task(self._run()) + + async def _run(self) -> None: + error: BaseException | None = None + cancelled = False + events = self._events + try: + async for event in events: + await self._handle_event(event) + except asyncio.CancelledError: + cancelled = True + except BaseException as exc: + error = normalize_exception(exc) + finally: + try: + await events.aclose() + except BaseException as exc: + if error is None and not cancelled: + error = normalize_exception(exc) + self._finish(error) + + async def _handle_event(self, event: _ReadSessionEvent) -> None: + if isinstance(event, _ReadSessionRetrying): + self._mark_not_caught_up() + elif isinstance(event, _ReadSessionHeartbeat): + self._mark_caught_up(event.tail) + else: + batch = event.batch + caught_up_tail = ( + batch.tail + if batch.tail is not None + and batch.records[-1].seq_num + 1 == batch.tail.seq_num + else None + ) + + if self._ignore_command_records: + batch = ReadBatch( + records=[r for r in batch.records if not r.is_command_record()], + tail=batch.tail, + ) + + if batch.records: + self._mark_not_caught_up() + delivery = _ReadDelivery(batch, caught_up_tail) + await self._delivery_queue.put(delivery) + await delivery.delivered.wait() + elif caught_up_tail is not None: + self._mark_caught_up(caught_up_tail) + else: + self._mark_not_caught_up() + + def _mark_caught_up(self, tail: StreamPosition) -> None: + if self._closed: + return + self._caught_up_tail = _copy_position(tail) + caught_up_futs = tuple(self._caught_up_futs) + self._caught_up_futs.clear() + for caught_up_fut in caught_up_futs: + if not caught_up_fut.done(): + caught_up_fut.set_result(_copy_position(tail)) + + def _mark_not_caught_up(self) -> None: + if not self._closed: + self._caught_up_tail = None + + def _finish(self, error: BaseException | None = None) -> None: + if self._closed: + return + self._closed = True + self._error = error + if self._error is not None: + self._caught_up_tail = None + caught_up_futs = tuple(self._caught_up_futs) + self._caught_up_futs.clear() + for caught_up_fut in caught_up_futs: + if not caught_up_fut.done(): + caught_up_fut.set_exception( + self._error if self._error is not None else _session_closed_error() + ) + + while not self._delivery_queue.empty(): + self._delivery_queue.get_nowait() + self._delivery_queue.put_nowait(self._error) + + +def _copy_position(position: StreamPosition) -> StreamPosition: + return StreamPosition(position.seq_num, position.timestamp) + + +def _session_closed_error() -> S2ClientError: + return S2ClientError("ReadSession is closed") diff --git a/src/s2_sdk/_s2s/_read_session.py b/src/s2_sdk/_s2s/_read_session.py index 3e01b76..59c0d1a 100644 --- a/src/s2_sdk/_s2s/_read_session.py +++ b/src/s2_sdk/_s2s/_read_session.py @@ -2,18 +2,24 @@ import logging import math import time -from typing import Any, AsyncIterable +from collections.abc import AsyncGenerator +from typing import Any import s2_sdk._generated.s2.v1.s2_pb2 as pb from s2_sdk._client import HttpClient -from s2_sdk._exceptions import ReadTimeoutError +from s2_sdk._exceptions import ReadTimeoutError, S2ClientError from s2_sdk._mappers import read_batch_from_proto, read_limit_params, read_start_params +from s2_sdk._read_session import ( + _ReadSessionBatch, + _ReadSessionEvent, + _ReadSessionHeartbeat, + _ReadSessionRetrying, +) from s2_sdk._retrier import Attempt, compute_backoff, http_retry_on from s2_sdk._s2s import _stream_records_path from s2_sdk._s2s._protocol import parse_error_info, read_messages from s2_sdk._types import ( _S2_ENCRYPTION_KEY_HEADER, - ReadBatch, ReadLimit, Retry, SeqNum, @@ -35,10 +41,9 @@ async def run_read_session( until_timestamp: int | None, clamp_to_tail: bool, wait: int | None, - ignore_command_records: bool, retry: Retry, encryption_key: str | None = None, -) -> AsyncIterable[ReadBatch]: +) -> AsyncGenerator[_ReadSessionEvent, None]: params = _build_read_params(start, limit, until_timestamp, clamp_to_tail, wait) max_retries = retry._max_retries() min_base_delay = retry.min_base_delay.total_seconds() @@ -90,37 +95,36 @@ async def run_read_session( if batch.tail is not None: last_tail_at = time.monotonic() - if not batch.records: - continue - - last_record = batch.records[-1] - params["seq_num"] = last_record.seq_num + 1 - params.pop("timestamp", None) - params.pop("tail_offset", None) - - if remaining_count is not None: - remaining_count = max(remaining_count - len(batch.records), 0) - params["count"] = remaining_count - if remaining_bytes is not None: - remaining_bytes = max( - remaining_bytes - metered_bytes(batch.records), 0 - ) - params["bytes"] = remaining_bytes - - if ignore_command_records: - batch = ReadBatch( - records=[ - r for r in batch.records if not r.is_command_record() - ], - tail=batch.tail, - ) + if batch.records: + last_record = batch.records[-1] + params["seq_num"] = last_record.seq_num + 1 + params.pop("timestamp", None) + params.pop("tail_offset", None) + + if remaining_count is not None: + remaining_count = max( + remaining_count - len(batch.records), 0 + ) + params["count"] = remaining_count + if remaining_bytes is not None: + remaining_bytes = max( + remaining_bytes - metered_bytes(batch.records), 0 + ) + params["bytes"] = remaining_bytes if batch.records: - yield batch + yield _ReadSessionBatch(batch) + elif batch.tail is not None: + yield _ReadSessionHeartbeat(batch.tail) + else: + raise S2ClientError( + "Read session received an empty batch without a tail" + ) return except Exception as e: if attempt.value < max_retries and http_retry_on(e): + yield _ReadSessionRetrying() backoff = compute_backoff( attempt.value, min_base_delay=min_base_delay, diff --git a/tests/test_correctness.py b/tests/test_correctness.py index 5293ca2..8f1f5d1 100644 --- a/tests/test_correctness.py +++ b/tests/test_correctness.py @@ -3,7 +3,16 @@ import pytest -from s2_sdk import Batching, Record, Retry, S2Stream, SeqNum, SequencedRecord +from s2_sdk import ( + AppendInput, + Batching, + Record, + Retry, + S2Stream, + SeqNum, + SequencedRecord, + TailOffset, +) _NUM_RECORDS = 1024 _RECORD_IDX_HEADER = b"record-idx" @@ -29,27 +38,28 @@ async def read_records() -> None: last_seq_num: int | None = None num_records_read = 0 - async for batch in stream.read_session(start=SeqNum(0), wait=60): - for record in batch.records: - seq_num = record.seq_num - if last_seq_num is None: - assert seq_num == 0 - else: - assert seq_num == last_seq_num + 1 - last_seq_num = seq_num - - record_idx = _record_idx(record) - assert 0 <= record_idx < _NUM_RECORDS - assert record_idx <= next_record_idx - - if record_idx == next_record_idx: - next_record_idx = record_idx + 1 - num_records_read += 1 - - if next_record_idx == _NUM_RECORDS: - assert last_seq_num + 1 == num_records_read - assert num_records_read >= _NUM_RECORDS - return + async with stream.read_session(start=SeqNum(0), wait=60) as session: + async for batch in session: + for record in batch.records: + seq_num = record.seq_num + if last_seq_num is None: + assert seq_num == 0 + else: + assert seq_num == last_seq_num + 1 + last_seq_num = seq_num + + record_idx = _record_idx(record) + assert 0 <= record_idx < _NUM_RECORDS + assert record_idx <= next_record_idx + + if record_idx == next_record_idx: + next_record_idx = record_idx + 1 + num_records_read += 1 + + if next_record_idx == _NUM_RECORDS: + assert last_seq_num + 1 == num_records_read + assert num_records_read >= _NUM_RECORDS + return pytest.fail( "read session ended before all records were read: " @@ -72,6 +82,28 @@ async def append_records() -> None: task_group.create_task(append_records()) +@pytest.mark.correctness +@pytest.mark.asyncio +async def test_read_session_reports_caught_up_after_delivering_tail(stream: S2Stream): + await stream.append( + AppendInput(records=[Record(body=b"first"), Record(body=b"second")]) + ) + + async with stream.read_session(start=TailOffset(2), wait=60) as session: + caught_up = session.caught_up() + records = [] + + assert not session.is_caught_up() + while not session.is_caught_up(): + records.extend((await anext(session)).records) + + tail = await caught_up + assert [record.seq_num for record in records] == [ + tail.seq_num - 2, + tail.seq_num - 1, + ] + + def _indexed_record(idx: int) -> Record: return Record( body=b"", diff --git a/tests/test_read_session.py b/tests/test_read_session.py new file mode 100644 index 0000000..1ea8754 --- /dev/null +++ b/tests/test_read_session.py @@ -0,0 +1,204 @@ +import asyncio +from collections.abc import AsyncGenerator + +import pytest + +from s2_sdk import S2ClientError +from s2_sdk._read_session import ( + ReadSession, + _ReadSessionBatch, + _ReadSessionEvent, + _ReadSessionHeartbeat, + _ReadSessionRetrying, +) +from s2_sdk._types import ReadBatch, SequencedRecord, StreamPosition + + +def _record(seq_num: int, *, command: bool = False) -> SequencedRecord: + return SequencedRecord( + seq_num=seq_num, + body=b"command" if command else b"record", + headers=[(b"", b"fence")] if command else [], + timestamp=1, + ) + + +async def _events( + *events: _ReadSessionEvent, +) -> AsyncGenerator[_ReadSessionEvent, None]: + for event in events: + yield event + + +async def test_batch_with_tail_marks_caught_up(): + tail = StreamPosition(2, 1) + batch = ReadBatch([_record(0), _record(1)], tail) + session = ReadSession(_events(_ReadSessionBatch(batch))) + caught_up = session.caught_up() + + await asyncio.sleep(0) + assert not session.is_caught_up() + + assert await anext(session) == batch + assert session.is_caught_up() + assert await caught_up == tail + assert await session.caught_up() == tail + await session.close() + + +async def test_heartbeat_marks_caught_up_after_batch_without_tail(): + tail = StreamPosition(2, 1) + batch = ReadBatch([_record(0), _record(1)]) + session = ReadSession( + _events( + _ReadSessionBatch(batch), + _ReadSessionHeartbeat(tail), + ) + ) + caught_up = session.caught_up() + + await asyncio.sleep(0) + assert not session.is_caught_up() + + assert await anext(session) == batch + assert not session.is_caught_up() + assert await caught_up == tail + assert session.is_caught_up() + with pytest.raises(StopAsyncIteration): + await anext(session) + + +async def test_filtered_command_record_still_counts_toward_caught_up(): + tail = StreamPosition(2, 1) + batch = ReadBatch([_record(0), _record(1, command=True)], tail) + session = ReadSession( + _events(_ReadSessionBatch(batch)), ignore_command_records=True + ) + caught_up = session.caught_up() + + batch = await anext(session) + assert [record.seq_num for record in batch.records] == [0] + assert session.is_caught_up() + assert await caught_up == tail + await session.close() + + +async def test_caught_up_resolves_after_session_retry(): + allow_heartbeat = asyncio.Event() + tail = StreamPosition(3, 1) + + async def retry_then_heartbeat() -> AsyncGenerator[_ReadSessionEvent, None]: + yield _ReadSessionRetrying() + await allow_heartbeat.wait() + yield _ReadSessionHeartbeat(tail) + + session = ReadSession(retry_then_heartbeat()) + caught_up = session.caught_up() + + await asyncio.sleep(0) + assert not session.is_caught_up() + + allow_heartbeat.set() + + assert await caught_up == tail + assert session.is_caught_up() + await session.close() + + +async def test_caught_up_fails_when_session_ends(): + session = ReadSession(_events()) + + with pytest.raises(S2ClientError, match="ReadSession is closed"): + await session.caught_up() + with pytest.raises(StopAsyncIteration): + await anext(session) + with pytest.raises(StopAsyncIteration): + await anext(session) + + +async def test_context_exit_stops_read_events(): + stopped = asyncio.Event() + + async def read_events_until_stopped() -> AsyncGenerator[_ReadSessionEvent, None]: + try: + yield _ReadSessionBatch(ReadBatch([_record(0)])) + await asyncio.Event().wait() + finally: + stopped.set() + + session = ReadSession(read_events_until_stopped()) + async with session: + await asyncio.sleep(0) + + assert stopped.is_set() + + +async def test_close_before_read_starts(): + session = ReadSession(_events()) + + await session.close() + + with pytest.raises(StopAsyncIteration): + await anext(session) + + +async def test_cancelled_caught_up_await_does_not_affect_other_awaits(): + allow_heartbeat = asyncio.Event() + tail = StreamPosition(1, 1) + + async def delayed_heartbeat() -> AsyncGenerator[_ReadSessionEvent, None]: + await allow_heartbeat.wait() + yield _ReadSessionHeartbeat(tail) + + session = ReadSession(delayed_heartbeat()) + caught_up_1 = asyncio.ensure_future(session.caught_up()) + caught_up_2 = asyncio.ensure_future(session.caught_up()) + await asyncio.sleep(0) + caught_up_1.cancel() + + with pytest.raises(asyncio.CancelledError): + await caught_up_1 + + allow_heartbeat.set() + assert await caught_up_2 == tail + await session.close() + + +async def test_close_does_not_swallow_caller_cancellation(): + read_cleanup_started = asyncio.Event() + + async def blocked_read_events() -> AsyncGenerator[_ReadSessionEvent, None]: + try: + await asyncio.Event().wait() + finally: + read_cleanup_started.set() + await asyncio.Event().wait() + yield _ReadSessionRetrying() + + session = ReadSession(blocked_read_events()) + aiter(session) + await asyncio.sleep(0) + + close_task = asyncio.create_task(session.close()) + await read_cleanup_started.wait() + close_task.cancel() + with pytest.raises(asyncio.CancelledError): + await close_task + with pytest.raises(StopAsyncIteration): + await anext(session) + + +async def test_read_error_fails_caught_up_and_iteration(): + async def failing_read_events() -> AsyncGenerator[_ReadSessionEvent, None]: + raise ValueError("read failed") + yield _ReadSessionRetrying() + + session = ReadSession(failing_read_events()) + caught_up = session.caught_up() + + with pytest.raises(S2ClientError, match="read failed") as caught_up_error: + await caught_up + with pytest.raises(S2ClientError, match="read failed") as read_error: + await anext(session) + + assert read_error.value is caught_up_error.value