Wake deadline consumers on enqueue instead of polling every second

2026-09-29 23:39:42 +00:00
parent e344ce2e30
commit b43e074728
7 changed files with 162 additions and 19 deletions
@@ -11,8 +11,11 @@ machinery; the *meaning* of each deadline belongs to the game engine:
  token), and the scheduler enqueues it in the store's shared deadline
  queue (a Redis sorted set in production, see
  :mod:`tavolo.platform.store`); enqueueing is idempotent;
-* a background consumer running on **every** worker polls the queue for
-  due entries and hands them to the engine's
+* enqueueing a deadline announces it on a shared pub/sub channel, so
+  every worker's consumer wakes immediately instead of waiting for its
+  next poll; the consumer also wakes on a slow heartbeat, which is the
+  backstop recovering notifications lost to a reconnect or a worker
+  crash. Due entries are handed to the engine's
  :meth:`~tavolo.platform.engine.GameEngine.fire_deadline` under the
  per-game lock; the engine revalidates the token against the live state
  and raises :class:`~tavolo.platform.errors.GameError` when the entry
@@ -65,8 +68,8 @@ class DeadlineScheduler:
    """Enqueue and fire the deadlines game engines ask for.

    Holds the store and the registry so a single instance serves every
-    game type. One consumer task (and its wake-up event) is kept per
-    event loop — tests run each test on a fresh loop.
+    game type. One consumer task, its listener task and its wake-up
+    event are kept per event loop — tests run each test on a fresh loop.

    ``context_binder`` is an optional async callable the consumer awaits
    before persisting a finished match. The consumer is a long-lived task
@@ -80,14 +83,19 @@ class DeadlineScheduler:
        store: GameStore,
        registry: GameRegistry,
        heartbeat_ms: int = 1000,
+        retry_ms: int = 1000,
        context_binder: Optional[Callable[[], Awaitable[None]]] = None,
    ) -> None:
        self._store = store
        self._registry = registry
        self._heartbeat = heartbeat_ms / 1000
+        # Backoff for entries that failed to process and for listener
+        # reconnects; deliberately independent of the (slow) heartbeat.
+        self._retry_delay = retry_ms / 1000
        self._context_binder = context_binder
        self._consumers: Dict[asyncio.AbstractEventLoop, asyncio.Task] = {}
        self._wake_events: Dict[asyncio.AbstractEventLoop, asyncio.Event] = {}
+        self._listeners: Dict[asyncio.AbstractEventLoop, asyncio.Task] = {}

    async def sync_deadline(self, session: GameSession) -> None:
        """Enqueue the deadline the current state implies, if any.
@@ -109,6 +117,16 @@ class DeadlineScheduler:
        self.ensure_consumer()
        due_at = deadline.due_at.timestamp()
        await self._store.add_deadline(encode(entry), due_at)
+        # Announce the enqueue so consumers on other workers wake up too.
+        # The entry is already queued, so a lost notification only delays
+        # their discovery to the next poll backstop.
+        try:
+            await self._store.notify_deadline()
+        except Exception:
+            log.warning(
+                "deadline enqueue notification failed; relying on the poll backstop",
+                exc_info=True,
+            )
        wake = self._wake_events.get(asyncio.get_running_loop())
        if wake is not None:
            wake.set()
@@ -190,15 +208,22 @@ class DeadlineScheduler:
            if old.is_closed():
                self._consumers.pop(old, None)
                self._wake_events.pop(old, None)
+                self._listeners.pop(old, None)
        task = self._consumers.get(loop)
        if task is None or task.done():
            self._wake_events[loop] = asyncio.Event()
            self._consumers[loop] = loop.create_task(self._run(loop))
+            listener = self._listeners.get(loop)
+            if listener is None or listener.done():
+                self._listeners[loop] = loop.create_task(self._listen(loop))
            log.debug("deadline consumer started")

    def stop_consumer(self, loop: asyncio.AbstractEventLoop) -> None:
        task = self._consumers.pop(loop, None)
        self._wake_events.pop(loop, None)
+        listener = self._listeners.pop(loop, None)
+        if listener is not None:
+            listener.cancel()
        if task is not None:
            task.cancel()

@@ -228,21 +253,48 @@ class DeadlineScheduler:
                log.exception("deadline consumer: poll failed; retrying")
                failed = True
            if failed:
-                # A failed entry is still due: back off to the heartbeat
-                # instead of polling it again immediately.
-                delay = self._heartbeat
+                # A failed entry is still due: back off briefly instead of
+                # polling it again immediately. Kept independent of the
+                # heartbeat, which is only a slow backstop.
+                delay = self._retry_delay
            try:
                await asyncio.wait_for(wake.wait(), timeout=delay)
            except asyncio.TimeoutError:
                pass

+    async def _listen(self, loop: asyncio.AbstractEventLoop) -> None:
+        """Wake the consumer on deadline-enqueue notifications.
+
+        The subscription is best effort: a dropped connection or a lost
+        message is recovered by the consumer's heartbeat backstop, so
+        this task only has to reconnect and keep forwarding signals.
+        """
+        while True:
+            try:
+                async with self._store.subscribe_deadlines() as events:
+                    async for _ in events:
+                        wake = self._wake_events.get(loop)
+                        if wake is None:
+                            # Shut down: the consumer is gone.
+                            return
+                        wake.set()
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                log.exception("deadline listener failed; reconnecting")
+            try:
+                await asyncio.sleep(self._retry_delay)
+            except asyncio.CancelledError:
+                raise
+

 class DeadlineSchedulerMixin(KayaMixin):
    """Run the deadline consumer for the whole app lifetime.

-    Every worker (and every pod) runs the same consumer; coordination
-    happens exclusively through the shared deadline queue and the
-    per-game locks, so any worker may fire any game's deadline.
+    Every worker (and every pod) runs the same consumer and listener;
+    coordination happens exclusively through the shared deadline queue,
+    its notification channel and the per-game locks, so any worker may
+    fire any game's deadline.
    """

    def __init__(self, scheduler: DeadlineScheduler) -> None:
@@ -26,8 +26,9 @@ signal (never the state) means updated state reaches connections on
 every worker without leaking hidden information into the channel.

 Timeouts are driven by a shared delayed-deadline queue: producers
-enqueue an opaque ``member`` string with a due timestamp, and a consumer
-on every worker polls for due entries (see
+enqueue an opaque ``member`` string with a due timestamp and announce it
+on a shared "deadline enqueued" channel, and a consumer on every worker
+wakes on that signal to process due entries (see
 :mod:`tavolo.platform.deadlines`). Delivery is at-least-once — entries
 are removed only after they are processed — so a worker dying
 mid-processing cannot lose a deadline; engines revalidate entries
@@ -55,6 +56,9 @@ GAME_KEY_PREFIX = "tavolo:game:"
 CODE_KEY_PREFIX = "tavolo:code:"
 CHANNEL_PREFIX = "tavolo:game:"
 DEADLINES_KEY = "tavolo:deadlines"
+# Channel announcing newly enqueued deadline entries, so every worker's
+# consumer can wake immediately instead of polling on its heartbeat.
+DEADLINES_CHANNEL = "tavolo:deadlines:events"
 # Sorted set of running games (score = expiry timestamp), maintained by the
 # stores so the telemetry gauge can count ongoing matches without loading
 # every session. Finished games are removed on their final save.
@@ -165,6 +169,18 @@ class GameStore(ABC):
    async def remove_deadline(self, member: str) -> None:
        """Remove ``member`` from the queue; a no-op when absent."""

+    @abstractmethod
+    async def notify_deadline(self) -> None:
+        """Wake every worker's deadline consumer (best effort).
+
+        Notification is advisory: consumers stay correct when a signal
+        is lost, because their poll backstop recovers it.
+        """
+
+    @abstractmethod
+    def subscribe_deadlines(self) -> AsyncContextManager[AsyncIterator[None]]:
+        """Async context manager yielding deadline-enqueue signals."""
+

 def _channel(game_id: str) -> str:
    return f"{CHANNEL_PREFIX}{game_id}:events"
@@ -266,6 +282,20 @@ class RedisGameStore(GameStore):
    async def remove_deadline(self, member: str) -> None:
        await self._redis.zrem(DEADLINES_KEY, member)

+    async def notify_deadline(self) -> None:
+        await self._redis.publish(DEADLINES_CHANNEL, "deadline")
+
+    @contextlib.asynccontextmanager
+    async def subscribe_deadlines(self) -> AsyncIterator[AsyncIterator[None]]:
+        pubsub = self._redis.pubsub()
+        await pubsub.subscribe(DEADLINES_CHANNEL)
+        try:
+            yield _redis_events(pubsub)
+        finally:
+            with contextlib.suppress(Exception):
+                await pubsub.unsubscribe(DEADLINES_CHANNEL)
+                await pubsub.aclose()
+

 async def _redis_events(pubsub) -> AsyncIterator[None]:
    async for message in pubsub.listen():
@@ -284,6 +314,7 @@ class InMemoryGameStore(GameStore):
        self._locks: Dict[str, asyncio.Lock] = {}
        self._subscribers: Dict[str, Set[asyncio.Queue]] = {}
        self._deadlines: Dict[str, float] = {}
+        self._deadline_subscribers: Set[asyncio.Queue] = set()

    def _lock_for(self, game_id: str) -> asyncio.Lock:
        lock = self._locks.get(game_id)
@@ -352,6 +383,19 @@ class InMemoryGameStore(GameStore):
    async def remove_deadline(self, member: str) -> None:
        self._deadlines.pop(member, None)

+    async def notify_deadline(self) -> None:
+        for queue in list(self._deadline_subscribers):
+            queue.put_nowait(_BUMP)
+
+    @contextlib.asynccontextmanager
+    async def subscribe_deadlines(self) -> AsyncIterator[AsyncIterator[None]]:
+        queue: asyncio.Queue = asyncio.Queue()
+        self._deadline_subscribers.add(queue)
+        try:
+            yield _queue_events(queue)
+        finally:
+            self._deadline_subscribers.discard(queue)
+

 async def _queue_events(queue: asyncio.Queue) -> AsyncIterator[None]:
    while True:
@@ -220,6 +220,7 @@ class DummyEngine(GameEngine):

 def make_platform(
    engines: Sequence[GameEngine] = (DummyEngine(),),
+    heartbeat_ms: int = 50,
 ) -> Tuple[KayaApp, Platform, TortoiseMixin]:
    """Build a throwaway app + platform wired with in-memory stores."""
    registry = GameRegistry(engines)
@@ -251,7 +252,7 @@ def make_platform(
    scheduler = DeadlineScheduler(
        game_store,
        registry,
-        heartbeat_ms=50,
+        heartbeat_ms=heartbeat_ms,
        context_binder=tortoise_mixin.ensure_context,
    )
    platform = Platform(
@@ -112,6 +112,38 @@ async def _finished(store, game_id: str) -> Optional[GameSession]:
    return None


+class NotificationTest(unittest.TestCase):
+    """Deadline-enqueue notifications wake consumers on other workers."""
+
+    @async_test
+    async def test_notification_wakes_a_long_heartbeat_consumer(self) -> None:
+        # The consumer's heartbeat is a minute away and nothing sets its
+        # local wake event: only the pub/sub notification can deliver the
+        # deadline in time. This is the cross-worker path.
+        _, platform, _ = make_platform(heartbeat_ms=60_000)
+        store = platform.game_store
+        scheduler = platform.scheduler
+        session = _started_session(deadline_in_seconds=0.1)
+        await store.save(session)
+        scheduler.ensure_consumer()
+        # Let the listener task subscribe before announcing.
+        await asyncio.sleep(0.05)
+
+        engine = platform.registry.require(session.game_type)
+        deadline = engine.next_deadline(session)
+        assert deadline is not None
+        member = encode({
+            "game_id": session.id,
+            "kind": deadline.kind,
+            "token": deadline.token,
+        })
+        await store.add_deadline(member, deadline.due_at.timestamp())
+        await store.notify_deadline()
+
+        result = await _wait_for(lambda: _plays_is(store, session.id, 1))
+        self.assertIsNotNone(result, "notification did not wake the consumer")
+
+
 class ProcessDueTest(unittest.TestCase):
    """Direct ``process_due`` behaviour: engine revalidation and idempotency."""

@@ -270,6 +270,17 @@ class InMemoryGameStoreTest(unittest.TestCase):
        self.assertEqual(60.0, await store.next_deadline())
        self.assertEqual(["b"], await store.due_deadlines(now=100.0))

+    @async_test
+    async def test_deadline_notifications_reach_subscribers(self) -> None:
+        store = InMemoryGameStore(_registry())
+        async with store.subscribe_deadlines() as events:
+            await store.notify_deadline()
+            signal = await asyncio.wait_for(anext(events), timeout=1.0)
+            self.assertIsNone(signal)
+        # The subscription is gone after the context exits: notifying
+        # again is a no-op rather than an error.
+        await store.notify_deadline()
+

 class RedisGameStoreTest(unittest.TestCase):
    @async_test