stonks-oracle/services/specialist/batching.py

"""Dynamic batcher — collects requests up to max_batch_size or max_wait_ms.

Bounded: rejects requests when the queue exceeds max_queue_size to prevent
unbounded memory growth under sustained load.
"""

from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass, field
from typing import Any, Callable, TypeVar

logger = logging.getLogger(__name__)

T = TypeVar("T")


class QueueFullError(Exception):
    """Raised when the batcher queue is at capacity."""

    pass


@dataclass
class _PendingRequest:
    """A single request waiting in the batch queue."""

    payload: Any
    future: asyncio.Future[Any] = field(default_factory=lambda: asyncio.get_event_loop().create_future())
    enqueued_at: float = field(default_factory=time.monotonic)


class DynamicBatcher:
    """Bounded dynamic batcher that collects incoming requests and processes
    them together when either max_batch_size is reached or max_wait_ms elapses.

    Bounded: if the internal queue exceeds max_queue_size, new submissions
    are rejected with QueueFullError.

    Usage:
        batcher = DynamicBatcher(max_batch_size=32, max_wait_ms=50.0, max_queue_size=256)
        batcher.start(process_batch_fn)
        result = await batcher.submit(payload)
        await batcher.stop()
    """

    def __init__(
        self,
        max_batch_size: int = 32,
        max_wait_ms: float = 50.0,
        max_queue_size: int = 256,
    ) -> None:
        self.max_batch_size = max_batch_size
        self.max_wait_ms = max_wait_ms
        self.max_queue_size = max_queue_size
        self._queue: asyncio.Queue[_PendingRequest] = asyncio.Queue(
            maxsize=max_queue_size
        )
        self._process_fn: Callable[..., Any] | None = None
        self._task: asyncio.Task[None] | None = None
        self._running = False

        # Metrics
        self.total_batches_processed: int = 0
        self.total_items_processed: int = 0
        self.total_rejections: int = 0

    @property
    def is_running(self) -> bool:
        return self._running

    @property
    def queue_size(self) -> int:
        """Current number of pending items in the queue."""
        return self._queue.qsize()

    def start(self, process_fn: Callable[..., Any]) -> None:
        """Start the batcher background loop.

        Args:
            process_fn: Callable that accepts a list of payloads and returns
                        a list of results (same length, same order).
        """
        self._process_fn = process_fn
        self._running = True
        self._task = asyncio.ensure_future(self._loop())

    async def stop(self) -> None:
        """Stop the batcher and drain remaining requests."""
        self._running = False
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None

        # Drain anything left
        await self._drain()

    async def submit(self, payload: Any) -> Any:
        """Submit a single request and await its result.

        The request will be batched with other concurrent requests.

        Raises:
            QueueFullError: If the queue has reached max_queue_size.
        """
        if self._queue.full():
            self.total_rejections += 1
            raise QueueFullError(
                f"Batcher queue is full ({self.max_queue_size} items). "
                "Request rejected — try again later."
            )

        loop = asyncio.get_running_loop()
        pending = _PendingRequest(
            payload=payload,
            future=loop.create_future(),
            enqueued_at=time.monotonic(),
        )
        try:
            self._queue.put_nowait(pending)
        except asyncio.QueueFull:
            self.total_rejections += 1
            raise QueueFullError(
                f"Batcher queue is full ({self.max_queue_size} items). "
                "Request rejected — try again later."
            )
        return await pending.future

    async def _loop(self) -> None:
        """Background loop that collects and dispatches batches."""
        while self._running:
            batch: list[_PendingRequest] = []

            try:
                # Wait for the first item
                first = await asyncio.wait_for(
                    self._queue.get(), timeout=0.1
                )
                batch.append(first)
            except asyncio.TimeoutError:
                continue

            # Collect more items up to batch size or wait timeout
            deadline = time.monotonic() + (self.max_wait_ms / 1000.0)
            while len(batch) < self.max_batch_size:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    break
                try:
                    item = await asyncio.wait_for(
                        self._queue.get(), timeout=remaining
                    )
                    batch.append(item)
                except asyncio.TimeoutError:
                    break

            # Process the batch
            await self._process_batch(batch)

    async def _drain(self) -> None:
        """Drain and process remaining queue items."""
        batch: list[_PendingRequest] = []
        while not self._queue.empty():
            try:
                item = self._queue.get_nowait()
                batch.append(item)
            except asyncio.QueueEmpty:
                break

        if batch:
            await self._process_batch(batch)

    async def _process_batch(self, batch: list[_PendingRequest]) -> None:
        """Invoke the process function and resolve futures."""
        if not batch or not self._process_fn:
            return

        payloads = [req.payload for req in batch]
        try:
            results = self._process_fn(payloads)
            if len(results) != len(batch):
                raise ValueError(
                    f"Process function returned {len(results)} results "
                    f"for {len(batch)} inputs"
                )
            for req, result in zip(batch, results):
                if not req.future.done():
                    req.future.set_result(result)
            self.total_batches_processed += 1
            self.total_items_processed += len(batch)
        except Exception as exc:
            logger.exception("Batch processing failed for %d items", len(batch))
            for req in batch:
                if not req.future.done():
                    req.future.set_exception(exc)