Performance Trade-offs & Optimization Strategy Questions

import asyncio
import time
from typing import Any, Callable, List, Awaitable, Tuple

class BatchScheduler:
    def __init__(self, batch_size: int, max_wait_ms: int, worker_fn: Callable[[List[Any]], Awaitable[List[Any]]]):
        self.batch_size = batch_size
        self.max_wait_s = max_wait_ms / 1000.0
        self.worker_fn = worker_fn

        self._lock = asyncio.Lock()
        self._queue: List[Tuple[int, float, Any, asyncio.Future]] = []  # (priority, enqueue_time, payload, future)
        self._timer_task: asyncio.Task | None = None

    async def submit(self, payload, priority: int = 0) -> Any:
        fut = asyncio.get_running_loop().create_future()
        now = time.monotonic()
        async with self._lock:
            # append; keep insertion order for same priority
            self._queue.append((priority, now, payload, fut))
            # start timer if first
            if len(self._queue) == 1:
                self._start_timer()
            # preemption: if new request has higher priority than any existing (lower number = lower priority)
            min_priority = min(p for p,_,_,_ in self._queue)
            if priority > min_priority:
                # flush earlier to avoid high-priority waiting behind many low-priority ones
                await self._flush_locked()
            elif len(self._queue) >= self.batch_size:
                await self._flush_locked()

        return await fut

    def _start_timer(self):
        if self._timer_task and not self._timer_task.done():
            self._timer_task.cancel()
        self._timer_task = asyncio.create_task(self._timer_worker())

    async def _timer_worker(self):
        try:
            async with self._lock:
                if not self._queue:
                    return
                oldest_time = min(t for _, t, _, _ in self._queue)
                wait = max(0.0, self.max_wait_s - (time.monotonic() - oldest_time))
            if wait > 0:
                await asyncio.sleep(wait)
            async with self._lock:
                if self._queue:
                    await self._flush_locked()
        except asyncio.CancelledError:
            return

    async def _flush_locked(self):
        # must be called while holding self._lock
        if not self._queue:
            return
        batch = self._queue[:self.batch_size]
        self._queue = self._queue[self.batch_size:]
        # cancel & restart timer for remaining
        if self._timer_task and not self._timer_task.done():
            self._timer_task.cancel()
            self._timer_task = None
        if self._queue:
            self._start_timer()

        priorities, times, payloads, futures = zip(*batch)
        # release lock while calling worker_fn
        self._lock.release()
        try:
            results = await self.worker_fn(list(payloads))
            if not isinstance(results, list) or len(results) != len(futures):
                raise RuntimeError("worker_fn must return a list with same length as input")
            for fut, res in zip(futures, results):
                if not fut.done():
                    fut.set_result(res)
        except Exception as exc:
            for fut in futures:
                if not fut.done():
                    fut.set_exception(exc)
        finally:
            # reacquire lock to preserve caller expectations
            await self._lock.acquire()