ENG-8212: Redis Oplock implementation (#5932)

* Token manager tracks instance_id in token_to_socket

* RedisTokenManager: keep local dicts globally updated via pub/sub

* Implement lost+found for StateUpdate without websocket

When an update is emitted for a token, but the websocket for that token is on
another instance of the app, post it to the lost+found channel where other
instances are listening for updates to send to their clients.

* Implement `enumerate_tokens` for TokenManager

Set the groundwork for being able to broadcast updates to all connected states.

* Consolidate on `_get_token_owner`

* fix test_connection_banner.py: expect SocketRecord JSON

* Use a single lock waiter

For more efficient and fair lock queueing, each StateManagerRedis uses a single
task to monitor the keyspace for lock release/expire and then wakes up the next
caller that was waiting in the queue (no fairness between separate processes
though).

Now lockers will wait for an `asyncio.Event` which is set by the redis pubsub
waiter. If any locker waits longer than the lock_expiration, it will just try
to get the lock in case there was some mixup with the pub/sub, the locker won't
be blocked forever.

* Redis Oplock implementation

* When taking a lock from redis, hold it for 80% of the lock expiration timeout
* While the lock is held, other events processed against the instance will use
  the cached in-memory copy of the state.
* When the timeout expires or another process signals intention to access a
  locked state, flush the modifed states to redis and release the lock.

Set REFLEX_OPLOCK_ENABLED=1 to use this feature

* add test_background_task.py::test_fast_yielding

* Implement real redis-backed test cases for lost+found

* add some polling for the emit mocks since L+F doesn't happen immediately

* Fix up unit tests for OPLOCK_ENABLED mode

* support py3.10

* Do not track contended leases in-process

Always check redis for contended leases before granting a lease. It's a bit
slower, but much more reliable and avoids racy lock_expiration timeouts when
contention occurs before the lease is created or when the pubsub hasn't caught
up to reality.

Always start _lock_update_task in __post_init__ to avoid race where the lease
is granted, then contended, but the pubsub task hasn't started to catch the
contention.

* Add real+mock test cases for StateManagerRedis

* update test_state to use mock_redis when real redis is not available

* safe await cancelled task

* explicitly disable oplock for basic test_redis cases

* py3.10 support: asyncio.TimeoutError != TimeoutError

* break out of forever tasks when event loop goes away

No point in continually spamming "no running event loop" to the console.

* generalize "forever" tasks to centralize exception handling/retry

* remove unused arg

* less racy way test_ensure_task_limit_window_passed

* rename REFLEX_STATE_MANAGER_REDIS_DEBUG to match the class name

Author: masenf

.github/workflows/unit_tests.yml

@@ -62,6 +62,13 @@ jobs:
           export PYTHONUNBUFFERED=1
           export REFLEX_REDIS_URL=redis://localhost:6379
           uv run pytest tests/units --cov --no-cov-on-fail --cov-report=
+      - name: Run unit tests w/ redis and OPLOCK_ENABLED
+        if: ${{ matrix.os == 'ubuntu-latest' }}
+        run: |
+          export PYTHONUNBUFFERED=1
+          export REFLEX_REDIS_URL=redis://localhost:6379
+          export REFLEX_OPLOCK_ENABLED=true
+          uv run pytest tests/units --cov --no-cov-on-fail --cov-report=
       # Change to explicitly install v1 when reflex-hosting-cli is compatible with v2
       - name: Run unit tests w/ pydantic v1
         run: |

reflex/environment.py

@@ -727,6 +727,12 @@ class EnvironmentVariables:
     # How long to wait between automatic reload on frontend error to avoid reload loops.
     REFLEX_AUTO_RELOAD_COOLDOWN_TIME_MS: EnvVar[int] = env_var(10_000)
 
+    # Whether to enable debug logging for the redis state manager.
+    REFLEX_STATE_MANAGER_REDIS_DEBUG: EnvVar[bool] = env_var(False)
+
+    # Whether to opportunistically hold the redis lock to allow fast in-memory access while uncontended.
+    REFLEX_OPLOCK_ENABLED: EnvVar[bool] = env_var(False)
+
 
 environment = EnvironmentVariables()
 

reflex/istate/manager/redis.py

@@ -0,0 +1,109 @@
+"""Helpers for managing asyncio tasks."""
+
+import asyncio
+import time
+from collections.abc import Callable, Coroutine
+from typing import Any
+
+from reflex.utils import console
+
+
+async def _run_forever(
+    coro_function: Callable[..., Coroutine],
+    *args: Any,
+    suppress_exceptions: list[type[BaseException]],
+    exception_delay: float,
+    exception_limit: int,
+    exception_limit_window: float,
+    **kwargs: Any,
+):
+    """Wrapper to continuously run a coroutine function, suppressing certain exceptions.
+
+    Args:
+        coro_function: The coroutine function to run.
+        *args: The arguments to pass to the coroutine function.
+        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.
+        exception_limit_window: The time window in seconds for counting suppressed exceptions.
+        **kwargs: The keyword arguments to pass to the coroutine function.
+    """
+    last_regular_loop_start = 0
+    exception_count = 0
+
+    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():
+            exception_count = 0
+        if not exception_count:
+            last_regular_loop_start = time.monotonic()
+        try:
+            await coro_function(*args, **kwargs)
+        except (asyncio.CancelledError, RuntimeError):
+            raise
+        except Exception as e:
+            if any(isinstance(e, ex) for ex in suppress_exceptions):
+                exception_count += 1
+                if exception_count >= exception_limit:
+                    console.error(
+                        f"{coro_function.__name__}: task exceeded exception limit {exception_limit} within {exception_limit_window}s: {e}"
+                    )
+                    raise
+                console.error(f"{coro_function.__name__}: task error suppressed: {e}")
+                await asyncio.sleep(exception_delay)
+                continue
+            raise
+
+
+def ensure_task(
+    owner: Any,
+    task_attribute: str,
+    coro_function: Callable[..., Coroutine],
+    *args: Any,
+    suppress_exceptions: list[type[BaseException]] | None = None,
+    exception_delay: float = 1.0,
+    exception_limit: int = 5,
+    exception_limit_window: float = 60.0,
+    **kwargs: Any,
+) -> asyncio.Task:
+    """Ensure that a task is running for the given coroutine function.
+
+    Note: if the task is already running, args and kwargs are ignored.
+
+    Args:
+        owner: The owner of the task.
+        task_attribute: The attribute name to store/retrieve the task from the owner object.
+        coro_function: The coroutine function to run as a 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.
+        exception_limit_window: The time window in seconds for counting suppressed exceptions.
+        *args: The arguments to pass to the coroutine function.
+        **kwargs: The keyword arguments to pass to the coroutine function.
+
+    Returns:
+        The asyncio task running the coroutine function.
+    """
+    if suppress_exceptions is None:
+        suppress_exceptions = []
+    if RuntimeError in suppress_exceptions:
+        msg = "Cannot suppress RuntimeError exceptions which may be raised by asyncio machinery."
+        raise RuntimeError(msg)
+
+    task = getattr(owner, task_attribute, None)
+    if task is None or task.done():
+        asyncio.get_running_loop()  # Ensure we're in an event loop.
+        task = asyncio.create_task(
+            _run_forever(
+                coro_function,
+                *args,
+                suppress_exceptions=suppress_exceptions,
+                exception_delay=exception_delay,
+                exception_limit=exception_limit,
+                exception_limit_window=exception_limit_window,
+                **kwargs,
+            ),
+            name=f"reflex_ensure_task|{type(owner).__name__}.{task_attribute}={coro_function.__name__}|{time.time()}",
+        )
+        setattr(owner, task_attribute, task)
+    return task

reflex/utils/tasks.py

@@ -14,6 +14,7 @@
 from reflex.istate.manager.redis import StateManagerRedis
 from reflex.state import BaseState, StateUpdate
 from reflex.utils import console, prerequisites
+from reflex.utils.tasks import ensure_task
 
 if TYPE_CHECKING:
     from redis.asyncio import Redis
@@ -239,8 +240,13 @@ def _handle_socket_record_del(self, token: str) -> None:
         ) is not None and socket_record.instance_id != self.instance_id:
             self.sid_to_token.pop(socket_record.sid, None)
 
-    async def _subscribe_socket_record_updates(self, redis_db: int) -> None:
+    async def _subscribe_socket_record_updates(self) -> None:
         """Subscribe to Redis keyspace notifications for socket record updates."""
+        await StateManagerRedis(
+            state=BaseState, redis=self.redis
+        )._enable_keyspace_notifications()
+        redis_db = self.redis.get_connection_kwargs().get("db", 0)
+
         async with self.redis.pubsub() as pubsub:
             await pubsub.psubscribe(
                 f"__keyspace@{redis_db}__:{self._get_redis_key('*')}"
@@ -260,26 +266,14 @@ async def _subscribe_socket_record_updates(self, redis_db: int) -> None:
                     elif event == "set":
                         await self._get_token_owner(token, refresh=True)
 
-    async def _socket_record_updates_forever(self) -> None:
-        """Background task to monitor Redis keyspace notifications for socket record updates."""
-        await StateManagerRedis(
-            state=BaseState, redis=self.redis
-        )._enable_keyspace_notifications()
-        redis_db = self.redis.get_connection_kwargs().get("db", 0)
-        while True:
-            try:
-                await self._subscribe_socket_record_updates(redis_db)
-            except asyncio.CancelledError:  # noqa: PERF203
-                break
-            except Exception as e:
-                console.error(f"RedisTokenManager socket record update task error: {e}")
-
     def _ensure_socket_record_task(self) -> None:
         """Ensure the socket record updates subscriber task is running."""
-        if self._socket_record_task is None or self._socket_record_task.done():
-            self._socket_record_task = asyncio.create_task(
-                self._socket_record_updates_forever()
-            )
+        ensure_task(
+            owner=self,
+            task_attribute="_socket_record_task",
+            coro_function=self._subscribe_socket_record_updates,
+            suppress_exceptions=[Exception],
+        )
 
     async def link_token_to_sid(self, token: str, sid: str) -> str | None:
         """Link a token to a session ID with Redis-based duplicate detection.
@@ -386,23 +380,6 @@ async def _subscribe_lost_and_found_updates(
                     record = LostAndFoundRecord(**json.loads(message["data"].decode()))
                     await emit_update(StateUpdate(**record.update), record.token)
 
-    async def _lost_and_found_updates_forever(
-        self,
-        emit_update: Callable[[StateUpdate, str], Coroutine[None, None, None]],
-    ):
-        """Background task to monitor Redis lost and found deltas.
-
-        Args:
-            emit_update: The function to emit state updates.
-        """
-        while True:
-            try:
-                await self._subscribe_lost_and_found_updates(emit_update)
-            except asyncio.CancelledError:  # noqa: PERF203
-                break
-            except Exception as e:
-                console.error(f"RedisTokenManager lost and found task error: {e}")
-
     def ensure_lost_and_found_task(
         self,
         emit_update: Callable[[StateUpdate, str], Coroutine[None, None, None]],
@@ -412,10 +389,13 @@ def ensure_lost_and_found_task(
         Args:
             emit_update: The function to emit state updates.
         """
-        if self._lost_and_found_task is None or self._lost_and_found_task.done():
-            self._lost_and_found_task = asyncio.create_task(
-                self._lost_and_found_updates_forever(emit_update)
-            )
+        ensure_task(
+            owner=self,
+            task_attribute="_lost_and_found_task",
+            coro_function=self._subscribe_lost_and_found_updates,
+            suppress_exceptions=[Exception],
+            emit_update=emit_update,
+        )
 
     async def _get_token_owner(self, token: str, refresh: bool = False) -> str | None:
         """Get the instance ID of the owner of a token.

reflex/utils/token_manager.py

@@ -55,6 +55,11 @@ async def handle_event_yield_only(self):
                     yield State.increment()
                 await asyncio.sleep(0.005)
 
+        @rx.event(background=True)
+        async def fast_yielding(self):
+            for _ in range(1000):
+                yield State.increment()
+
         @rx.event
         def increment(self):
             self.counter += 1
@@ -202,6 +207,11 @@ def index() -> rx.Component:
                 on_click=State.disconnect_reconnect_background,
                 id="disconnect-reconnect-background",
             ),
+            rx.button(
+                "Fast Yielding",
+                on_click=State.fast_yielding,
+                id="fast-yielding",
+            ),
             rx.button("Reset", on_click=State.reset_counter, id="reset"),
         )
 
@@ -451,3 +461,28 @@ def test_disconnect_reconnect(
     )
     # Final update should come through on the new websocket connection
     AppHarness.expect(lambda: counter.text == "3", timeout=5)
+
+
+def test_fast_yielding(
+    background_task: AppHarness,
+    driver: WebDriver,
+    token: str,
+) -> None:
+    """Test that fast yielding works as expected.
+
+    Args:
+        background_task: harness for BackgroundTask app.
+        driver: WebDriver instance.
+        token: The token for the connected client.
+    """
+    assert background_task.app_instance is not None
+
+    # get a reference to all buttons
+    fast_yielding_button = driver.find_element(By.ID, "fast-yielding")
+
+    # get a reference to the counter
+    counter = driver.find_element(By.ID, "counter")
+    assert background_task._poll_for(lambda: counter.text == "0", timeout=5)
+
+    fast_yielding_button.click()
+    assert background_task._poll_for(lambda: counter.text == "1000", timeout=50)