remove background task deadlock updating linked state

perform the subsequent updates in an asyncio.Task to allow the original caller
to drop the lock for the other shared states.

Author: masenf

reflex/istate/shared.py

@@ -1,12 +1,56 @@
 """Base classes for shared / linked states."""
 
+import asyncio
 import contextlib
 from collections.abc import AsyncIterator
 
 from reflex.event import Event, get_hydrate_event
 from reflex.state import BaseState, State, _override_base_method, _substate_key
 from reflex.utils.exceptions import ReflexRuntimeError
 
+UPDATE_OTHER_CLIENT_TASKS: set[asyncio.Task] = set()
+
+
+def _do_update_other_tokens(
+    affected_tokens: set[str],
+    previous_dirty_vars: dict[str, set[str]],
+    state_type: type[BaseState],
+) -> list[asyncio.Task]:
+    """Update other clients after a shared state update.
+
+    Submit the updates in separate asyncio tasks to avoid deadlocking.
+
+    Args:
+        affected_tokens: The tokens to update.
+        previous_dirty_vars: The dirty vars to apply to other clients.
+        state_type: The type of the shared state.
+
+    Returns:
+        The list of asyncio tasks created to perform the updates.
+    """
+    from reflex.utils.prerequisites import get_app
+
+    app = get_app().app
+
+    async def _update_client(token: str):
+        async with app.modify_state(
+            _substate_key(token, state_type),
+            previous_dirty_vars=previous_dirty_vars,
+        ):
+            pass
+
+    tasks = []
+    for affected_token in affected_tokens:
+        # Don't send updates for disconnected clients.
+        if affected_token not in app.event_namespace._token_manager.token_to_socket:
+            continue
+        # TODO: remove disconnected client's after some time.
+        t = asyncio.create_task(_update_client(affected_token))
+        UPDATE_OTHER_CLIENT_TASKS.add(t)
+        t.add_done_callback(UPDATE_OTHER_CLIENT_TASKS.discard)
+        tasks.append(t)
+    return tasks
+
 
 class SharedStateBaseInternal(State):
     """The private base state for all shared states."""
@@ -215,23 +259,11 @@ async def _modify_linked_states(
 
         # Only propagate dirty vars when we are not already propagating from another state.
         if previous_dirty_vars is None:
-            from reflex.utils.prerequisites import get_app
-
-            app = get_app().app
-
-            for affected_token in affected_tokens:
-                # Don't send updates for disconnected clients.
-                if (
-                    affected_token
-                    not in app.event_namespace._token_manager.token_to_socket
-                ):
-                    continue
-                # TODO: remove disconnected client's after some time.
-                async with app.modify_state(
-                    _substate_key(affected_token, type(self)),
-                    previous_dirty_vars=current_dirty_vars,
-                ):
-                    pass
+            _do_update_other_tokens(
+                affected_tokens=affected_tokens,
+                previous_dirty_vars=current_dirty_vars,
+                state_type=type(self),
+            )
 
 
 class SharedState(SharedStateBaseInternal, mixin=True):

tests/integration/test_linked_state.py

@@ -16,7 +16,6 @@
 
 def LinkedStateApp():
     """Test that linked state works as expected."""
-    import asyncio
     from typing import Any
 
     import reflex as rx
@@ -72,13 +71,11 @@ async def bump_counter_bg(self):
                 async with self:
                     ss = await self.get_state(SharedState)
                     ss.counter += 1
-                await asyncio.sleep(0)
             async with self:
                 ss = await self.get_state(SharedState)
             for _ in range(5):
                 async with ss:
                     ss.counter += 1
-                await asyncio.sleep(0)
 
         @rx.event
         async def bump_counter_yield(self):