From ec775cbfa95bb5d2cbaff0daa92dafad16eca2c2 Mon Sep 17 00:00:00 2001 From: YusefSyed <211442445+YusefSyed@users.noreply.github.com> Date: Sat, 14 Nov 2026 04:06:14 -0500 Subject: [PATCH 1/2] Keep the Redis pub/sub subscribers retrying through an outage RedisTokenManager's two subscriber tasks ran under ensure_task's default limit, so five errors within a minute ended them. After a Redis outage of more than a few seconds, clients already connected to a worker stopped receiving updates produced on other workers until a new socket connected to that worker and restarted the tasks. ensure_task now accepts exception_limit=None: the task retries without a limit, the delay doubles after each error in a row up to exception_delay_max (30 seconds by default), only the first error of a run is logged, and the coroutine function is passed a `recovered` callable that logs the end of the run and resets the delay. Both subscribers use it and call `recovered` when Redis confirms the subscription. The defaults are unchanged for other callers, and CancelledError and RuntimeError still propagate. On every confirmed subscription the socket record subscriber also stores this instance's own socket records that are missing from Redis, since a record that expired while it was not listening is not announced again. Fixes #7835 --- news/+redis-subscriber-outage.bugfix.md | 1 + reflex/utils/tasks.py | 38 ++- reflex/utils/token_manager.py | 43 ++- tests/units/mock_redis.py | 5 + tests/units/test_mock_redis.py | 21 ++ tests/units/utils/test_tasks.py | 174 ++++++++++ tests/units/utils/test_token_manager.py | 437 +++++++++++++++++++++++- 7 files changed, 709 insertions(+), 10 deletions(-) create mode 100644 news/+redis-subscriber-outage.bugfix.md diff --git a/news/+redis-subscriber-outage.bugfix.md b/news/+redis-subscriber-outage.bugfix.md new file mode 100644 index 00000000000..d92ac5a8592 --- /dev/null +++ b/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. diff --git a/reflex/utils/tasks.py b/reflex/utils/tasks.py index 282cadb7325..9c6dce14a5c 100644 --- a/reflex/utils/tasks.py +++ b/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,33 @@ 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 the delay may grow to when there is no exception limit. **kwargs: The keyword arguments to pass to the coroutine function. """ last_regular_loop_start = 0 exception_count = 0 + retry_delay = exception_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 = exception_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 +65,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 +91,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 +108,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 the delay may grow to 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 +135,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()}" diff --git a/reflex/utils/token_manager.py b/reflex/utils/token_manager.py index afb1685a5ad..153753d3e31 100644 --- a/reflex/utils/token_manager.py +++ b/reflex/utils/token_manager.py @@ -279,8 +279,34 @@ 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. + """ + pipeline = self.redis.pipeline(transaction=False) + for token, socket_record in self.token_to_socket.items(): + if socket_record.instance_id == self.instance_id: + pipeline.set( + self._get_redis_key(token), + pickle.dumps(socket_record), + ex=self.token_expiration, + nx=True, + ) + try: + await pipeline.execute() + except Exception as e: + logger.error(f"Redis error restoring socket records: {e}") + + 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 +315,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 +339,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 +450,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 +483,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, ) diff --git a/tests/units/mock_redis.py b/tests/units/mock_redis.py index 693da4192f9..ee28557346a 100644 --- a/tests/units/mock_redis.py +++ b/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 diff --git a/tests/units/test_mock_redis.py b/tests/units/test_mock_redis.py index d0c81ea908e..a3173a84892 100644 --- a/tests/units/test_mock_redis.py +++ b/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", + } diff --git a/tests/units/utils/test_tasks.py b/tests/units/utils/test_tasks.py index 5d2d7174581..9438d73eff1 100644 --- a/tests/units/utils/test_tasks.py +++ b/tests/units/utils/test_tasks.py @@ -1,5 +1,8 @@ import asyncio import contextlib +import logging +from collections.abc import Callable +from typing import Any from unittest.mock import Mock import pytest @@ -11,6 +14,43 @@ class NotSuppressedError(Exception): """An exception that should not be suppressed.""" +@pytest.fixture +def sleeps(monkeypatch: pytest.MonkeyPatch) -> list[float]: + """Record the delays passed to asyncio.sleep instead of waiting for them. + + Args: + monkeypatch: The pytest monkeypatch fixture. + + Returns: + The recorded delays, in call order. + """ + delays: list[float] = [] + real_sleep = asyncio.sleep + + async def sleep(delay: float) -> None: + delays.append(delay) + await real_sleep(0) + + monkeypatch.setattr(asyncio, "sleep", sleep) + return delays + + +def _task_logs(caplog: pytest.LogCaptureFixture) -> list[tuple[str, str]]: + """Get the level and message of every record logged by reflex.utils.tasks. + + Args: + caplog: The pytest log capture fixture. + + Returns: + The (level name, message) pairs, in log order. + """ + return [ + (record.levelname, record.getMessage()) + for record in caplog.records + if record.name == "reflex.utils.tasks" + ] + + @pytest.mark.asyncio async def test_ensure_task_suppresses_exceptions(): """Test that ensure_task suppresses specified exceptions.""" @@ -116,3 +156,137 @@ def test_ensure_task_no_runtime_error_suppression(): exception_limit=5, exception_limit_window=1.0, ) + + +async def test_ensure_task_default_limit( + sleeps: list[float], caplog: pytest.LogCaptureFixture +): + """By default every error is logged and the task raises on the fifth within a minute. + + Args: + sleeps: The recorded retry delays. + caplog: The pytest log capture fixture. + """ + call_count = 0 + + async def faulty_coro(): # noqa: RUF029 + nonlocal call_count + call_count += 1 + raise ValueError("Intentional error") # noqa: EM101 + + owner = Mock() + ensure_task( + owner=owner, + task_attribute="task", + coro_function=faulty_coro, + suppress_exceptions=[ValueError], + ) + + with pytest.raises(ValueError) as exc_info: + await asyncio.wait_for(owner.task, timeout=1) + + assert str(exc_info.value) == "Intentional error" + assert call_count == 5 + assert sleeps == [1.0, 1.0, 1.0, 1.0] + assert _task_logs(caplog) == [ + *[("ERROR", "faulty_coro: task error suppressed: Intentional error")] * 4, + ( + "ERROR", + "faulty_coro: task exceeded exception limit 5 within 60.0s: Intentional error", + ), + ] + + +@pytest.mark.parametrize( + ("delays", "expected_sleeps"), + [ + ({}, [1.0, 2.0, 4.0, 8.0, 16.0, 30.0, 30.0, 1.0]), + ( + {"exception_delay": 0.5, "exception_delay_max": 3.0}, + [0.5, 1.0, 2.0, 3.0, 3.0, 3.0, 3.0, 0.5], + ), + ], +) +async def test_ensure_task_no_limit( + delays: dict[str, Any], + expected_sleeps: list[float], + sleeps: list[float], + caplog: pytest.LogCaptureFixture, +): + """Without a limit the task backs off up to the cap and logs state changes only. + + Args: + delays: The delay arguments passed to ensure_task. + expected_sleeps: The retry delays the task should wait for. + sleeps: The recorded retry delays. + caplog: The pytest log capture fixture. + """ + caplog.set_level(logging.INFO, logger="reflex.utils.tasks") + call_count = 0 + + async def flaky_coro(recovered: Callable[[], None]): # noqa: RUF029 + nonlocal call_count + call_count += 1 + if call_count == 1: + # Nothing has failed yet, so there is no recovery to log. + recovered() + if call_count <= 7: + msg = f"down {call_count}" + raise ValueError(msg) + recovered() + if call_count == 8: + raise ValueError("down again") # noqa: EM101 + # A second call within the same run logs nothing. + recovered() + raise RuntimeError("Test Passed") # noqa: EM101 + + owner = Mock() + ensure_task( + owner=owner, + task_attribute="task", + coro_function=flaky_coro, + suppress_exceptions=[ValueError], + exception_limit=None, + # There is no window without a limit: an elapsed one does not end a run of errors. + exception_limit_window=0, + **delays, + ) + + with pytest.raises(RuntimeError) as exc_info: + await asyncio.wait_for(owner.task, timeout=1) + + assert str(exc_info.value) == "Test Passed" + assert call_count == 9 + assert sleeps == expected_sleeps + assert _task_logs(caplog) == [ + ("ERROR", "flaky_coro: task error, retrying until it recovers: down 1"), + ("INFO", "flaky_coro: task recovered"), + ("ERROR", "flaky_coro: task error, retrying until it recovers: down again"), + ("INFO", "flaky_coro: task recovered"), + ] + + +async def test_ensure_task_no_limit_cancel_during_retry_delay(): + """Cancelling a task that is waiting to retry ends it without waiting out the delay.""" + failed = asyncio.Event() + + async def faulty_coro(recovered: Callable[[], None]): # noqa: RUF029 + failed.set() + raise ValueError("Intentional error") # noqa: EM101 + + owner = Mock() + task = ensure_task( + owner=owner, + task_attribute="task", + coro_function=faulty_coro, + suppress_exceptions=[ValueError], + exception_delay=30.0, + exception_limit=None, + ) + await asyncio.wait_for(failed.wait(), timeout=1) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + + assert task.cancelled() diff --git a/tests/units/utils/test_token_manager.py b/tests/units/utils/test_token_manager.py index c9871d5dfb1..c001f6b1014 100644 --- a/tests/units/utils/test_token_manager.py +++ b/tests/units/utils/test_token_manager.py @@ -1,19 +1,27 @@ """Unit tests for TokenManager implementations.""" import asyncio +import dataclasses +import logging +import math import pickle import time -from collections.abc import Callable, Generator +from collections.abc import AsyncGenerator, Callable, Generator from contextlib import asynccontextmanager -from typing import cast -from unittest.mock import AsyncMock, Mock, patch +from typing import Any, cast +from unittest.mock import AsyncMock, Mock, call, patch import pytest +from redis.asyncio import Redis +from redis.exceptions import ConnectionError as RedisConnectionError +from redis.exceptions import ResponseError from reflex import config from reflex.app import EventNamespace from reflex.istate.data import RouterData from reflex.state import StateUpdate +from reflex.utils import token_manager +from reflex.utils.tasks import ensure_task from reflex.utils.token_manager import ( LocalTokenManager, RedisTokenManager, @@ -730,6 +738,429 @@ async def logged(channel: bytes, data: bytes, start: int = 0): assert pickle.loads(cast("bytes", await redis.get(key))) == record +@dataclasses.dataclass +class RedisOutage: + """Switch that makes the subscribers of a mock redis fail.""" + + # Raised by config_set, psubscribe and listen until the limit is reached. + error: Exception | None = None + # How many times the error is raised before redis works again. + limit: float = math.inf + # How many times the error was raised. + failures: int = 0 + # The most recently opened pubsub. + pubsub: Any = None + + +def breakable_pubsub(redis: Redis) -> RedisOutage: + """Let a test take the pubsub of a mock redis down. + + During an outage config_set and psubscribe raise its error, and listen + raises it in place of its next message, as redis-py does when the server + cannot be reached. + + Args: + redis: The mock redis client to patch. + + Returns: + The switch controlling the outage. + """ + outage = RedisOutage() + real_pubsub = redis.pubsub + # No token is connected elsewhere, so link_token_to_sid keeps the token. + redis.exists = AsyncMock(return_value=False) # pyright: ignore[reportAttributeAccessIssue] + + def fail() -> None: + if outage.error is not None and outage.failures < outage.limit: + outage.failures += 1 + raise outage.error + + async def config_set(name: str, value: str) -> None: # noqa: RUF029 + fail() + + @asynccontextmanager + async def pubsub() -> AsyncGenerator[Any, None]: + async with real_pubsub() as pubsub_mock: + real_psubscribe, real_listen = pubsub_mock.psubscribe, pubsub_mock.listen + + async def psubscribe(*patterns: str) -> None: + fail() + await real_psubscribe(*patterns) + + async def listen() -> AsyncGenerator[dict[str, Any], None]: + async for message in real_listen(): + fail() + yield message + + pubsub_mock.psubscribe = psubscribe # pyright: ignore[reportAttributeAccessIssue] + pubsub_mock.listen = listen + outage.pubsub = pubsub_mock + yield pubsub_mock + + redis.config_set = config_set # pyright: ignore[reportAttributeAccessIssue] + redis.pubsub = pubsub # pyright: ignore[reportAttributeAccessIssue] + return outage + + +def forget(redis: Redis, key: str) -> None: + """Drop a key from a mock redis without a keyspace notification. + + This is what a subscriber sees when a key expires while it is disconnected. + + Args: + redis: The mock redis client. + key: The key to drop. + """ + for internal in ("keys", "expire_times"): + redis._internals[internal].pop(key.encode(), None) # pyright: ignore[reportAttributeAccessIssue] + + +async def wait_until(condition: Callable[[], bool], timeout: float = 5.0) -> None: + """Poll until a condition holds. + + Args: + condition: The condition to wait for. + timeout: The maximum time to wait in seconds. + """ + async with asyncio.timeout(timeout): + while not condition(): # noqa: ASYNC110 + await asyncio.sleep(0.001) + + +def subscriptions(redis: Redis, pattern: str) -> int: + """Count the subscriptions made to a pattern on a mock redis. + + Args: + redis: The mock redis client. + pattern: The subscribed pattern. + + Returns: + How many times the pattern was subscribed to. + """ + return redis._internals["event_log"].count({ # pyright: ignore[reportAttributeAccessIssue] + "channel": b"psubscribe", + "data": pattern.encode(), + }) + + +def subscriber_logs(caplog: pytest.LogCaptureFixture) -> list[tuple[str, str, str]]: + """Get every record logged by the subscriber tasks and the token manager. + + Args: + caplog: The pytest log capture fixture. + + Returns: + The (logger name, level name, message) triples, in log order. + """ + return [ + (record.name, record.levelname, record.getMessage()) + for record in caplog.records + if record.name in ("reflex.utils.tasks", "reflex.utils.token_manager") + ] + + +def outage_logs(subscriber: str, error: str) -> list[tuple[str, str, str]]: + """Get the records a subscriber task logs over one outage. + + Args: + subscriber: The name of the subscriber's coroutine function. + error: The message of the first error of the outage. + + Returns: + The record for the start of the outage and the record for its end. + """ + return [ + ( + "reflex.utils.tasks", + "ERROR", + f"{subscriber}: task error, retrying until it recovers: {error}", + ), + ("reflex.utils.tasks", "INFO", f"{subscriber}: task recovered"), + ] + + +@pytest.fixture +def subscriber_retry_delay(monkeypatch: pytest.MonkeyPatch) -> Callable[[float], None]: + """Set the delay before the subscriber tasks retry for the first time. + + Args: + monkeypatch: The pytest monkeypatch fixture. + + Returns: + A function taking the delay, in seconds, for subscribers started after it. + """ + + def set_delay(exception_delay: float) -> None: + def ensure_task_with_delay(*args: Any, **kwargs: Any) -> asyncio.Task: + return ensure_task(*args, **{**kwargs, "exception_delay": exception_delay}) + + monkeypatch.setattr(token_manager, "ensure_task", ensure_task_with_delay) + + return set_delay + + +async def test_lost_and_found_subscriber_survives_redis_outage( + subscriber_retry_delay: Callable[[float], None], + caplog: pytest.LogCaptureFixture, +): + """Deltas from another instance are delivered again after a long redis outage. + + Args: + subscriber_retry_delay: Sets the first retry delay of the subscribers. + caplog: The pytest log capture fixture. + """ + caplog.set_level(logging.INFO, logger="reflex.utils.tasks") + subscriber_retry_delay(0.001) + redis = mock_redis() + outage = breakable_pubsub(redis) + owner, publisher = RedisTokenManager(redis), RedisTokenManager(redis) + channel = f"channel:{owner._get_lost_and_found_key(owner.instance_id)}" + expected_logs = outage_logs("_subscribe_lost_and_found_updates", "redis down") + emit_update = AsyncMock() + before, after = ( + StateUpdate(delta={"state": {"count": 1}}), + StateUpdate(delta={"state": {"count": 2}}), + ) + + owner.ensure_lost_and_found_task(emit_update) + task = owner._lost_and_found_task + assert task is not None + try: + assert await owner.link_token_to_sid("token1", "sid1") is None + await wait_until(lambda: subscriptions(redis, channel) == 1) + assert await publisher.emit_lost_and_found("token1", before) + await wait_until(lambda: emit_update.await_count == 1) + + # Redis goes away: the subscriber notices on its next read, then fails + # to subscribe six times, more than the default exception limit allows. + outage.error, outage.limit = RedisConnectionError("redis down"), 7 + await redis.publish(channel, b"") + # mock_redis replays the last event before a subscription to it: make + # that an unrelated one, as redis does not replay anything. + await redis.set("unrelated", b"") + await wait_until( + lambda: task.done() or subscriber_logs(caplog)[-1:] == expected_logs[1:] + ) + assert subscriber_logs(caplog) == expected_logs + assert outage.failures == 7 + + assert await publisher.emit_lost_and_found("token1", after) + await wait_until(lambda: emit_update.await_count == 2) + assert not task.done() + finally: + await owner.close() + await publisher.close() + + assert emit_update.await_args_list == [ + call(before, "token1"), + call(after, "token1"), + ] + assert subscriber_logs(caplog) == expected_logs + + +async def test_socket_record_subscriber_survives_redis_outage( + subscriber_retry_delay: Callable[[float], None], + caplog: pytest.LogCaptureFixture, +): + """After a long redis outage, records that expired during it are stored again. + + Args: + subscriber_retry_delay: Sets the first retry delay of the subscribers. + caplog: The pytest log capture fixture. + """ + caplog.set_level(logging.INFO, logger="reflex.utils.tasks") + subscriber_retry_delay(0.001) + redis = mock_redis() + outage = breakable_pubsub(redis) + manager = RedisTokenManager(redis) + pattern = f"__keyspace@1__:{manager._get_redis_key('*')}" + expected_logs = outage_logs("_subscribe_socket_record_updates", "redis down") + expired_key, claimed_key, foreign_key = ( + manager._get_redis_key(token) for token in ("token1", "token2", "token3") + ) + keys = redis._internals["keys"] # pyright: ignore[reportAttributeAccessIssue] + + try: + assert await manager.link_token_to_sid("token1", "sid1") is None + assert await manager.link_token_to_sid("token2", "sid2") is None + task = manager._socket_record_task + assert task is not None + await wait_until(lambda: subscriptions(redis, pattern) == 1) + expired_record = pickle.dumps(manager.token_to_socket["token1"]) + # A cached record of a socket on another instance. + manager.token_to_socket["token3"] = SocketRecord( + instance_id="other-instance", sid="sid3" + ) + + # Redis goes away: the subscriber notices on its next read, then fails + # to subscribe six times, more than the default exception limit allows. + # It misses that one record is deleted and that another instance + # claims the other token. + outage.error, outage.limit = RedisConnectionError("redis down"), 7 + await redis.delete(expired_key) + claim = pickle.dumps(SocketRecord(instance_id="other-instance", sid="sid4")) + await redis.set(claimed_key, claim) + # mock_redis replays the last event before a subscription to it: make + # that an unrelated one, as redis does not replay anything. + await redis.set("unrelated", b"") + await wait_until( + lambda: task.done() or subscriber_logs(caplog)[-1:] == expected_logs[1:] + ) + assert subscriber_logs(caplog) == expected_logs + assert outage.failures == 7 + + await wait_until(lambda: expired_key.encode() in keys) + assert await redis.get(expired_key) == expired_record + assert redis._internals["expire_times"][expired_key.encode()] > time.monotonic() # pyright: ignore[reportAttributeAccessIssue] + # A record that still exists is not overwritten. + assert await redis.get(claimed_key) == claim + # Only this instance's own records are stored. + assert await redis.get(foreign_key) is None + + # The subscriber handles notifications again. + await redis.delete(expired_key) + await wait_until(lambda: expired_key.encode() in keys) + assert not task.done() + finally: + await manager.close() + + assert subscriber_logs(caplog) == expected_logs + + +async def test_socket_record_subscriber_retries_refused_config( + subscriber_retry_delay: Callable[[float], None], + caplog: pytest.LogCaptureFixture, +): + """A redis that refuses CONFIG SET is logged once and retried without a limit. + + Args: + subscriber_retry_delay: Sets the first retry delay of the subscribers. + caplog: The pytest log capture fixture. + """ + caplog.set_level(logging.INFO, logger="reflex.utils.tasks") + subscriber_retry_delay(0.001) + redis = mock_redis() + outage = breakable_pubsub(redis) + outage.error = ResponseError("unknown command 'CONFIG'") + manager = RedisTokenManager(redis) + + manager._ensure_socket_record_task() + task = manager._socket_record_task + assert task is not None + try: + await wait_until(lambda: outage.failures > 5 or task.done()) + assert subscriber_logs(caplog) == [ + outage_logs("_subscribe_socket_record_updates", "unknown command 'CONFIG'")[ + 0 + ] + ] + assert not task.done() + finally: + await manager.close() + + +async def test_socket_record_subscriber_restores_records_on_resubscribe( + caplog: pytest.LogCaptureFixture, +): + """Every confirmed subscription stores this instance's missing records again. + + Args: + caplog: The pytest log capture fixture. + """ + caplog.set_level(logging.INFO, logger="reflex.utils.tasks") + redis = mock_redis() + outage = breakable_pubsub(redis) + manager = RedisTokenManager(redis) + pattern = f"__keyspace@1__:{manager._get_redis_key('*')}" + lost_key, marker_key = ( + manager._get_redis_key(token) for token in ("token1", "token2") + ) + keys = redis._internals["keys"] # pyright: ignore[reportAttributeAccessIssue] + + async def resubscribe() -> None: + """Lose a record, then resubscribe the way redis-py does after it reconnects. + + Returns once the subscriber has handled a later notification, so the + confirmation before it has been handled too. + """ + forget(redis, lost_key) + await outage.pubsub.psubscribe(pattern) + await redis.delete(marker_key) + await wait_until(lambda: marker_key.encode() in keys) + + try: + assert await manager.link_token_to_sid("token1", "sid1") is None + assert await manager.link_token_to_sid("token2", "sid2") is None + task = manager._socket_record_task + assert task is not None + await wait_until(lambda: subscriptions(redis, pattern) == 1) + record, marker_record = ( + pickle.dumps(manager.token_to_socket[token]) + for token in ("token1", "token2") + ) + + await resubscribe() + assert await redis.get(lost_key) == record + + # A restore that fails is logged once and leaves the subscriber running. + pipeline = redis.pipeline + failing_pipeline = Mock() + failing_pipeline.execute = AsyncMock( + side_effect=RedisConnectionError("redis down") + ) + redis.pipeline = Mock(return_value=failing_pipeline) # pyright: ignore[reportAttributeAccessIssue] + await resubscribe() + assert subscriber_logs(caplog) == [ + ( + "reflex.utils.token_manager", + "ERROR", + "Redis error restoring socket records: redis down", + ), + ] + assert await redis.get(lost_key) is None + # Both records were sent in one round trip, to expire like any other. + redis.pipeline.assert_called_once_with(transaction=False) + assert failing_pipeline.set.call_args_list == [ + call(lost_key, record, ex=manager.token_expiration, nx=True), + call(marker_key, marker_record, ex=manager.token_expiration, nx=True), + ] + failing_pipeline.execute.assert_awaited_once_with() + + # The next confirmed subscription restores the record. + redis.pipeline = pipeline # pyright: ignore[reportAttributeAccessIssue] + await resubscribe() + assert await redis.get(lost_key) == record + assert not task.done() + finally: + await manager.close() + + assert len(subscriber_logs(caplog)) == 1 + + +async def test_close_during_subscriber_retry_delay( + subscriber_retry_delay: Callable[[float], None], +): + """Closing the manager does not wait for subscribers that are waiting to retry. + + Args: + subscriber_retry_delay: Sets the first retry delay of the subscribers. + """ + subscriber_retry_delay(30.0) + redis = mock_redis() + outage = breakable_pubsub(redis) + outage.error = RedisConnectionError("redis down") + manager = RedisTokenManager(redis) + + manager.ensure_lost_and_found_task(AsyncMock()) + manager._ensure_socket_record_task() + tasks = (manager._lost_and_found_task, manager._socket_record_task) + await wait_until(lambda: outage.failures == 2) + + await asyncio.wait_for(manager.close(), timeout=1) + + assert all(task is not None and task.cancelled() for task in tasks) + + @pytest.fixture def redis_url(): """Returns the Redis URL from the environment.""" From 2c47be417846777919b608d52dd472425ec247f6 Mon Sep 17 00:00:00 2001 From: YusefSyed <211442445+YusefSyed@users.noreply.github.com> Date: Sat, 14 Nov 2026 13:40:19 -0500 Subject: [PATCH 2/2] Adopt socket records claimed during an outage, cap the first retry delay The restore's SET NX keeps a record that another instance stored while the subscriber was not listening, but this instance kept its own local record for that token: updates produced here still went to its own socket, and a disconnect of that socket deleted the other instance's record. The restore now reads each key back in the same pipeline and, where Redis holds a different record than the one it wrote, refreshes the local entry through _get_token_owner, as _handle_socket_record_del does when its own SET NX is rejected. Without an exception limit, an exception_delay above exception_delay_max was used as is for the first retry and again after each recovery. Both now use the smaller of the two. --- reflex/utils/tasks.py | 9 +- reflex/utils/token_manager.py | 25 ++++-- tests/units/utils/test_tasks.py | 2 + tests/units/utils/test_token_manager.py | 111 +++++++++++++++++++++++- 4 files changed, 133 insertions(+), 14 deletions(-) diff --git a/reflex/utils/tasks.py b/reflex/utils/tasks.py index 9c6dce14a5c..5725d27841b 100644 --- a/reflex/utils/tasks.py +++ b/reflex/utils/tasks.py @@ -31,12 +31,13 @@ async def _run_forever( 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 the delay may grow to when there is no exception limit. + 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 - retry_delay = exception_delay + 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.""" @@ -44,7 +45,7 @@ def recovered() -> None: if exception_count: logger.info(f"{coro_function.__name__}: task recovered") exception_count = 0 - retry_delay = exception_delay + retry_delay = first_retry_delay if exception_limit is None: kwargs["recovered"] = recovered @@ -111,7 +112,7 @@ def ensure_task( 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 the delay may grow to when there is no exception limit. + 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. diff --git a/reflex/utils/token_manager.py b/reflex/utils/token_manager.py index 153753d3e31..6851ae3c6e8 100644 --- a/reflex/utils/token_manager.py +++ b/reflex/utils/token_manager.py @@ -283,21 +283,30 @@ 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. + 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: - pipeline.set( - self._get_redis_key(token), - pickle.dumps(socket_record), - ex=self.token_expiration, - nx=True, - ) + 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: - await pipeline.execute() + 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] diff --git a/tests/units/utils/test_tasks.py b/tests/units/utils/test_tasks.py index 9438d73eff1..42bae626fd4 100644 --- a/tests/units/utils/test_tasks.py +++ b/tests/units/utils/test_tasks.py @@ -205,6 +205,8 @@ async def faulty_coro(): # noqa: RUF029 {"exception_delay": 0.5, "exception_delay_max": 3.0}, [0.5, 1.0, 2.0, 3.0, 3.0, 3.0, 3.0, 0.5], ), + # A first delay above the cap is capped too, also after a recovery. + ({"exception_delay": 10.0, "exception_delay_max": 3.0}, [3.0] * 8), ], ) async def test_ensure_task_no_limit( diff --git a/tests/units/utils/test_token_manager.py b/tests/units/utils/test_token_manager.py index c001f6b1014..fcc06ea8e19 100644 --- a/tests/units/utils/test_token_manager.py +++ b/tests/units/utils/test_token_manager.py @@ -738,6 +738,105 @@ async def logged(channel: bytes, data: bytes, start: int = 0): assert pickle.loads(cast("bytes", await redis.get(key))) == record +async def test_restore_socket_records_adopts_newer_records( + caplog: pytest.LogCaptureFixture, +): + """A restore writes missing records back and adopts newer ones it finds instead. + + Args: + caplog: The pytest log capture fixture. + """ + redis = mock_redis() + manager = RedisTokenManager(redis) + key = manager._get_redis_key + own = { + token: SocketRecord(instance_id=manager.instance_id, sid=f"sid-{token}") + for token in ("lost", "kept", "claimed", "vanished") + } + foreign = SocketRecord(instance_id="other-instance", sid="sid-foreign") + newer = { + token: SocketRecord(instance_id="other-instance", sid=f"sid-newer-{token}") + for token in ("claimed", "vanished") + } + for token, record in {**own, "foreign": foreign}.items(): + manager.token_to_socket[token] = record + manager.sid_to_token[record.sid] = token + for token, record in {"kept": own["kept"], **newer}.items(): + await redis.set(key(token), pickle.dumps(record)) + real_pipeline = redis.pipeline + + def pipeline(transaction: bool = True) -> Any: + """Make the newer record of "vanished" disappear after the round trip. + + Args: + transaction: Passed on to the mock pipeline. + + Returns: + The mock pipeline. + """ + restore: Any = real_pipeline(transaction=transaction) + execute = restore.execute + + async def execute_then_disconnect() -> list[Any]: + results = await execute() + # Its socket on the other instance disconnects before the read. + await redis.delete(key("vanished")) + return results + + restore.execute = execute_then_disconnect + return restore + + redis.pipeline = pipeline # pyright: ignore[reportAttributeAccessIssue] + with patch.object( + manager, "_get_token_owner", wraps=manager._get_token_owner + ) as get_token_owner: + await manager._restore_socket_records() + + # Only the records that differ from this instance's are read again. + assert get_token_owner.await_args_list == [ + call("claimed", refresh=True), + call("vanished", refresh=True), + ] + assert await redis.get(key("lost")) == pickle.dumps(own["lost"]) + assert redis._internals["expire_times"][key("lost").encode()] > time.monotonic() # pyright: ignore[reportAttributeAccessIssue] + assert await redis.get(key("kept")) == pickle.dumps(own["kept"]) + assert await redis.get(key("claimed")) == pickle.dumps(newer["claimed"]) + assert await redis.get(key("vanished")) is None + assert await redis.get(key("foreign")) is None + assert manager.token_to_socket == { + **own, + "claimed": newer["claimed"], + "foreign": foreign, + } + assert manager.sid_to_token == { + "sid-lost": "lost", + "sid-kept": "kept", + "sid-newer-claimed": "claimed", + "sid-vanished": "vanished", + "sid-foreign": "foreign", + } + + # The socket that lost the claimed token cannot delete the newer record. + await manager.disconnect_token("claimed", "sid-claimed") + assert await redis.get(key("claimed")) == pickle.dumps(newer["claimed"]) + + # A restore that fails is logged once and reads nothing back. + failing_pipeline = Mock() + failing_pipeline.execute = AsyncMock( + side_effect=RedisConnectionError("redis down") + ) + redis.pipeline = Mock(return_value=failing_pipeline) # pyright: ignore[reportAttributeAccessIssue] + await manager._restore_socket_records() + assert get_token_owner.await_count == 2 + + assert [ + (record.levelname, record.getMessage()) + for record in caplog.records + if record.name == "reflex.utils.token_manager" + ] == [("ERROR", "Redis error restoring socket records: redis down")] + assert manager.token_to_socket["vanished"] == own["vanished"] + + @dataclasses.dataclass class RedisOutage: """Switch that makes the subscribers of a mock redis fail.""" @@ -962,7 +1061,7 @@ async def test_socket_record_subscriber_survives_redis_outage( subscriber_retry_delay: Callable[[float], None], caplog: pytest.LogCaptureFixture, ): - """After a long redis outage, records that expired during it are stored again. + """After a long redis outage, expired records are stored again and claimed ones adopted. Args: subscriber_retry_delay: Sets the first retry delay of the subscribers. @@ -998,7 +1097,8 @@ async def test_socket_record_subscriber_survives_redis_outage( # claims the other token. outage.error, outage.limit = RedisConnectionError("redis down"), 7 await redis.delete(expired_key) - claim = pickle.dumps(SocketRecord(instance_id="other-instance", sid="sid4")) + claim_record = SocketRecord(instance_id="other-instance", sid="sid4") + claim = pickle.dumps(claim_record) await redis.set(claimed_key, claim) # mock_redis replays the last event before a subscription to it: make # that an unrelated one, as redis does not replay anything. @@ -1021,6 +1121,13 @@ async def test_socket_record_subscriber_survives_redis_outage( await redis.delete(expired_key) await wait_until(lambda: expired_key.encode() in keys) assert not task.done() + + # The claim was adopted, so the socket that lost the token cannot + # delete it when it disconnects. + assert manager.token_to_socket["token2"] == claim_record + assert "sid2" not in manager.sid_to_token + await manager.disconnect_token("token2", "sid2") + assert await redis.get(claimed_key) == claim finally: await manager.close()