From 54ac85aca0fc548faee6913529ebdeef43f5a888 Mon Sep 17 00:00:00 2001 From: Mehul Arora Date: Tue, 11 Aug 2026 13:39:21 -0400 Subject: [PATCH 1/2] feat: add durable producer flush --- examples/docs/streams.py | 3 + src/s2_sdk/_producer.py | 229 +++++++++++++++++++++++------------ tests/test_producer.py | 254 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 412 insertions(+), 74 deletions(-) create mode 100644 tests/test_producer.py diff --git a/examples/docs/streams.py b/examples/docs/streams.py index ea0f5a1..cf52c79 100644 --- a/examples/docs/streams.py +++ b/examples/docs/streams.py @@ -99,6 +99,9 @@ async def producer_example(stream): # Submit individual records ticket = await producer.submit(Record(body=b"my event")) + # Make all records submitted so far durable without closing the producer + await producer.flush() + # Get the exact sequence number for this record ack = await ticket print(f"Record durable at seq_num {ack.seq_num}") diff --git a/src/s2_sdk/_producer.py b/src/s2_sdk/_producer.py index 986b6d6..8f0b95b 100644 --- a/src/s2_sdk/_producer.py +++ b/src/s2_sdk/_producer.py @@ -27,6 +27,11 @@ class _UnackedBatch: indexed_ack_futs: tuple[asyncio.Future[IndexedAppendAck], ...] +def _retrieve_task_exception(task: asyncio.Task[None]) -> None: + with suppress(asyncio.CancelledError): + task.exception() + + class Producer: """High-level interface for submitting individual records. @@ -48,9 +53,10 @@ class Producer: "_error", "_final_flush_done", "_fencing_token", - "_flush_lock", "_linger_task", "_match_seq_num", + "_operation_lock", + "_pending_ack_futs", "_unacked", "_session", ) @@ -81,7 +87,8 @@ def __init__( self._accumulator = BatchAccumulator(batching) self._indexed_ack_futs: list[asyncio.Future[IndexedAppendAck]] = [] - self._flush_lock = asyncio.Lock() + self._operation_lock = asyncio.Lock() + self._pending_ack_futs: set[asyncio.Future[IndexedAppendAck]] = set() self._linger_task: asyncio.Task[None] | None = None self._unacked: deque[_UnackedBatch] = deque() self._batch_ready = asyncio.Event() @@ -94,41 +101,94 @@ def __init__( async def submit(self, record: Record) -> RecordSubmitTicket: """Submit a record for appending. - Waits when backpressure limits are reached. + Waits when backpressure limits are reached. Await the returned ticket + for this record's acknowledgement, or call :meth:`flush` to wait for + all previously submitted records. """ - if self._closed: - raise S2ClientError("Producer is closed") - if self._error is not None: - raise self._error + async with self._operation_lock: + self._check_ready() - loop = asyncio.get_running_loop() - ack_fut: asyncio.Future[IndexedAppendAck] = loop.create_future() - self._indexed_ack_futs.append(ack_fut) + loop = asyncio.get_running_loop() + ack_fut: asyncio.Future[IndexedAppendAck] = loop.create_future() + self._indexed_ack_futs.append(ack_fut) + self._pending_ack_futs.add(ack_fut) + ack_fut.add_done_callback(self._pending_ack_futs.discard) - first_in_batch = self._accumulator.is_empty() - self._accumulator.add(record) - if self._accumulator.is_full(): - await self._flush() - elif first_in_batch and self._accumulator.linger > 0: - self._linger_task = loop.create_task(self._flush_after_linger()) + first_in_batch = self._accumulator.is_empty() + self._accumulator.add(record) + if self._accumulator.is_full(): + await self._flush_current_batch() + elif first_in_batch and self._accumulator.linger > 0: + self._linger_task = loop.create_task(self._flush_after_linger()) - return RecordSubmitTicket(ack_fut) + return RecordSubmitTicket(ack_fut) @fallible - async def close(self) -> None: - """Close the producer and wait for all submitted records to be appended.""" - if self._closed: - return - self._closed = True + async def flush(self) -> None: + """Flush pending records and wait for prior submissions to become durable. + + Records submitted before this operation are included. Concurrent + submissions may fall on either side of the flush boundary. + + An empty flush returns immediately. A successful flush leaves the + producer open for further submissions. An append failure is terminal, + and subsequent calls to :meth:`submit`, :meth:`flush`, and + :meth:`close` raise the same error. Canceling the caller's wait does not + cancel the flush operation or the submitted records. + """ + flush_task = asyncio.create_task(self._flush_and_wait()) try: - await self._flush() - await self._session.close() - finally: - self._final_flush_done = True - self._batch_ready.set() - await self._drain_task - if self._error is not None: - raise self._error + await asyncio.shield(flush_task) + except asyncio.CancelledError: + flush_task.add_done_callback(_retrieve_task_exception) + raise + + async def _flush_and_wait(self) -> None: + async with self._operation_lock: + self._check_ready() + ack_futs_to_wait = tuple(self._pending_ack_futs) + await self._flush_current_batch() + + if self._error is not None: + raise self._error + if not ack_futs_to_wait: + return + + results = await asyncio.shield( + asyncio.gather(*ack_futs_to_wait, return_exceptions=True) + ) + for result in results: + if isinstance(result, BaseException): + raise result + + @fallible + async def close(self) -> None: + """Close the producer and wait for all submitted records to become durable.""" + async with self._operation_lock: + if self._closed: + if self._error is not None: + raise self._error + return + self._closed = True + try: + if self._error is None: + try: + await self._flush_current_batch() + except BaseException as e: + self._fail(e) + try: + await self._session.close() + except BaseException as e: + self._fail(e) + finally: + self._final_flush_done = True + self._batch_ready.set() + try: + await self._drain_task + except BaseException as e: + self._fail(e) + if self._error is not None: + raise self._error async def __aenter__(self) -> Self: return self @@ -137,44 +197,46 @@ async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool: await self.close() return False - async def _flush(self) -> None: + def _check_ready(self) -> None: + if self._error is not None: + raise self._error + if self._closed: + raise S2ClientError("Producer is closed") + + async def _flush_current_batch(self) -> None: await self._cancel_linger_task() await self._submit_accumulated_records() async def _submit_accumulated_records(self) -> None: - async with self._flush_lock: - if self._accumulator.is_empty(): - return + if self._accumulator.is_empty(): + return - records = self._accumulator.take() - indexed_ack_futs = tuple(self._indexed_ack_futs) - self._indexed_ack_futs.clear() + records = self._accumulator.take() + indexed_ack_futs = tuple(self._indexed_ack_futs) + self._indexed_ack_futs.clear() - batch = AppendInput( - records=records, - fencing_token=self._fencing_token, - match_seq_num=self._match_seq_num, - ) - if self._match_seq_num is not None: - self._match_seq_num += len(records) + batch = AppendInput( + records=records, + fencing_token=self._fencing_token, + match_seq_num=self._match_seq_num, + ) + if self._match_seq_num is not None: + self._match_seq_num += len(records) - try: - ticket = await self._session.submit(batch) - except BaseException as e: - e = normalize_exception(e) - self._error = e - for ack_fut in indexed_ack_futs: - if not ack_fut.done(): - ack_fut.set_exception(e) - # Suppress "Future exception was never retrieved" for - # futures the caller never got back (submit raised). - ack_fut.exception() - raise e - - self._unacked.append( - _UnackedBatch(ticket=ticket, indexed_ack_futs=indexed_ack_futs) - ) - self._batch_ready.set() + try: + ticket = await self._session.submit(batch) + except BaseException as e: + error = self._fail(e) + for ack_fut in indexed_ack_futs: + # Suppress "Future exception was never retrieved" for the + # record whose submit call raised before returning its ticket. + ack_fut.exception() + raise error + + self._unacked.append( + _UnackedBatch(ticket=ticket, indexed_ack_futs=indexed_ack_futs) + ) + self._batch_ready.set() async def _cancel_linger_task(self) -> None: linger_task = self._linger_task @@ -210,28 +272,47 @@ async def _drain_acks(self) -> None: ) ) except BaseException as e: - e = normalize_exception(e) - self._error = e - for ack_fut in unacked.indexed_ack_futs: - if not ack_fut.done(): - ack_fut.set_exception(e) - # Fail all remaining unacked batches too - for remaining in self._unacked: - for ack_fut in remaining.indexed_ack_futs: - if not ack_fut.done(): - ack_fut.set_exception(e) + self._fail(e) self._unacked.clear() return async def _flush_after_linger(self) -> None: assert self._accumulator.linger is not None await asyncio.sleep(self._accumulator.linger) + async with self._operation_lock: + if self._linger_task is not asyncio.current_task(): + return + self._linger_task = None + if self._closed or self._error is not None: + return + try: + await self._submit_accumulated_records() + except asyncio.CancelledError: + raise + except BaseException as e: + self._fail(e) + + def _fail(self, cause: BaseException) -> BaseException: + if self._error is None: + self._error = normalize_exception(cause) + error = self._error + + linger_task = self._linger_task self._linger_task = None - await self._submit_accumulated_records() + if linger_task is not None and linger_task is not asyncio.current_task(): + linger_task.cancel() + + self._accumulator.take() + self._indexed_ack_futs.clear() + for ack_fut in tuple(self._pending_ack_futs): + if not ack_fut.done(): + ack_fut.set_exception(error) + + return error class RecordSubmitTicket: - """Awaitable that resolves to an :class:`IndexedAppendAck` once the record is appended.""" + """Awaitable that resolves to an :class:`IndexedAppendAck` once the record is durable.""" __slots__ = ("_ack_fut",) @@ -239,4 +320,4 @@ def __init__(self, ack_fut: asyncio.Future[IndexedAppendAck]) -> None: self._ack_fut = ack_fut def __await__(self): - return self._ack_fut.__await__() + return asyncio.shield(self._ack_fut).__await__() diff --git a/tests/test_producer.py b/tests/test_producer.py new file mode 100644 index 0000000..eb36592 --- /dev/null +++ b/tests/test_producer.py @@ -0,0 +1,254 @@ +import asyncio +from collections.abc import Awaitable +from dataclasses import dataclass +from datetime import timedelta +from typing import TypeVar +from unittest.mock import MagicMock, patch + +import pytest + +from s2_sdk import ( + AppendAck, + AppendInput, + Batching, + BatchSubmitTicket, + Compression, + Producer, + Record, + RecordSubmitTicket, + Retry, + S2ClientError, + S2Error, + SeqNumMismatchError, + StreamPosition, +) + +_TEST_TIMEOUT = 1 +_T = TypeVar("_T") + + +@dataclass(slots=True) +class _PendingAppend: + append_input: AppendInput + ack_fut: asyncio.Future[AppendAck] + + def resolve(self, ack: AppendAck) -> None: + self.ack_fut.set_result(ack) + + def reject(self, error: BaseException) -> None: + self.ack_fut.set_exception(error) + + +class _TestAppendSession: + def __init__( + self, + *, + fail_on_submit: int | None = None, + submit_error: BaseException | None = None, + ) -> None: + self.pending_appends: asyncio.Queue[_PendingAppend] = asyncio.Queue() + self._fail_on_submit = fail_on_submit + self._submit_error = submit_error + self._submit_count = 0 + + async def submit(self, append_input: AppendInput) -> BatchSubmitTicket: + self._submit_count += 1 + if self._submit_count == self._fail_on_submit: + assert self._submit_error is not None + raise self._submit_error + + ack_fut: asyncio.Future[AppendAck] = asyncio.get_running_loop().create_future() + self.pending_appends.put_nowait(_PendingAppend(append_input, ack_fut)) + return BatchSubmitTicket(ack_fut) + + async def close(self) -> None: + pass + + +def _producer( + session: _TestAppendSession, + *, + batching: Batching, +) -> Producer: + with patch("s2_sdk._producer.AppendSession", return_value=session): + return Producer( + client=MagicMock(), + stream_name="test-stream", + retry=Retry(), + compression=Compression.NONE, + fencing_token=None, + match_seq_num=None, + max_unacked_bytes=5 * 1024 * 1024, + batching=batching, + ) + + +async def _next_append(session: _TestAppendSession) -> _PendingAppend: + async with asyncio.timeout(_TEST_TIMEOUT): + return await session.pending_appends.get() + + +async def _wait_for(awaitable: Awaitable[_T]) -> _T: + async with asyncio.timeout(_TEST_TIMEOUT): + return await awaitable + + +def _make_ack(start_seq_num: int, record_count: int) -> AppendAck: + end_seq_num = start_seq_num + record_count + return AppendAck( + start=StreamPosition(seq_num=start_seq_num, timestamp=1), + end=StreamPosition(seq_num=end_seq_num, timestamp=1), + tail=StreamPosition(seq_num=end_seq_num, timestamp=1), + ) + + +async def test_flush_emits_partial_batch_and_waits_for_prior_acks(): + session = _TestAppendSession() + producer = _producer( + session, + batching=Batching(max_records=2, linger=timedelta(hours=1)), + ) + + tickets = [ + await producer.submit(Record(body=body)) for body in (b"one", b"two", b"three") + ] + + cancelled_waiter = asyncio.ensure_future(tickets[2]) + await asyncio.sleep(0) + cancelled_waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled_waiter + + flush_task = asyncio.create_task(producer.flush()) + + full_append = await _next_append(session) + partial_append = await _next_append(session) + assert [ + len(full_append.append_input.records), + len(partial_append.append_input.records), + ] == [2, 1] + + full_ack = _make_ack(41, 2) + full_append.resolve(full_ack) + first_ack, second_ack = await asyncio.gather(*tickets[:2]) + assert not flush_task.done() + + partial_ack = _make_ack(43, 1) + partial_append.resolve(partial_ack) + await _wait_for(flush_task) + third_ack = await tickets[2] + + assert [first_ack.seq_num, second_ack.seq_num, third_ack.seq_num] == [41, 42, 43] + assert first_ack.batch is full_ack + assert second_ack.batch is full_ack + assert third_ack.batch is partial_ack + await producer.close() + + +async def test_flush_is_reusable_and_excludes_later_submissions(): + session = _TestAppendSession() + producer = _producer( + session, + batching=Batching(max_records=100, linger=timedelta(hours=1)), + ) + + await _wait_for(asyncio.create_task(producer.flush())) + assert session.pending_appends.empty() + + first_ticket = await producer.submit(Record(body=b"before")) + first_flush_task = asyncio.create_task(producer.flush()) + first_append = await _next_append(session) + + later_submit_task = asyncio.create_task(producer.submit(Record(body=b"after"))) + await asyncio.sleep(0) + first_flush_task.cancel() + with pytest.raises(asyncio.CancelledError): + await first_flush_task + assert not later_submit_task.done() + + first_append.resolve(_make_ack(0, 1)) + later_ticket = await _wait_for(later_submit_task) + assert session.pending_appends.empty() + + second_flush_task = asyncio.create_task(producer.flush()) + second_append = await _next_append(session) + second_append.resolve(_make_ack(1, 1)) + await _wait_for(second_flush_task) + + assert (await first_ticket).seq_num == 0 + assert (await later_ticket).seq_num == 1 + await producer.close() + + +@pytest.mark.parametrize("failure_phase", ["submit", "ack"]) +async def test_flush_propagates_terminal_failure_to_pending_tickets( + failure_phase: str, +): + error = ( + SeqNumMismatchError( + "seq_num_mismatch", + "expected sequence number 7", + 412, + expected_seq_num=7, + ) + if failure_phase == "submit" + else S2ClientError("connection closed before acknowledgement") + ) + session = _TestAppendSession( + fail_on_submit=2 if failure_phase == "submit" else None, + submit_error=error, + ) + producer = _producer( + session, + batching=Batching( + max_records=2 if failure_phase == "submit" else 1, + linger=timedelta(hours=1), + ), + ) + + tickets: list[RecordSubmitTicket] = [] + tickets.append(await producer.submit(Record(body=b"one"))) + tickets.append(await producer.submit(Record(body=b"two"))) + first_append = await _next_append(session) + close_task: asyncio.Task[None] | None = None + + if failure_phase == "submit": + tickets.append(await producer.submit(Record(body=b"three"))) + flush_task = asyncio.create_task(producer.flush()) + else: + await _next_append(session) + flush_task = asyncio.create_task(producer.flush()) + await asyncio.sleep(0) + assert not flush_task.done() + close_task = asyncio.create_task(producer.close()) + await asyncio.sleep(0) + assert not close_task.done() + first_append.reject(error) + + with pytest.raises(S2Error) as flush_exc: + await _wait_for(flush_task) + assert flush_exc.value is error + + for ticket in tickets: + with pytest.raises(S2Error) as ticket_exc: + await ticket + assert ticket_exc.value is error + + with pytest.raises(S2Error) as repeated_flush_exc: + await producer.flush() + assert repeated_flush_exc.value is error + + with pytest.raises(S2Error) as submit_exc: + await producer.submit(Record(body=b"rejected")) + assert submit_exc.value is error + + if failure_phase == "submit": + first_append.resolve(_make_ack(0, 2)) + close_task = asyncio.create_task(producer.close()) + assert close_task is not None + with pytest.raises(S2Error) as close_exc: + await _wait_for(close_task) + assert close_exc.value is error + if failure_phase == "submit": + assert isinstance(error, SeqNumMismatchError) + assert error.expected_seq_num == 7 From 428eb70cacb2915dc839a92f307db3a18552bb66 Mon Sep 17 00:00:00 2001 From: Mehul Arora Date: Tue, 11 Aug 2026 16:46:52 -0400 Subject: [PATCH 2/2] refactor: clarify producer flush state --- src/s2_sdk/_producer.py | 36 ++++++++++++++++++------------------ tests/test_producer.py | 8 ++++---- 2 files changed, 22 insertions(+), 22 deletions(-) diff --git a/src/s2_sdk/_producer.py b/src/s2_sdk/_producer.py index 8f0b95b..dd528e4 100644 --- a/src/s2_sdk/_producer.py +++ b/src/s2_sdk/_producer.py @@ -24,7 +24,7 @@ @dataclass(slots=True) class _UnackedBatch: ticket: BatchSubmitTicket - indexed_ack_futs: tuple[asyncio.Future[IndexedAppendAck], ...] + record_ack_futs: tuple[asyncio.Future[IndexedAppendAck], ...] def _retrieve_task_exception(task: asyncio.Task[None]) -> None: @@ -46,7 +46,7 @@ class Producer: __slots__ = ( "_accumulator", - "_indexed_ack_futs", + "_buffered_ack_futs", "_batch_ready", "_closed", "_drain_task", @@ -56,7 +56,7 @@ class Producer: "_linger_task", "_match_seq_num", "_operation_lock", - "_pending_ack_futs", + "_outstanding_ack_futs", "_unacked", "_session", ) @@ -86,9 +86,9 @@ def __init__( self._match_seq_num = match_seq_num self._accumulator = BatchAccumulator(batching) - self._indexed_ack_futs: list[asyncio.Future[IndexedAppendAck]] = [] + self._buffered_ack_futs: list[asyncio.Future[IndexedAppendAck]] = [] self._operation_lock = asyncio.Lock() - self._pending_ack_futs: set[asyncio.Future[IndexedAppendAck]] = set() + self._outstanding_ack_futs: set[asyncio.Future[IndexedAppendAck]] = set() self._linger_task: asyncio.Task[None] | None = None self._unacked: deque[_UnackedBatch] = deque() self._batch_ready = asyncio.Event() @@ -110,9 +110,9 @@ async def submit(self, record: Record) -> RecordSubmitTicket: loop = asyncio.get_running_loop() ack_fut: asyncio.Future[IndexedAppendAck] = loop.create_future() - self._indexed_ack_futs.append(ack_fut) - self._pending_ack_futs.add(ack_fut) - ack_fut.add_done_callback(self._pending_ack_futs.discard) + self._buffered_ack_futs.append(ack_fut) + self._outstanding_ack_futs.add(ack_fut) + ack_fut.add_done_callback(self._outstanding_ack_futs.discard) first_in_batch = self._accumulator.is_empty() self._accumulator.add(record) @@ -146,16 +146,16 @@ async def flush(self) -> None: async def _flush_and_wait(self) -> None: async with self._operation_lock: self._check_ready() - ack_futs_to_wait = tuple(self._pending_ack_futs) + prior_ack_futs = tuple(self._outstanding_ack_futs) await self._flush_current_batch() if self._error is not None: raise self._error - if not ack_futs_to_wait: + if not prior_ack_futs: return results = await asyncio.shield( - asyncio.gather(*ack_futs_to_wait, return_exceptions=True) + asyncio.gather(*prior_ack_futs, return_exceptions=True) ) for result in results: if isinstance(result, BaseException): @@ -212,8 +212,8 @@ async def _submit_accumulated_records(self) -> None: return records = self._accumulator.take() - indexed_ack_futs = tuple(self._indexed_ack_futs) - self._indexed_ack_futs.clear() + record_ack_futs = tuple(self._buffered_ack_futs) + self._buffered_ack_futs.clear() batch = AppendInput( records=records, @@ -227,14 +227,14 @@ async def _submit_accumulated_records(self) -> None: ticket = await self._session.submit(batch) except BaseException as e: error = self._fail(e) - for ack_fut in indexed_ack_futs: + for ack_fut in record_ack_futs: # Suppress "Future exception was never retrieved" for the # record whose submit call raised before returning its ticket. ack_fut.exception() raise error self._unacked.append( - _UnackedBatch(ticket=ticket, indexed_ack_futs=indexed_ack_futs) + _UnackedBatch(ticket=ticket, record_ack_futs=record_ack_futs) ) self._batch_ready.set() @@ -263,7 +263,7 @@ async def _drain_acks(self) -> None: unacked = self._unacked.popleft() try: ack: AppendAck = await unacked.ticket # type: ignore[assignment] - for i, ack_fut in enumerate(unacked.indexed_ack_futs): + for i, ack_fut in enumerate(unacked.record_ack_futs): if not ack_fut.done(): ack_fut.set_result( IndexedAppendAck( @@ -303,8 +303,8 @@ def _fail(self, cause: BaseException) -> BaseException: linger_task.cancel() self._accumulator.take() - self._indexed_ack_futs.clear() - for ack_fut in tuple(self._pending_ack_futs): + self._buffered_ack_futs.clear() + for ack_fut in tuple(self._outstanding_ack_futs): if not ack_fut.done(): ack_fut.set_exception(error) diff --git a/tests/test_producer.py b/tests/test_producer.py index eb36592..5259b90 100644 --- a/tests/test_producer.py +++ b/tests/test_producer.py @@ -50,8 +50,11 @@ def __init__( self._fail_on_submit = fail_on_submit self._submit_error = submit_error self._submit_count = 0 + self._closed = False async def submit(self, append_input: AppendInput) -> BatchSubmitTicket: + if self._closed: + raise S2ClientError("AppendSession is closed") self._submit_count += 1 if self._submit_count == self._fail_on_submit: assert self._submit_error is not None @@ -62,7 +65,7 @@ async def submit(self, append_input: AppendInput) -> BatchSubmitTicket: return BatchSubmitTicket(ack_fut) async def close(self) -> None: - pass + self._closed = True def _producer( @@ -249,6 +252,3 @@ async def test_flush_propagates_terminal_failure_to_pending_tickets( with pytest.raises(S2Error) as close_exc: await _wait_for(close_task) assert close_exc.value is error - if failure_phase == "submit": - assert isinstance(error, SeqNumMismatchError) - assert error.expected_seq_num == 7