Approach: use an asyncio.Lock to protect a growing current batch (items + per-call Futures). When the first item arrives in an empty batch, schedule a flush task after 100ms. add(item) appends (item, Future) under the lock and returns await future. The flush swaps out the batch atomically, then calls process_batch with retries (up to 3) and maps results back to futures preserving order. On failure after retries, set exceptions on all futures.python
import asyncio
from typing import Any, Callable, List, Tuple
class Batcher:
def __init__(self, process_batch: Callable[[List[Any]], asyncio.Future], window_ms: int = 100):
self._process_batch = process_batch
self._window = window_ms / 1000.0
self._lock = asyncio.Lock()
self._items: List[Any] = []
self._futures: List[asyncio.Future] = []
self._flush_task: asyncio.Task | None = None
async def add(self, item: Any) -> Any:
fut = asyncio.get_running_loop().create_future()
async with self._lock:
self._items.append(item)
self._futures.append(fut)
if self._flush_task is None or self._flush_task.done():
# schedule flush after window
self._flush_task = asyncio.create_task(self._schedule_flush())
return await fut
async def _schedule_flush(self):
await asyncio.sleep(self._window)
await self._flush()
async def _flush(self):
# swap batch atomically
async with self._lock:
items, futures = self._items, self._futures
self._items, self._futures = [], []
self._flush_task = None
if not items:
return
max_retries = 3
backoff = 0.05
last_exc = None
for attempt in range(1, max_retries + 1):
try:
# call process_batch and expect an iterable of results aligned with items
results = await self._process_batch(items)
if len(results) != len(futures):
raise RuntimeError("process_batch returned wrong number of results")
for fut, res in zip(futures, results):
if not fut.done():
fut.set_result(res)
return
except Exception as e:
last_exc = e
if attempt < max_retries:
await asyncio.sleep(backoff * (2 ** (attempt - 1)))
else:
break
# failed after retries -> set exception on all pending futures
for fut in futures:
if not fut.done():
fut.set_exception(last_exc)
