feat: add flush method to Producer

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,26 +117,39 @@ 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."""
         if self._closed:
             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

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