fix: linger timeout closes async iterator in append_record_batches (#28)

Author: quettabit

src/s2_sdk/_batching.py

@@ -50,41 +50,47 @@ async def append_record_batches(
     validate_batching(batching.max_records, batching.max_bytes)
     acc = BatchAccumulator(batching)
     linger_secs = batching.linger.total_seconds()
-    aiter = records.__aiter__()
+    record_iter = records.__aiter__()
+    pending_next = None
 
-    while True:
-        try:
-            record = await anext(aiter)
-        except StopAsyncIteration:
-            break
+    try:
+        while True:
+            if pending_next is not None:
+                record = await pending_next
+                pending_next = None
+            else:
+                record = await anext(record_iter, None)
+            if record is None:
+                break
 
-        acc.add(record)
-        if acc.is_full():
-            yield acc.take()
-            continue
+            acc.add(record)
 
-        try:
             deadline = (
-                asyncio.get_event_loop().time() + linger_secs
+                asyncio.get_running_loop().time() + linger_secs
                 if linger_secs > 0
                 else None
             )
             while not acc.is_full():
                 if deadline is not None:
-                    remaining = deadline - asyncio.get_event_loop().time()
+                    remaining = deadline - asyncio.get_running_loop().time()
                     if remaining <= 0:
                         break
-                    record = await asyncio.wait_for(anext(aiter), timeout=remaining)
+                    next_task = asyncio.create_task(anext(record_iter, None))
+                    done, _ = await asyncio.wait({next_task}, timeout=remaining)
+                    if not done:
+                        pending_next = next_task
+                        break
+                    record = next_task.result()
                 else:
-                    record = await anext(aiter)
+                    record = await anext(record_iter, None)
+                if record is None:
+                    break
                 acc.add(record)
-        except StopAsyncIteration:
-            pass
-        except TimeoutError:
-            pass
 
-        if not acc.is_empty():
             yield acc.take()
+    finally:
+        if pending_next is not None:
+            pending_next.cancel()
 
 
 async def append_inputs(

src/s2_sdk/_types.py

@@ -208,7 +208,7 @@ class Batching:
     """Maximum metered bytes per batch. Must be between 8 and 1 MiB. Default is 1 MiB."""
 
     linger: timedelta = field(default_factory=lambda: timedelta(milliseconds=5))
-    """Maximum time to wait for more records before flushing a partial batch. Default is 5 ms.
+    """Maximum time to wait for more records before flushing a batch. Default is 5 ms.
 
     Note:
         If set to 0, batches are flushed only when ``max_records`` or ``max_bytes`` is reached.

tests/test_batching.py

@@ -1,3 +1,4 @@
+import asyncio
 from datetime import timedelta
 
 import pytest
@@ -63,6 +64,26 @@ async def test_oversized_record_passes():
     assert len(batches[0]) == 1
 
 
+@pytest.mark.asyncio
+async def test_linger_flushes_batches():
+    async def delayed_records():
+        yield Record(body=b"r1")
+        await asyncio.sleep(0.1)  # Longer than linger
+        yield Record(body=b"r2")
+        yield Record(body=b"r3")
+
+    batches = []
+    async for batch in append_record_batches(
+        delayed_records(),
+        batching=Batching(max_records=10, linger=timedelta(seconds=0.01)),
+    ):
+        batches.append(batch)
+
+    assert len(batches) == 2
+    assert len(batches[0]) == 1  # r1 (linger expired)
+    assert len(batches[1]) == 2  # r2, r3
+
+
 @pytest.mark.asyncio
 async def test_append_inputs_skips_empty_batches():
     inputs = []