diff --git a/src/volcano_sdk/_realtime_fetch_worker.py b/src/volcano_sdk/_realtime_fetch_worker.py new file mode 100644 index 00000000..7f803d41 --- /dev/null +++ b/src/volcano_sdk/_realtime_fetch_worker.py @@ -0,0 +1,155 @@ +"""Bounded, ordered row fetching for realtime Postgres changes.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Generic, TypeVar + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + from .realtime import _PostgresFetchRequest + +FallbackT = TypeVar("FallbackT") +PostgresRecord = dict[str, Any] +_WORKER_CLOSED = "Postgres fetch worker is closed" +_INVALID_QUEUE_LIMIT = "queue_limit must be positive" + + +@dataclass(frozen=True, slots=True) +class PostgresFetchJob(Generic[FallbackT]): + """Pair a captured row request with its lightweight fallback.""" + + request: _PostgresFetchRequest + fallback: FallbackT + + +@dataclass(frozen=True, slots=True) +class PostgresFetchOutcome(Generic[FallbackT]): + """Describe either a fetched record or a failed fetch.""" + + job: PostgresFetchJob[FallbackT] + record: PostgresRecord | None = None + error: Exception | None = None + + +@dataclass(frozen=True, slots=True) +class _StopWorker: + pass + + +_STOP_WORKER = _StopWorker() + + +class PostgresFetchWorker(Generic[FallbackT]): + """Fetch queued rows serially and deliver outcomes in enqueue order.""" + + def __init__( + self, + fetch: Callable[[_PostgresFetchRequest], PostgresRecord | None], + deliver: Callable[[PostgresFetchOutcome[FallbackT]], Awaitable[None]], + *, + queue_limit: int, + ) -> None: + """Create a worker with a fixed pending-job limit.""" + if queue_limit <= 0: + raise ValueError(_INVALID_QUEUE_LIMIT) + self._fetch = fetch + self._deliver = deliver + self._queue: asyncio.Queue[PostgresFetchJob[FallbackT] | _StopWorker] = ( + asyncio.Queue(maxsize=queue_limit) + ) + self._state_lock = asyncio.Lock() + self._task: asyncio.Task[None] | None = None + self._stop_task: asyncio.Task[None] | None = None + self._closed = False + + async def enqueue(self, job: PostgresFetchJob[FallbackT]) -> None: + """Queue one fetch, applying backpressure when the queue is full.""" + async with self._state_lock: + if self._closed: + raise RuntimeError(_WORKER_CLOSED) + self._raise_worker_failure() + if self._task is None: + self._task = asyncio.create_task(self._run()) + await self._put_while_running(job, self._task) + + async def close(self) -> None: + """Drain accepted jobs and stop the worker.""" + async with self._state_lock: + task = self._task + if not self._closed: + self._closed = True + self._raise_worker_failure() + if task is not None and self._stop_task is None: + self._stop_task = asyncio.create_task(self._queue.put(_STOP_WORKER)) + stop_task = self._stop_task + if task is not None and stop_task is not None: + await self._wait_for_close(task, stop_task) + + def _raise_worker_failure(self) -> None: + task = self._task + if task is not None and task.done(): + task.result() + + async def _put_while_running( + self, + job: PostgresFetchJob[FallbackT], + task: asyncio.Task[None], + ) -> None: + put_task = asyncio.create_task(self._queue.put(job)) + try: + completed, _pending = await asyncio.wait( + (task, put_task), + return_when=asyncio.FIRST_COMPLETED, + ) + finally: + if not put_task.done(): + put_task.cancel() + await asyncio.gather(put_task, return_exceptions=True) + if task in completed: + task.result() + await put_task + + async def _wait_for_close( + self, + task: asyncio.Task[None], + stop_task: asyncio.Task[None], + ) -> None: + completed, _pending = await asyncio.wait( + (task, stop_task), + return_when=asyncio.FIRST_COMPLETED, + ) + if task in completed: + try: + task.result() + except BaseException: + stop_task.cancel() + await asyncio.gather(stop_task, return_exceptions=True) + raise + await asyncio.shield(stop_task) + await asyncio.shield(task) + + async def _run(self) -> None: + while True: + item = await self._queue.get() + try: + if isinstance(item, _StopWorker): + return + await self._fetch_and_deliver(item) + finally: + self._queue.task_done() + + async def _fetch_and_deliver(self, job: PostgresFetchJob[FallbackT]) -> None: + (result,) = await asyncio.gather( + asyncio.to_thread(self._fetch, job.request), + return_exceptions=True, + ) + if isinstance(result, BaseException): + if not isinstance(result, Exception): + raise result + outcome = PostgresFetchOutcome(job=job, error=result) + else: + outcome = PostgresFetchOutcome(job=job, record=result) + await self._deliver(outcome) diff --git a/tests/unit/test_realtime_fetch_worker.py b/tests/unit/test_realtime_fetch_worker.py new file mode 100644 index 00000000..e0b3f0c2 --- /dev/null +++ b/tests/unit/test_realtime_fetch_worker.py @@ -0,0 +1,232 @@ +from __future__ import annotations + +import asyncio +import threading +from contextlib import suppress +from typing import Any + +import pytest + +from volcano_sdk import realtime as realtime_module +from volcano_sdk._realtime_fetch_worker import ( + PostgresFetchJob, + PostgresFetchOutcome, + PostgresFetchWorker, +) + + +class BlockingRowFetch: + def __init__(self) -> None: + self.started = threading.Event() + self.release = threading.Event() + self.row_ids: list[int] = [] + + def __call__( + self, + request: realtime_module._PostgresFetchRequest, + ) -> dict[str, Any]: + row_id = request.row_id + assert isinstance(row_id, int) + self.row_ids.append(row_id) + if row_id == 1: + self.started.set() + if not self.release.wait(timeout=1): + message = "blocked row fetch was not released" + raise TimeoutError(message) + return {"id": row_id} + + +class FailingThenSuccessfulFetch: + def __init__(self, failure: Exception) -> None: + self.failure = failure + + def __call__( + self, + request: realtime_module._PostgresFetchRequest, + ) -> dict[str, Any]: + row_id = request.row_id + assert isinstance(row_id, int) + if row_id == 1: + raise self.failure + return {"id": row_id} + + +def fetch_job(row_id: int) -> PostgresFetchJob[str]: + return PostgresFetchJob( + request=realtime_module._PostgresFetchRequest( + database_name="app", + access_token="captured-token", + table="messages", + row_id=row_id, + ), + fallback=f"lightweight-{row_id}", + ) + + +async def wait_for_thread(event: threading.Event) -> None: + assert await asyncio.to_thread(event.wait, 0.2) + + +def test_postgres_fetch_worker_bounds_and_orders_fetches() -> None: + async def scenario() -> None: + fetch = BlockingRowFetch() + outcomes: list[PostgresFetchOutcome[str]] = [] + + async def deliver(outcome: PostgresFetchOutcome[str]) -> None: + outcomes.append(outcome) + + worker = PostgresFetchWorker(fetch, deliver, queue_limit=1) + await worker.enqueue(fetch_job(1)) + await wait_for_thread(fetch.started) + await worker.enqueue(fetch_job(2)) + third_enqueue = asyncio.create_task(worker.enqueue(fetch_job(3))) + await asyncio.sleep(0) + + assert not third_enqueue.done() + + fetch.release.set() + await asyncio.wait_for(third_enqueue, timeout=0.2) + await asyncio.wait_for(worker.close(), timeout=0.2) + + assert fetch.row_ids == [1, 2, 3] + assert [outcome.record for outcome in outcomes] == [ + {"id": 1}, + {"id": 2}, + {"id": 3}, + ] + assert [outcome.job.fallback for outcome in outcomes] == [ + "lightweight-1", + "lightweight-2", + "lightweight-3", + ] + + asyncio.run(scenario()) + + +def test_postgres_fetch_worker_rejects_an_unbounded_queue() -> None: + async def deliver(_outcome: PostgresFetchOutcome[str]) -> None: + return None + + with pytest.raises(ValueError, match="queue_limit must be positive"): + PostgresFetchWorker(BlockingRowFetch(), deliver, queue_limit=0) + + +def test_postgres_fetch_worker_delivers_fallback_and_continues_after_failure() -> None: + async def scenario() -> None: + failure = RuntimeError("database unavailable") + outcomes: list[PostgresFetchOutcome[str]] = [] + + async def deliver(outcome: PostgresFetchOutcome[str]) -> None: + outcomes.append(outcome) + + worker = PostgresFetchWorker( + FailingThenSuccessfulFetch(failure), + deliver, + queue_limit=2, + ) + await worker.enqueue(fetch_job(1)) + await worker.enqueue(fetch_job(2)) + await asyncio.wait_for(worker.close(), timeout=0.2) + + assert outcomes == [ + PostgresFetchOutcome(job=fetch_job(1), error=failure), + PostgresFetchOutcome(job=fetch_job(2), record={"id": 2}), + ] + + asyncio.run(scenario()) + + +def test_postgres_fetch_worker_rejects_work_after_delivery_failure() -> None: + async def scenario() -> None: + failure = RuntimeError("delivery failed") + delivery_started = asyncio.Event() + + async def fail_delivery(_outcome: PostgresFetchOutcome[str]) -> None: + delivery_started.set() + raise failure + + worker = PostgresFetchWorker( + BlockingRowFetch(), + fail_delivery, + queue_limit=1, + ) + await worker.enqueue(fetch_job(2)) + await delivery_started.wait() + await asyncio.sleep(0) + + with pytest.raises(RuntimeError) as raised: + await worker.enqueue(fetch_job(3)) + + assert raised.value is failure + + asyncio.run(scenario()) + + +def test_postgres_fetch_worker_recovers_from_cancelled_close() -> None: + async def scenario() -> None: + fetch = BlockingRowFetch() + outcomes: list[PostgresFetchOutcome[str]] = [] + + async def deliver(outcome: PostgresFetchOutcome[str]) -> None: + outcomes.append(outcome) + + worker = PostgresFetchWorker(fetch, deliver, queue_limit=1) + await worker.enqueue(fetch_job(1)) + await wait_for_thread(fetch.started) + await worker.enqueue(fetch_job(2)) + + interrupted_close = asyncio.create_task(worker.close()) + await asyncio.sleep(0) + interrupted_close.cancel() + with suppress(asyncio.CancelledError): + await interrupted_close + + fetch.release.set() + try: + await asyncio.wait_for(worker.close(), timeout=0.2) + finally: + task = worker._task + if task is not None and not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + assert [outcome.record for outcome in outcomes] == [{"id": 1}, {"id": 2}] + + asyncio.run(scenario()) + + +def test_postgres_fetch_worker_unblocks_a_full_enqueue_after_failure() -> None: + async def scenario() -> None: + failure = RuntimeError("delivery failed") + delivery_started = asyncio.Event() + release_delivery = asyncio.Event() + + async def fail_delivery(_outcome: PostgresFetchOutcome[str]) -> None: + delivery_started.set() + await release_delivery.wait() + raise failure + + worker = PostgresFetchWorker( + BlockingRowFetch(), + fail_delivery, + queue_limit=1, + ) + await worker.enqueue(fetch_job(2)) + await delivery_started.wait() + await worker.enqueue(fetch_job(3)) + blocked_enqueue = asyncio.create_task(worker.enqueue(fetch_job(4))) + await asyncio.sleep(0) + assert not blocked_enqueue.done() + + release_delivery.set() + try: + with pytest.raises(RuntimeError) as raised: + await asyncio.wait_for(blocked_enqueue, timeout=0.2) + finally: + task = worker._task + if task is not None: + await asyncio.gather(task, return_exceptions=True) + + assert raised.value is failure + + asyncio.run(scenario())