feat!: introduce ReadSession with caught-up state tracking

docs/source/api-reference.md

@@ -22,6 +22,9 @@
 .. autoclass:: BatchSubmitTicket()
     :members:
 
+.. autoclass:: ReadSession()
+    :members:
+
 .. autoclass:: Producer()
     :members:
 

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",

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.
     """

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()

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:

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.
     """

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")