Keep the Redis pub/sub subscribers retrying through an outage

news/+redis-subscriber-outage.bugfix.md

@@ -0,0 +1 @@
+Fix state updates produced on another backend worker not reaching already-connected clients after a Redis outage longer than a few seconds. The Redis pub/sub subscribers now retry without a limit (after 1 second, doubling up to 30), log an outage once when it starts and once when it ends, and restore the worker's socket records that expired in the meantime.

reflex/utils/tasks.py

@@ -15,8 +15,9 @@ async def _run_forever(
     *args: Any,
     suppress_exceptions: list[type[BaseException]],
     exception_delay: float,
-    exception_limit: int,
+    exception_limit: int | None,
     exception_limit_window: float,
+    exception_delay_max: float,
     **kwargs: Any,
 ):
     """Wrapper to continuously run a coroutine function, suppressing certain exceptions.
@@ -27,15 +28,34 @@ async def _run_forever(
         suppress_exceptions: The exceptions to suppress.
         exception_delay: The delay between retries when an exception is suppressed.
         exception_limit: The maximum number of suppressed exceptions within the limit window before raising.
+            Without a limit (None), the delay doubles after each exception in a row, only the first of them
+            is logged, and the coroutine function is passed a `recovered` callable to call once it works again.
         exception_limit_window: The time window in seconds for counting suppressed exceptions.
+        exception_delay_max: The longest delay between retries when there is no exception limit.
         **kwargs: The keyword arguments to pass to the coroutine function.
     """
     last_regular_loop_start = 0
     exception_count = 0
+    first_retry_delay = min(exception_delay, exception_delay_max)
+    retry_delay = first_retry_delay
+
+    def recovered() -> None:
+        """End a run of suppressed exceptions: log it and start over from the first delay."""
+        nonlocal exception_count, retry_delay
+        if exception_count:
+            logger.info(f"{coro_function.__name__}: task recovered")
+            exception_count = 0
+            retry_delay = first_retry_delay
+
+    if exception_limit is None:
+        kwargs["recovered"] = recovered
 
     while True:
         # Reset the exception count when the limit window has elapsed since the last non-exception loop started.
-        if last_regular_loop_start + exception_limit_window < time.monotonic():
+        if (
+            exception_limit is not None
+            and last_regular_loop_start + exception_limit_window < time.monotonic()
+        ):
             exception_count = 0
         if not exception_count:
             last_regular_loop_start = time.monotonic()
@@ -46,6 +66,14 @@ async def _run_forever(
         except Exception as e:
             if any(isinstance(e, ex) for ex in suppress_exceptions):
                 exception_count += 1
+                if exception_limit is None:
+                    if exception_count == 1:
+                        logger.error(
+                            f"{coro_function.__name__}: task error, retrying until it recovers: {e}"
+                        )
+                    await asyncio.sleep(retry_delay)
+                    retry_delay = min(retry_delay * 2, exception_delay_max)
+                    continue
                 if exception_count >= exception_limit:
                     logger.error(
                         f"{coro_function.__name__}: task exceeded exception limit {exception_limit} within {exception_limit_window}s: {e}"
@@ -64,8 +92,9 @@ def ensure_task(
     *args: Any,
     suppress_exceptions: list[type[BaseException]] | None = None,
     exception_delay: float = 1.0,
-    exception_limit: int = 5,
+    exception_limit: int | None = 5,
     exception_limit_window: float = 60.0,
+    exception_delay_max: float = 30.0,
     task_context: Context | None = None,
     **kwargs: Any,
 ) -> asyncio.Task:
@@ -80,7 +109,10 @@ def ensure_task(
         suppress_exceptions: The exceptions to log and continue when running the coroutine.
         exception_delay: The delay between retries when an exception is suppressed.
         exception_limit: The maximum number of suppressed exceptions within the limit window before raising.
+            Without a limit (None), the delay doubles after each exception in a row, only the first of them
+            is logged, and the coroutine function is passed a `recovered` callable to call once it works again.
         exception_limit_window: The time window in seconds for counting suppressed exceptions.
+        exception_delay_max: The longest delay between retries when there is no exception limit.
         task_context: The context to use for the task.
         *args: The arguments to pass to the coroutine function.
         **kwargs: The keyword arguments to pass to the coroutine function.
@@ -104,6 +136,7 @@ def ensure_task(
             exception_delay=exception_delay,
             exception_limit=exception_limit,
             exception_limit_window=exception_limit_window,
+            exception_delay_max=exception_delay_max,
             **kwargs,
         )
         task_name = f"reflex_ensure_task|{type(owner).__name__}.{task_attribute}={coro_function.__name__}|{time.time()}"

reflex/utils/token_manager.py

@@ -279,8 +279,43 @@ async def _handle_socket_record_del(self, token: str) -> None:
         else:
             self.token_to_socket.pop(token, None)
 
-    async def _subscribe_socket_record_updates(self) -> None:
-        """Subscribe to Redis keyspace notifications for socket record updates."""
+    async def _restore_socket_records(self) -> None:
+        """Store this instance's socket records that are missing from Redis.
+
+        A record that expires while the subscriber is not listening is not
+        announced again, so the records are written back once it listens. A
+        newer record stored meanwhile is kept and adopted, as when a deleted
+        record cannot be restored.
+        """
+        pipeline = self.redis.pipeline(transaction=False)
+        records: dict[str, bytes] = {}
+        for token, socket_record in self.token_to_socket.items():
+            if socket_record.instance_id == self.instance_id:
+                key = self._get_redis_key(token)
+                records[token] = record_pkl = pickle.dumps(socket_record)
+                pipeline.set(key, record_pkl, ex=self.token_expiration, nx=True)
+                pipeline.get(key)
+        try:
+            results = await pipeline.execute()
+        except Exception as e:
+            logger.error(f"Redis error restoring socket records: {e}")
+            return
+        for (token, record_pkl), stored_pkl in zip(
+            records.items(), results[1::2], strict=True
+        ):
+            if stored_pkl != record_pkl:
+                # Adopt the newer record without waiting for its set
+                # notification, which may have been sent while not listening.
+                await self._get_token_owner(token, refresh=True)
+
+    async def _subscribe_socket_record_updates(
+        self, recovered: Callable[[], None]
+    ) -> None:
+        """Subscribe to Redis keyspace notifications for socket record updates.
+
+        Args:
+            recovered: Called each time Redis confirms the subscription.
+        """
         await enable_keyspace_notifications(self.redis)
         redis_db = self.redis.get_connection_kwargs().get("db", 0)
 
@@ -289,7 +324,10 @@ async def _subscribe_socket_record_updates(self) -> None:
                 f"__keyspace@{redis_db}__:{self._get_redis_key('*')}"
             )
             async for message in pubsub.listen():
-                if message["type"] == "pmessage":
+                if message["type"] == "psubscribe":
+                    recovered()
+                    await self._restore_socket_records()
+                elif message["type"] == "pmessage":
                     key = message["channel"].split(b":", 1)[1].decode()
                     token = key.replace(self._token_socket_record_prefix, "")
 
@@ -310,6 +348,7 @@ def _ensure_socket_record_task(self) -> None:
             task_attribute="_socket_record_task",
             coro_function=self._subscribe_socket_record_updates,
             suppress_exceptions=[Exception],
+            exception_limit=None,
         )
 
     async def link_token_to_sid(self, token: str, sid: str) -> str | None:
@@ -420,18 +459,22 @@ def _get_lost_and_found_key(instance_id: str) -> str:
     async def _subscribe_lost_and_found_updates(
         self,
         emit_update: Callable[[StateUpdate, str], Coroutine[None, None, None]],
+        recovered: Callable[[], None],
     ) -> None:
         """Subscribe to Redis channel notifications for lost and found deltas.
 
         Args:
             emit_update: The function to emit state updates.
+            recovered: Called each time Redis confirms the subscription.
         """
         async with self.redis.pubsub() as pubsub:
             await pubsub.psubscribe(
                 f"channel:{self._get_lost_and_found_key(self.instance_id)}"
             )
             async for message in pubsub.listen():
-                if message["type"] == "pmessage":
+                if message["type"] == "psubscribe":
+                    recovered()
+                elif message["type"] == "pmessage":
                     record = pickle.loads(message["data"])
                     await emit_update(record.update, record.token)
 
@@ -449,6 +492,7 @@ def ensure_lost_and_found_task(
             task_attribute="_lost_and_found_task",
             coro_function=self._subscribe_lost_and_found_updates,
             suppress_exceptions=[Exception],
+            exception_limit=None,
             emit_update=emit_update,
         )
 

tests/units/mock_redis.py

@@ -154,6 +154,10 @@ async def mock_getdel(key: KeyT) -> Any:
         await redis_mock.delete(key)
         return value
 
+    async def mock_publish(channel: KeyT, message: bytes) -> None:  # noqa: RUF029
+        # Unlike redis, the subscriptions that receive the message are not counted.
+        _event_log_append_notify({"channel": _key_bytes(channel), "data": message})
+
     async def mock_pexpire(key: KeyT, px: int, xx: bool = False) -> bool:  # noqa: RUF029
         _expire_keys()
         key = _key_bytes(key)
@@ -325,6 +329,7 @@ async def listen() -> AsyncGenerator[dict[str, Any], None]:
     redis_mock.srem = mock_srem
     redis_mock.scard = mock_scard
     redis_mock.pexpire = mock_pexpire
+    redis_mock.publish = mock_publish
     redis_mock.pipeline = pipeline
     redis_mock.eval = mock_eval
     redis_mock.pttl = pttl

tests/units/test_mock_redis.py

@@ -28,3 +28,24 @@ async def test_pubsub_listen_yields_pmessage_for_pattern_without_handler():
         "channel": b"__keyspace@1__:foo1",
         "data": b"set",
     }
+
+
+async def test_publish_reaches_pattern_subscribers():
+    """A published message is delivered to a pubsub whose pattern matches the channel."""
+    redis = mock_redis()
+    async with redis.pubsub() as pubsub:
+        await pubsub.psubscribe("channel:foo*")
+        messages = pubsub.listen()
+        confirmation = await anext(messages)
+        assert confirmation["type"] == "psubscribe"
+
+        await redis.publish("channel:bar", b"skipped")
+        await redis.publish("channel:foo1", b"payload")
+        message = await asyncio.wait_for(anext(messages), timeout=1)
+
+    assert message == {
+        "type": "pmessage",
+        "pattern": b"channel:foo*",
+        "channel": b"channel:foo1",
+        "data": b"payload",
+    }