diff --git a/news/+shared-state-disconnect-reap.feature.md b/news/+shared-state-disconnect-reap.feature.md new file mode 100644 index 00000000000..8db8e6b581c --- /dev/null +++ b/news/+shared-state-disconnect-reap.feature.md @@ -0,0 +1 @@ +Disconnected clients are unsubscribed from their linked shared states after a reconnect grace period (`REFLEX_SHARED_STATE_DISCONNECT_GRACE`, default 30s, 0 disables). The new `SharedState._on_subscriber_disconnected` hook lets shared states clean up per-client data (e.g. presence bookkeeping); a client reconnecting later re-subscribes automatically with its next event. diff --git a/packages/reflex-base/news/+shared-state-disconnect-reap.feature.md b/packages/reflex-base/news/+shared-state-disconnect-reap.feature.md new file mode 100644 index 00000000000..bdd01ef34da --- /dev/null +++ b/packages/reflex-base/news/+shared-state-disconnect-reap.feature.md @@ -0,0 +1 @@ +Add `REFLEX_SHARED_STATE_DISCONNECT_GRACE` (default 30s, 0 disables), the number of seconds a disconnected client may reconnect before it is unsubscribed from its linked shared states. diff --git a/packages/reflex-base/src/reflex_base/environment.py b/packages/reflex-base/src/reflex_base/environment.py index 00e38b9e882..131d8539690 100644 --- a/packages/reflex-base/src/reflex_base/environment.py +++ b/packages/reflex-base/src/reflex_base/environment.py @@ -797,6 +797,10 @@ class EnvironmentVariables: # The address to bind the HTTP client to. You can set this to "::" to enable IPv6. REFLEX_HTTP_CLIENT_BIND_ADDRESS: EnvVar[str | None] = env_var(None) + # Seconds a disconnected client may reconnect before it is unsubscribed + # from the shared states it was linked to. 0 disables the unsubscription. + REFLEX_SHARED_STATE_DISCONNECT_GRACE: EnvVar[int] = env_var(30) + # Maximum size of the message in the websocket server in bytes. REFLEX_SOCKET_MAX_HTTP_BUFFER_SIZE: EnvVar[int] = env_var( constants.POLLING_MAX_HTTP_BUFFER_SIZE diff --git a/reflex/app.py b/reflex/app.py index 97d7d7472e5..6ddbc55c6ac 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -2312,6 +2312,8 @@ def on_disconnect(self, sid: str) -> asyncio.Task | None: Returns: An asyncio Task for cleaning up the token, or None. """ + from reflex.istate.shared import schedule_disconnect_reap + if otel.enabled: otel.record_connection(-1) self._client_error_counts.pop(sid, None) @@ -2327,6 +2329,9 @@ def on_disconnect(self, sid: str) -> asyncio.Task | None: and (processor := self.app._event_processor) is not None ): processor._on_disconnect(disconnect_token) + # Unsubscribe the client from its linked shared states unless it + # reconnects within the grace period. + schedule_disconnect_reap(self.app, disconnect_token) # Use async cleanup through token manager task = asyncio.create_task( self._token_manager.disconnect_token(disconnect_token, sid), diff --git a/reflex/istate/shared.py b/reflex/istate/shared.py index a4d24de0657..c2305dc721c 100644 --- a/reflex/istate/shared.py +++ b/reflex/istate/shared.py @@ -3,11 +3,13 @@ import asyncio import contextlib import logging +import time from collections.abc import AsyncIterator from dataclasses import MISSING from typing import TYPE_CHECKING, Self, TypeVar from reflex_base.constants import ROUTER_DATA, ROUTER_VARS +from reflex_base.environment import environment from reflex_base.event import Event, get_hydrate_event from reflex_base.registry import RegistrationContext from reflex_base.utils.exceptions import ReflexRuntimeError @@ -17,9 +19,16 @@ from reflex.istate.manager.token import BaseStateToken from reflex.state import BaseState, OnLoadInternalState, State, _override_base_method +if TYPE_CHECKING: + from reflex.app import App + logger = logging.getLogger(__name__) UPDATE_OTHER_CLIENT_TASKS: set[asyncio.Task] = set() +# The latest reap per client token, the only one allowed to unlink it. +DISCONNECT_REAP_TASKS: dict[str, asyncio.Task] = {} +# Reaps that started unlinking; a newer disconnect no longer cancels them. +STARTED_REAP_TASKS: set[asyncio.Task] = set() LINKED_STATE = TypeVar("LINKED_STATE", bound="SharedStateBaseInternal") @@ -73,7 +82,9 @@ async def _update_client(token: str): pass for affected_token in affected_tokens: - # TODO: remove disconnected clients after some time. + # Disconnected clients are removed by schedule_disconnect_reap after + # the reconnect grace; until then the connectivity check above skips + # them. t = asyncio.create_task(_update_client(affected_token)) UPDATE_OTHER_CLIENT_TASKS.add(t) t.add_done_callback(_log_update_client_errors) @@ -81,6 +92,101 @@ async def _update_client(token: str): return tasks +def schedule_disconnect_reap(app: "App", token: str) -> asyncio.Task | None: + """Schedule unsubscribing a disconnected client from its linked shared states. + + Called on client disconnect. The reap runs after a grace period + (REFLEX_SHARED_STATE_DISCONNECT_GRACE) so brief reconnects (page reloads, + network blips) are no-ops; a client reaped too eagerly re-subscribes on its + next event through _internal_patch_linked_state. A new disconnect for the + same token restarts the grace; a cleanup already under way finishes the + shared state it is unlinking and leaves the rest to the new reap. + + Args: + app: The application object. + token: The client token that disconnected. + + Returns: + The scheduled reap task, or None when reaping is disabled. + """ + grace = environment.REFLEX_SHARED_STATE_DISCONNECT_GRACE.get() + if grace <= 0 or app._state is None: + return None + if ( + previous := DISCONNECT_REAP_TASKS.get(token) + ) is not None and previous not in STARTED_REAP_TASKS: + previous.cancel() + + task = asyncio.create_task( + _reap_disconnected_client(app, token, grace), + name=f"reflex_shared_state_reap|{token}|{time.time()}", + ) + DISCONNECT_REAP_TASKS[token] = task + + def _on_done(task: asyncio.Task) -> None: + if DISCONNECT_REAP_TASKS.get(token) is task: + DISCONNECT_REAP_TASKS.pop(token, None) + STARTED_REAP_TASKS.discard(task) + if not task.cancelled() and (exc := task.exception()) is not None: + logger.warning(f"Error reaping disconnected shared state client: {exc}") + + task.add_done_callback(_on_done) + return task + + +async def _reap_disconnected_client(app: "App", token: str, grace: float) -> None: + """Unsubscribe a client from its linked shared states unless it reconnected. + + Args: + app: The application object. + token: The client token that disconnected. + grace: Seconds to wait for a reconnect before unsubscribing. + """ + await asyncio.sleep(grace) + if (event_namespace := app.event_namespace) is None: + return + token_manager = event_namespace._token_manager + if await token_manager.is_token_connected(token): + # The client reconnected, possibly to another instance. + return + if app._state is None: + return + # Read-only peek at the client's links; racing a concurrent event is fine, + # since any later event through the link re-adds the client anyway. + root_state = await app.state_manager.get_state( + BaseStateToken(ident=token, cls=app._state) + ) + if not isinstance(root_state, State): + return + links = dict(root_state._reflex_internal_links or {}) + for state_name, linked_token in links.items(): + try: + state_cls = app._state.get_class_substate(state_name) + async with app.modify_state( + BaseStateToken(ident=linked_token, cls=state_cls) + ) as shared_root: + shared = await shared_root.get_state(state_cls) + if ( + not isinstance(shared, SharedState) + or token not in shared._linked_from + ): + continue + # Recheck under the lock: the client may have reconnected, or + # disconnected again and left this to a newer reap. + if ( + await token_manager.is_token_connected(token) + or (task := asyncio.current_task()) is None + or DISCONNECT_REAP_TASKS.get(token, task) is not task + ): + return + # Unlinking starts; a newer disconnect must not cancel it midway. + STARTED_REAP_TASKS.add(task) + shared._linked_from.discard(token) + await shared._on_subscriber_disconnected(token) + except Exception as e: + logger.warning(f"Error unsubscribing disconnected shared state client: {e}") + + @contextlib.asynccontextmanager async def _patch_state( original_state: BaseState, linked_state: BaseState, full_delta: bool = False @@ -536,6 +642,18 @@ def _init_bookkeeping(self, parent_state: BaseState | None = None) -> None: super()._init_bookkeeping(parent_state) self._previous_dirty_vars = set() + async def _on_subscriber_disconnected(self, client_token: str) -> None: + """Hook called when a subscribed client is unsubscribed after disconnecting. + + Override to clean up per-client data kept on the shared state (e.g. + presence bookkeeping). Runs with the shared token's state locked, after + the client token was removed from the subscriber set; mutations + propagate to the remaining linked clients. + + Args: + client_token: The client token that was unsubscribed. + """ + @classmethod def __init_subclass__(cls, **kwargs): """Initialize subclass and set up shared state fields. diff --git a/tests/units/istate/test_shared.py b/tests/units/istate/test_shared.py index e7cb2e73010..ce91e6322cc 100644 --- a/tests/units/istate/test_shared.py +++ b/tests/units/istate/test_shared.py @@ -9,9 +9,14 @@ import reflex as rx from reflex.istate.shared import ( + DISCONNECT_REAP_TASKS, + STARTED_REAP_TASKS, + SharedState, SharedStateBaseInternal, _do_update_other_tokens, _patch_state, + _reap_disconnected_client, + schedule_disconnect_reap, ) from reflex.state import BaseState, State from reflex.utils.token_manager import ( @@ -116,6 +121,236 @@ async def test_update_other_tokens_redis_cross_instance(redis_manager, mock_redi assert local_key not in [call.args[0] for call in mock_redis.get.call_args_list] +@pytest.fixture(autouse=True) +def _clear_reap_tasks(): + """Drop the reap bookkeeping left by tests calling the reap without scheduling it.""" + yield + DISCONNECT_REAP_TASKS.clear() + STARTED_REAP_TASKS.clear() + + +def _mock_reap_app(token_manager) -> tuple[Mock, Mock, list[str]]: + """Create a mock app with one client linked to one shared state. + + Returns: + The mock app, the mock shared state and the list collecting the + token idents passed to modify_state. + """ + modified_tokens: list[str] = [] + + shared = Mock(spec=SharedState) + shared._linked_from = {"gone", "other"} + shared._on_subscriber_disconnected = AsyncMock() + + @asynccontextmanager + async def modify_state(token, previous_dirty_vars=None): + modified_tokens.append(token.ident) + shared_root = Mock() + shared_root.get_state = AsyncMock(return_value=shared) + yield shared_root + + root_state = Mock(spec=State) + root_state._reflex_internal_links = {"state.shared_sub": "room-token"} + + app = Mock() + app.modify_state = modify_state + app.event_namespace = Mock() + app.event_namespace._token_manager = token_manager + app.state_manager.get_state = AsyncMock(return_value=root_state) + app._state = Mock() + app._state.get_class_substate = Mock(return_value=State) + return app, shared, modified_tokens + + +async def test_reap_disconnected_client_unsubscribes(): + """A client that stayed disconnected is removed from its linked states.""" + app, shared, modified_tokens = _mock_reap_app(LocalTokenManager()) + + await _reap_disconnected_client(app, "gone", grace=0) + + assert modified_tokens == ["room-token"] + assert shared._linked_from == {"other"} + shared._on_subscriber_disconnected.assert_awaited_once_with("gone") + + +async def test_reap_skips_reconnected_client(): + """A client that reconnected within the grace is left subscribed.""" + manager = LocalTokenManager() + manager.token_to_socket["gone"] = SocketRecord( + instance_id=manager.instance_id, sid="sid1" + ) + app, shared, modified_tokens = _mock_reap_app(manager) + + await _reap_disconnected_client(app, "gone", grace=0) + + assert modified_tokens == [] + assert shared._linked_from == {"gone", "other"} + shared._on_subscriber_disconnected.assert_not_awaited() + + +async def test_reap_skips_client_reconnected_before_lock(): + """A client reconnecting while the reap waits for the shared lock stays subscribed.""" + manager = LocalTokenManager() + app, shared, modified_tokens = _mock_reap_app(manager) + modify_state = app.modify_state + + @asynccontextmanager + async def reconnect_then_modify_state(token): + manager.token_to_socket["gone"] = SocketRecord( + instance_id=manager.instance_id, sid="sid2" + ) + async with modify_state(token) as shared_root: + yield shared_root + + app.modify_state = reconnect_then_modify_state + + await _reap_disconnected_client(app, "gone", grace=0) + + assert modified_tokens == ["room-token"] + assert shared._linked_from == {"gone", "other"} + shared._on_subscriber_disconnected.assert_not_awaited() + + +async def test_reap_skips_unsubscribed_client(): + """A token no longer in the subscriber set does not trigger the hook.""" + app, shared, modified_tokens = _mock_reap_app(LocalTokenManager()) + + await _reap_disconnected_client(app, "unknown", grace=0) + + assert modified_tokens == ["room-token"] + assert shared._linked_from == {"gone", "other"} + shared._on_subscriber_disconnected.assert_not_awaited() + + +def test_schedule_disconnect_reap_disabled(monkeypatch): + """Grace 0 disables scheduling entirely.""" + monkeypatch.setenv("REFLEX_SHARED_STATE_DISCONNECT_GRACE", "0") + app, _, _ = _mock_reap_app(LocalTokenManager()) + + assert schedule_disconnect_reap(app, "gone") is None + assert not DISCONNECT_REAP_TASKS + + +async def test_schedule_disconnect_reap_restarts_grace(monkeypatch): + """A new disconnect for the same token cancels the pending reap.""" + monkeypatch.setenv("REFLEX_SHARED_STATE_DISCONNECT_GRACE", "60") + app, _, modified_tokens = _mock_reap_app(LocalTokenManager()) + + first = schedule_disconnect_reap(app, "gone") + second = schedule_disconnect_reap(app, "gone") + assert first is not None + assert second is not None + with pytest.raises(asyncio.CancelledError): + await first + assert {"gone": second} == DISCONNECT_REAP_TASKS + + second.cancel() + with pytest.raises(asyncio.CancelledError): + await second + assert not DISCONNECT_REAP_TASKS + assert modified_tokens == [] + + +async def test_schedule_disconnect_reap_lets_started_cleanup_finish(monkeypatch): + """A new disconnect restarts the grace without cancelling a started cleanup.""" + monkeypatch.setenv("REFLEX_SHARED_STATE_DISCONNECT_GRACE", "60") + sleep = asyncio.sleep + + async def skip_grace(delay): + await sleep(0) + + monkeypatch.setattr(asyncio, "sleep", skip_grace) + app, shared, _ = _mock_reap_app(LocalTokenManager()) + hook_entered = asyncio.Event() + release_hook = asyncio.Event() + + async def slow_hook(token): + hook_entered.set() + await release_hook.wait() + + shared._on_subscriber_disconnected.side_effect = slow_hook + + first = schedule_disconnect_reap(app, "gone") + await hook_entered.wait() + second = schedule_disconnect_reap(app, "gone") + release_hook.set() + assert first is not None + assert second is not None + await first + await second + # Let the done callbacks release the tasks. + await sleep(0) + + assert shared._linked_from == {"other"} + shared._on_subscriber_disconnected.assert_awaited_once_with("gone") + assert not DISCONNECT_REAP_TASKS + assert not STARTED_REAP_TASKS + + +@pytest.mark.parametrize("blocked_link", [0, 1]) +async def test_superseded_reap_leaves_unlink_to_new_grace( + monkeypatch, blocked_link: int +): + """A reap waiting for a shared lock unlinks nothing once a newer disconnect restarts the grace.""" + monkeypatch.setenv("REFLEX_SHARED_STATE_DISCONNECT_GRACE", "60") + sleep = asyncio.sleep + first_grace, second_grace = asyncio.Event(), asyncio.Event() + graces = [first_grace, second_grace] + + async def wait_for_grace(delay): + await graces.pop(0).wait() + + monkeypatch.setattr(asyncio, "sleep", wait_for_grace) + app, _, _ = _mock_reap_app(LocalTokenManager()) + rooms = [Mock(spec=SharedState) for _ in range(2)] + for room in rooms: + room._linked_from = {"gone"} + room._on_subscriber_disconnected = AsyncMock() + app.state_manager.get_state.return_value._reflex_internal_links = { + "state.shared_0": "0", + "state.shared_1": "1", + } + lock_wanted = asyncio.Event() + release_lock = asyncio.Event() + + @asynccontextmanager + async def modify_state(token): + if token.ident == str(blocked_link): + lock_wanted.set() + await release_lock.wait() + shared_root = Mock() + shared_root.get_state = AsyncMock(return_value=rooms[int(token.ident)]) + yield shared_root + + app.modify_state = modify_state + + first_grace.set() + first = schedule_disconnect_reap(app, "gone") + await lock_wanted.wait() + # The client reconnects and drops again while the reap waits for the lock. + second = schedule_disconnect_reap(app, "gone") + release_lock.set() + assert first is not None + assert second is not None + await asyncio.wait([first]) + + for i, room in enumerate(rooms): + assert ("gone" in room._linked_from) is (i >= blocked_link) + rooms[blocked_link]._on_subscriber_disconnected.assert_not_awaited() + assert {"gone": second} == DISCONNECT_REAP_TASKS + + second_grace.set() + await second + # Let the done callbacks release the tasks. + await sleep(0) + + for room in rooms: + assert not room._linked_from + room._on_subscriber_disconnected.assert_awaited_once_with("gone") + assert not DISCONNECT_REAP_TASKS + assert not STARTED_REAP_TASKS + + async def test_no_fan_out_without_linked_clients(): """An event with nothing linked does not reach for the registered App.