feat(realtime): order Postgres deliveries (#88)

Author: swkeever

src/volcano_sdk/realtime.py

@@ -11,6 +11,11 @@
 from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias, cast
 from urllib.parse import quote, urlsplit, urlunsplit
 
+from ._realtime_fetch_worker import (
+    PostgresFetchJob,
+    PostgresFetchOutcome,
+    PostgresFetchWorker,
+)
 from .database import Database
 from .models import JSONValue, _freeze_json
 
@@ -34,6 +39,7 @@
     importlib.import_module("centrifuge").CentrifugeError,
 )
 CALLBACK_QUEUE_LIMIT = 128
+POSTGRES_QUEUE_LIMIT = 128
 NO_PENDING_CALLBACK = object()
 CALLBACK_QUEUE_FULL_MESSAGE = (
     "Volcano realtime callback queue is full; publication dropped"
@@ -154,6 +160,19 @@ class _PostgresDeliveryIdentity:
     subscription_epoch: int
 
 
+@dataclass(frozen=True, slots=True)
+class _PostgresDelivery:
+    change: PostgresChange
+    identity: _PostgresDeliveryIdentity
+
+
+@dataclass(frozen=True, slots=True)
+class _CallbackDelivery:
+    event: str
+    data: Any
+    postgres_identity: _PostgresDeliveryIdentity | None = None
+
+
 @dataclass(slots=True)
 class _SessionBoundDatabaseContext:
     _transport: Transport
@@ -397,20 +416,20 @@ async def on_publication(self, ctx: PublicationContext) -> None:
     async def on_subscribing(self, ctx: Any) -> None:
         del ctx
         self._channel._subscribed = False
-        self._channel._end_postgres_epoch()
+        await self._channel._end_postgres_epoch()
         await self._channel._presence_unsubscribed()
 
     async def on_subscribed(self, ctx: Any) -> None:
         del ctx
         self._channel._subscribed = True
-        self._channel._begin_postgres_epoch()
+        await self._channel._begin_postgres_epoch()
         if self._channel._type == "presence":
             self._channel._schedule_presence_sync()
 
     async def on_unsubscribed(self, ctx: Any) -> None:
         del ctx
         self._channel._subscribed = False
-        self._channel._end_postgres_epoch()
+        await self._channel._end_postgres_epoch()
         await self._channel._presence_unsubscribed()
 
     async def on_join(self, ctx: Any) -> None:
@@ -498,7 +517,7 @@ def __init__(
         self._presence_lock = asyncio.Lock()
         self._presence_sync_task: asyncio.Task[None] | None = None
         self._presence_sync_pending = False
-        self._callback_queue: asyncio.Queue[tuple[str, Any]] = asyncio.Queue(
+        self._callback_queue: asyncio.Queue[_CallbackDelivery] = asyncio.Queue(
             maxsize=CALLBACK_QUEUE_LIMIT
         )
         self._callback_task: asyncio.Task[None] | None = None
@@ -507,6 +526,8 @@ def __init__(
         self._pending_presence_sync: Any = NO_PENDING_CALLBACK
         self._postgres_epoch = 0
         self._postgres_session_lineage = 0
+        self._postgres_lock = asyncio.Lock()
+        self._postgres_worker: PostgresFetchWorker[_PostgresDelivery] | None = None
 
     @property
     def name(self) -> str:
@@ -598,14 +619,25 @@ def _capture_postgres_delivery_identity(self) -> _PostgresDeliveryIdentity:
             subscription_epoch=self._postgres_epoch,
         )
 
-    def _begin_postgres_epoch(self) -> None:
-        if self._type == "postgres":
-            self._postgres_epoch += 1
-            self._postgres_session_lineage = self._realtime._connection_lineage()
+    async def _begin_postgres_epoch(self) -> None:
+        if self._type != "postgres":
+            return
+        await self._stop_postgres_worker()
+        self._postgres_epoch += 1
+        self._postgres_session_lineage = self._realtime._connection_lineage()
 
-    def _end_postgres_epoch(self) -> None:
-        if self._type == "postgres":
-            self._postgres_epoch += 1
+    async def _end_postgres_epoch(self) -> None:
+        if self._type != "postgres":
+            return
+        self._postgres_epoch += 1
+        await self._stop_postgres_worker()
+
+    async def _stop_postgres_worker(self) -> None:
+        async with self._postgres_lock:
+            worker = self._postgres_worker
+            self._postgres_worker = None
+            if worker is not None:
+                await worker.close()
 
     def _postgres_delivery_is_current(
         self,
@@ -622,14 +654,41 @@ def _postgres_delivery_is_current(
 
     async def _receive_postgres_change(self, data: Any) -> None:
         change = _postgres_change(data)
-        if change is None:
+        if change is None or not self._callbacks.get("*"):
             return
         if change.mode == "lightweight" and change.type == "DELETE":
             old_record = change.old_record
             if old_record is None and change.id is not None:
                 old_record = {"id": change.id}
             change = replace(change, old_record=old_record, id=None, mode=None)
-        await self._emit("*", change)
+        identity = self._capture_postgres_delivery_identity()
+        if not self._postgres_delivery_is_current(identity):
+            return
+        delivery = _PostgresDelivery(change=change, identity=identity)
+        async with self._postgres_lock:
+            if not self._postgres_delivery_is_current(identity):
+                return
+            worker = self._postgres_worker
+            if worker is None:
+                worker = PostgresFetchWorker(
+                    self._realtime._fetch_postgres_row,
+                    self._deliver_postgres,
+                    queue_limit=POSTGRES_QUEUE_LIMIT,
+                )
+                self._postgres_worker = worker
+            await worker.enqueue(PostgresFetchJob(request=None, fallback=delivery))
+
+    async def _deliver_postgres(
+        self,
+        outcome: PostgresFetchOutcome[_PostgresDelivery],
+    ) -> None:
+        delivery = outcome.job.fallback
+        if self._postgres_delivery_is_current(delivery.identity):
+            await self._emit(
+                "*",
+                delivery.change,
+                postgres_identity=delivery.identity,
+            )
 
     async def subscribe(self) -> None:
         """Subscribe to this channel."""
@@ -643,7 +702,13 @@ async def unsubscribe(self) -> None:
         """Unsubscribe from this channel."""
         await self._realtime._unsubscribe(self)
 
-    async def _emit(self, event: str, data: Any) -> None:
+    async def _emit(
+        self,
+        event: str,
+        data: Any,
+        *,
+        postgres_identity: _PostgresDeliveryIdentity | None = None,
+    ) -> None:
         if not self._callbacks.get(event):
             return
         if event == "presence_sync":
@@ -657,7 +722,9 @@ async def _emit(self, event: str, data: Any) -> None:
             )
             self._callback_stop = None
         try:
-            self._callback_queue.put_nowait((event, data))
+            self._callback_queue.put_nowait(
+                _CallbackDelivery(event, data, postgres_identity)
+            )
         except asyncio.QueueFull:
             if event == "presence_sync":
                 self._pending_presence_sync = data
@@ -682,11 +749,15 @@ async def _restart_callback_dispatcher(self, previous: asyncio.Task[None]) -> No
 
     async def _dispatch_callbacks(self, stop: asyncio.Event) -> None:
         while not stop.is_set():
-            event, data = await self._callback_queue.get()
+            delivery = await self._callback_queue.get()
             try:
-                for callback in tuple(self._callbacks.get(event, [])):
+                if not self._callback_delivery_is_current(delivery):
+                    continue
+                for callback in tuple(self._callbacks.get(delivery.event, [])):
+                    if not self._callback_delivery_is_current(delivery):
+                        break
                     active_task = asyncio.create_task(
-                        self._run_callback(callback, data)
+                        self._run_callback(callback, delivery)
                     )
                     self._active_callback_task = active_task
                     try:
@@ -707,15 +778,25 @@ async def _dispatch_callbacks(self, stop: asyncio.Event) -> None:
                 self._callback_queue.task_done()
                 self._enqueue_pending_presence_sync()
 
+    def _callback_delivery_is_current(self, delivery: _CallbackDelivery) -> bool:
+        identity = delivery.postgres_identity
+        return identity is None or self._postgres_delivery_is_current(identity)
+
     def _enqueue_pending_presence_sync(self) -> None:
         pending = self._pending_presence_sync
         if pending is NO_PENDING_CALLBACK or self._callback_queue.full():
             return
         self._pending_presence_sync = NO_PENDING_CALLBACK
-        self._callback_queue.put_nowait(("presence_sync", pending))
+        self._callback_queue.put_nowait(_CallbackDelivery("presence_sync", pending))
 
-    async def _run_callback(self, callback: MessageCallback, data: Any) -> None:
-        result = callback(data)
+    async def _run_callback(
+        self,
+        callback: MessageCallback,
+        delivery: _CallbackDelivery,
+    ) -> None:
+        if not self._callback_delivery_is_current(delivery):
+            return
+        result = callback(delivery.data)
         if inspect.isawaitable(result):
             await result
 
@@ -850,6 +931,8 @@ async def _cancel_presence_sync(self) -> None:
 
     async def _reset(self) -> None:
         self._subscription = None
+        self._subscribed = False
+        await self._end_postgres_epoch()
         await self._cancel_presence_sync()
         task = self._callback_task
         active_task = self._active_callback_task

tests/unit/test_realtime.py

@@ -503,6 +503,73 @@ async def scenario() -> None:
     asyncio.run(scenario())
 
 
+def test_realtime_drops_queued_postgres_callbacks_from_an_old_epoch() -> None:
+    official = FakeCentrifugeClient()
+    client = VolcanoClient(
+        anon_key="anon-key",
+        _transport=AuthTransport(),
+        _realtime_client_factory=FakeCentrifugeFactory(official),
+    )
+    client.auth.sign_in(email="user@example.com", password="secret")
+
+    async def scenario() -> None:
+        def insert(record_id: int) -> dict[str, Any]:
+            return {
+                "type": "INSERT",
+                "schema": "public",
+                "table": "messages",
+                "record": {"id": record_id},
+                "timestamp": "2026-09-03T12:00:00Z",
+            }
+
+        received: list[int] = []
+        first_started = asyncio.Event()
+        release_first = asyncio.Event()
+        third_received = asyncio.Event()
+        channel = client.realtime.channel(
+            "public:messages",
+            channel_type="postgres",
+        )
+
+        async def on_insert(change: Any) -> None:
+            record_id = change.record["id"]
+            if record_id == 1:
+                first_started.set()
+                await release_first.wait()
+            received.append(record_id)
+            if record_id == 3:
+                third_received.set()
+
+        channel.on_postgres_changes(
+            "INSERT",
+            schema="public",
+            table="messages",
+            callback=on_insert,
+        )
+        await channel.subscribe()
+        subscription = official.subscription
+        assert subscription is not None
+
+        await subscription.emit(insert(1))
+        await first_started.wait()
+        first_worker = channel._postgres_worker
+        assert first_worker is not None
+        await subscription.emit(insert(2))
+
+        await subscription.emit_subscribing()
+        assert channel._postgres_worker is None
+        release_first.set()
+        await subscription.emit_subscribed()
+        await subscription.emit(insert(3))
+        await asyncio.wait_for(third_received.wait(), timeout=0.2)
+
+        assert received == [1, 3]
+        assert channel._postgres_worker is not first_worker
+        await client.realtime.disconnect()
+
+    asyncio.run(scenario())
+
+
 def test_realtime_routes_immutable_rls_scoped_postgres_changes() -> None:
     official = FakeCentrifugeClient()
     client = VolcanoClient(
@@ -1082,15 +1149,20 @@ async def wait_forever() -> None:
         blocker = asyncio.create_task(wait_forever())
         channel._callback_task = blocker
         for _ in range(channel._callback_queue.maxsize):
-            channel._callback_queue.put_nowait(("presence_sync", {"version": 0}))
+            channel._callback_queue.put_nowait(
+                realtime_module._CallbackDelivery(
+                    "presence_sync",
+                    {"version": 0},
+                )
+            )
 
         await channel._emit("presence_sync", {"version": 1})
         channel._callback_queue.get_nowait()
         channel._callback_queue.task_done()
         await channel._emit("presence_sync", {"version": 2})
         queued: list[Any] = []
         while not channel._callback_queue.empty():
-            queued.append(channel._callback_queue.get_nowait()[1])
+            queued.append(channel._callback_queue.get_nowait().data)
             channel._callback_queue.task_done()
         channel._enqueue_pending_presence_sync()
         blocker.cancel()