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..dd528e4 100644 --- a/src/s2_sdk/_producer.py +++ b/src/s2_sdk/_producer.py @@ -24,7 +24,12 @@ @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: + with suppress(asyncio.CancelledError): + task.exception() class Producer: @@ -41,16 +46,17 @@ class Producer: __slots__ = ( "_accumulator", - "_indexed_ack_futs", + "_buffered_ack_futs", "_batch_ready", "_closed", "_drain_task", "_error", "_final_flush_done", "_fencing_token", - "_flush_lock", "_linger_task", "_match_seq_num", + "_operation_lock", + "_outstanding_ack_futs", "_unacked", "_session", ) @@ -80,8 +86,9 @@ def __init__( self._match_seq_num = match_seq_num self._accumulator = BatchAccumulator(batching) - self._indexed_ack_futs: list[asyncio.Future[IndexedAppendAck]] = [] - self._flush_lock = asyncio.Lock() + self._buffered_ack_futs: list[asyncio.Future[IndexedAppendAck]] = [] + self._operation_lock = asyncio.Lock() + 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() @@ -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._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) - 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() + prior_ack_futs = tuple(self._outstanding_ack_futs) + await self._flush_current_batch() + + if self._error is not None: + raise self._error + if not prior_ack_futs: + return + + results = await asyncio.shield( + asyncio.gather(*prior_ack_futs, 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() + record_ack_futs = tuple(self._buffered_ack_futs) + self._buffered_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 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, record_ack_futs=record_ack_futs) + ) + self._batch_ready.set() async def _cancel_linger_task(self) -> None: linger_task = self._linger_task @@ -201,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( @@ -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._buffered_ack_futs.clear() + for ack_fut in tuple(self._outstanding_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..5259b90 --- /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 + 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 + 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: + self._closed = True + + +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