feat(realtime): add ordered fetch worker

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)

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())