From bc05aef50c51f7fbd61da8184afbb8cd3db97f11 Mon Sep 17 00:00:00 2001 From: quettabit <27509167+quettabit@users.noreply.github.com> Date: Sun, 16 Aug 2026 16:20:40 -0700 Subject: [PATCH] initial commit --- src/s2_sdk/_producer.py | 73 ++++++++++++++++++++++++++++------------ tests/test_stream_ops.py | 32 ++++++++++++++++++ 2 files changed, 83 insertions(+), 22 deletions(-) diff --git a/src/s2_sdk/_producer.py b/src/s2_sdk/_producer.py index 4013561..9bb9732 100644 --- a/src/s2_sdk/_producer.py +++ b/src/s2_sdk/_producer.py @@ -31,6 +31,7 @@ class _UnackedBatch: ticket: BatchSubmitTicket indexed_ack_futs: tuple[asyncio.Future[IndexedAppendAck], ...] + indexed_acks_resolved_fut: asyncio.Future[None] class Producer: @@ -48,17 +49,18 @@ class Producer: __slots__ = ( "_accumulator", "_indexed_ack_futs", - "_batch_ready", - "_closed", + "_last_batch_indexed_acks_resolved_fut", "_drain_task", - "_error", - "_final_flush_done", + "_drain_task_wakeup", + "_batch_submit_lock", + "_final_batch_submit_done", "_fencing_token", - "_flush_lock", "_linger_task", "_match_seq_num", "_unacked", "_session", + "_closed", + "_error", ) def __init__( @@ -87,13 +89,14 @@ def __init__( self._accumulator = BatchAccumulator(batching) self._indexed_ack_futs: list[asyncio.Future[IndexedAppendAck]] = [] - self._flush_lock = asyncio.Lock() + self._batch_submit_lock = asyncio.Lock() + self._last_batch_indexed_acks_resolved_fut: asyncio.Future[None] | None = None self._linger_task: asyncio.Task[None] | None = None self._unacked: deque[_UnackedBatch] = deque() - self._batch_ready = asyncio.Event() + self._drain_task_wakeup = asyncio.Event() self._drain_task = asyncio.get_running_loop().create_task(self._drain_acks()) self._closed = False - self._final_flush_done = False + self._final_batch_submit_done = False self._error: BaseException | None = None @fallible @@ -114,14 +117,27 @@ async def submit(self, record: Record) -> RecordSubmitTicket: first_in_batch = self._accumulator.is_empty() self._accumulator.add(record) if self._accumulator.is_full(): - await self._flush() + await self._submit_batch_now() elif first_in_batch and self._accumulator.linger > 0: - linger_task = loop.create_task(self._flush_after_linger()) + linger_task = loop.create_task(self._submit_batch_after_linger()) linger_task.add_done_callback(retrieve_task_exception_if_present) self._linger_task = linger_task return RecordSubmitTicket(ack_fut) + @fallible + async def flush(self) -> None: + """Wait for all submitted records to be appended.""" + if self._closed: + raise S2ClientError("Producer is closed") + if self._error is not None: + raise self._error + + await self._submit_batch_now() + indexed_acks_resolved_fut = self._last_batch_indexed_acks_resolved_fut + if indexed_acks_resolved_fut is not None: + await asyncio.shield(indexed_acks_resolved_fut) + @fallible async def close(self) -> None: """Close the producer and wait for all submitted records to be appended.""" @@ -129,11 +145,11 @@ async def close(self) -> None: return self._closed = True try: - await self._flush() + await self._submit_batch_now() await self._session.close() finally: - self._final_flush_done = True - self._batch_ready.set() + self._final_batch_submit_done = True + self._drain_task_wakeup.set() await self._drain_task if self._error is not None: raise self._error @@ -145,12 +161,12 @@ async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool: await self.close() return False - async def _flush(self) -> None: + async def _submit_batch_now(self) -> None: await self._cancel_linger_task() await self._submit_accumulated_records() async def _submit_accumulated_records(self) -> None: - async with self._flush_lock: + async with self._batch_submit_lock: if self._accumulator.is_empty(): return @@ -175,10 +191,18 @@ async def _submit_accumulated_records(self) -> None: set_and_retrieve_future_exception(ack_fut, e) raise e + indexed_acks_resolved_fut: asyncio.Future[None] = ( + asyncio.get_running_loop().create_future() + ) self._unacked.append( - _UnackedBatch(ticket=ticket, indexed_ack_futs=indexed_ack_futs) + _UnackedBatch( + ticket=ticket, + indexed_ack_futs=indexed_ack_futs, + indexed_acks_resolved_fut=indexed_acks_resolved_fut, + ) ) - self._batch_ready.set() + self._last_batch_indexed_acks_resolved_fut = indexed_acks_resolved_fut + self._drain_task_wakeup.set() async def _cancel_linger_task(self) -> None: linger_task = self._linger_task @@ -192,15 +216,14 @@ async def _cancel_linger_task(self) -> None: await linger_task async def _drain_acks(self) -> None: - """Single background task that resolves batches in FIFO order.""" while True: while not self._unacked: - if self._closed and self._final_flush_done: + if self._closed and self._final_batch_submit_done: return - self._batch_ready.clear() + self._drain_task_wakeup.clear() if self._unacked: break - await self._batch_ready.wait() + await self._drain_task_wakeup.wait() unacked = self._unacked.popleft() try: @@ -213,23 +236,29 @@ async def _drain_acks(self) -> None: batch=ack, ) ) + if not unacked.indexed_acks_resolved_fut.done(): + unacked.indexed_acks_resolved_fut.set_result(None) except BaseException as e: e = normalize_exception(e) self._error = e for ack_fut in unacked.indexed_ack_futs: set_and_retrieve_future_exception(ack_fut, e) + set_and_retrieve_future_exception(unacked.indexed_acks_resolved_fut, e) unacked_batches = tuple(self._unacked) self._unacked.clear() for batch in unacked_batches: for ack_fut in batch.indexed_ack_futs: set_and_retrieve_future_exception(ack_fut, e) + set_and_retrieve_future_exception( + batch.indexed_acks_resolved_fut, e + ) await asyncio.gather( *(batch.ticket for batch in unacked_batches), return_exceptions=True, ) return - async def _flush_after_linger(self) -> None: + async def _submit_batch_after_linger(self) -> None: assert self._accumulator.linger is not None await asyncio.sleep(self._accumulator.linger) self._linger_task = None diff --git a/tests/test_stream_ops.py b/tests/test_stream_ops.py index 2e34e94..48c69e4 100644 --- a/tests/test_stream_ops.py +++ b/tests/test_stream_ops.py @@ -687,6 +687,38 @@ async def test_producer_close_delivers_all_indexed_acks_from_different_acks( assert ack0.batch is not ack1.batch assert ack1.batch is not ack2.batch + async def test_producer_flush_submits_partial_batch_and_resolves_prior_record_tickets( + self, stream: S2Stream + ): + async with stream.producer( + batching=Batching(max_records=2, linger=timedelta(0)) + ) as p: + t0 = await p.submit(Record(body=b"lorem")) + t1 = await p.submit(Record(body=b"ipsum")) + t2 = await p.submit(Record(body=b"dolor")) + + await p.flush() + + assert t0._ack_fut.done() + assert t1._ack_fut.done() + assert t2._ack_fut.done() + + ack0 = await t0 + ack1 = await t1 + ack2 = await t2 + + assert ack0.seq_num == 0 + assert ack1.seq_num == 1 + assert ack2.seq_num == 2 + assert ack0.batch is ack1.batch + assert ack1.batch is not ack2.batch + + t3 = await p.submit(Record(body=b"sit")) + + ack3 = await t3 + assert ack3.seq_num == 3 + assert ack2.batch is not ack3.batch + async def test_producer_nonexistent_stream_errors(self, shared_basin: S2Basin): nonexistent = shared_basin.stream("nonexistent-stream-xyz")