Fix up unit tests for OPLOCK_ENABLED mode

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/istate/manager/redis.py

@@ -603,19 +603,21 @@ async def do_flush() -> None:
             async with state_lock:
                 # Write the state to redis while no one else can modify the cached copy.
                 state = self._cached_states.pop(client_token, None)
-                if state:
-                    if self._debug_enabled:
+                try:
+                    if state:
+                        if self._debug_enabled:
+                            console.debug(
+                                f"{SMR} [{time.monotonic() - start:.3f}] {client_token} lease breaker {lock_id.decode()} flushing state"
+                            )
+                        await self.set_state(token, state, lock_id=lock_id, **context)
+                finally:
+                    if (current_lease := self._local_leases.get(client_token)) is task:
+                        self._local_leases.pop(client_token, None)
+                        # TODO: clean up the cached states locks periodically
+                    elif self._debug_enabled:
                         console.debug(
-                            f"{SMR} [{time.monotonic() - start:.3f}] {client_token} lease breaker {lock_id.decode()} flushing state"
+                            f"{SMR} [{time.monotonic() - start:.3f}] {client_token} lease breaker {lock_id.decode()} cleanup of {task=} found different task in _local_leases {current_lease=}."
                         )
-                    await self.set_state(token, state, lock_id=lock_id, **context)
-                if (current_lease := self._local_leases.get(client_token)) is task:
-                    self._local_leases.pop(client_token, None)
-                    # TODO: clean up the cached states locks periodically
-                elif self._debug_enabled:
-                    console.debug(
-                        f"{SMR} [{time.monotonic() - start:.3f}] {client_token} lease breaker {lock_id.decode()} cleanup of {task=} found different task in _local_leases {current_lease=}."
-                    )
 
         async def lease_breaker():
             cancelled_error: asyncio.CancelledError | None = None
@@ -638,6 +640,14 @@ async def lease_breaker():
                 try:
                     # Shield the flush from cancellation to ensure it always runs to completion.
                     await asyncio.shield(do_flush())
+                except Exception as e:
+                    # Propagate exception to the main loop, since we have nowhere to catch it.
+                    if not isinstance(e, asyncio.CancelledError):
+                        asyncio.get_running_loop().call_exception_handler({
+                            "message": "Exception in Redis State Manager lease breaker",
+                            "exception": e,
+                        })
+                    raise
                 finally:
                     # Re-raise any cancellation error after cleaning up.
                     if cancelled_error is not None:
@@ -993,7 +1003,7 @@ async def _lock(self, token: str):
                         console.debug(
                             f"{SMR} [{time.monotonic() - start:.3f}] {lock_key.decode()} released by {lock_id.decode()}"
                         )
-                else:
+                elif deleted_lock_id is not None:
                     # This can happen if the caller never tried to `set_state` before the lock expired and is a pretty bad bug.
                     console.warn(
                         f"{lock_key.decode()} was released by {lock_id.decode()}, but it belonged to {deleted_lock_id.decode()}. This is a bug."
@@ -1020,5 +1030,6 @@ async def close(self):
             # Then cancel all outstanding leases and write the cached states to redis.
             for lease_task in self._local_leases.values():
                 lease_task.cancel()
+            await asyncio.gather(*self._local_leases.values(), return_exceptions=True)
         finally:
             await self.redis.aclose(close_connection_pool=True)

tests/units/test_app.py

@@ -32,6 +32,7 @@
 from reflex.components.core.cond import Cond
 from reflex.components.radix.themes.typography.text import Text
 from reflex.constants.state import FIELD_MARKER
+from reflex.environment import environment
 from reflex.event import Event
 from reflex.istate.manager.disk import StateManagerDisk
 from reflex.istate.manager.memory import StateManagerMemory
@@ -990,6 +991,9 @@ def getlist(key: str):
             == StateUpdate(delta=delta, events=[], final=True).json() + "\n"
         )
 
+    if environment.REFLEX_OPLOCK_ENABLED.get():
+        await app.state_manager.close()
+
     current_state = await app.state_manager.get_state(_substate_key(token, state))
     state_dict = current_state.dict()[state.get_full_name()]
     assert state_dict["img_list" + FIELD_MARKER] == [
@@ -1296,6 +1300,9 @@ def _dynamic_state_event(name, val, **kwargs):
         with pytest.raises(StopAsyncIteration):
             await process_coro.__anext__()
 
+        if environment.REFLEX_OPLOCK_ENABLED.get():
+            await app.state_manager.close()
+
         # check that router data was written to the state_manager store
         state = await app.state_manager.get_state(substate_token)
         assert state.dynamic == exp_val
@@ -1363,6 +1370,9 @@ def _dynamic_state_event(name, val, **kwargs):
             await process_coro.__anext__()
 
         prev_exp_val = exp_val
+
+    if environment.REFLEX_OPLOCK_ENABLED.get():
+        await app.state_manager.close()
     state = await app.state_manager.get_state(substate_token)
     assert state.loaded == len(exp_vals)
     assert state.counter == len(exp_vals)
@@ -1403,6 +1413,9 @@ async def test_process_events(mocker: MockerFixture, token: str):
     async for _update in process(app, event, "mock_sid", {}, "127.0.0.1"):
         pass
 
+    if environment.REFLEX_OPLOCK_ENABLED.get():
+        await app.state_manager.close()
+
     assert (await app.state_manager.get_state(event.substate_token)).value == 5
     assert app._postprocess.call_count == 6  # pyright: ignore [reportAttributeAccessIssue]
 

tests/units/test_state.py

@@ -26,6 +26,7 @@
 from reflex.base import Base
 from reflex.constants import CompileVars, RouteVar, SocketEvent
 from reflex.constants.state import FIELD_MARKER
+from reflex.environment import environment
 from reflex.event import Event, EventHandler
 from reflex.istate.manager import StateManager
 from reflex.istate.manager.disk import StateManagerDisk
@@ -1740,6 +1741,10 @@ async def test_state_manager_modify_state(
         complex_1 = state.complex[1]
         assert isinstance(complex_1, MutableProxy)
         state.complex[3] = complex_1
+
+    if environment.REFLEX_OPLOCK_ENABLED.get():
+        await state_manager.close()
+
     # lock should be dropped after exiting the context
     if isinstance(state_manager, StateManagerRedis):
         assert (await state_manager.redis.get(f"{token}_lock")) is None
@@ -1783,6 +1788,9 @@ async def _coro():
     for f in asyncio.as_completed(tasks):
         await f
 
+    if environment.REFLEX_OPLOCK_ENABLED.get():
+        await state_manager.close()
+
     assert (await state_manager.get_state(substate_token)).num1 == exp_num1
 
     if isinstance(state_manager, StateManagerRedis):
@@ -1837,12 +1845,35 @@ async def test_state_manager_lock_expire(
     state_manager_redis.lock_expiration = LOCK_EXPIRATION
     state_manager_redis.lock_warning_threshold = LOCK_WARNING_THRESHOLD
 
+    loop_exception = None
+
+    def loop_exception_handler(loop, context):
+        """Catch the LockExpiredError from the event loop.
+
+        Args:
+            loop: The event loop.
+            context: The exception context.
+        """
+        nonlocal loop_exception
+        loop_exception = context["exception"]
+
+    asyncio.get_event_loop().set_exception_handler(loop_exception_handler)
+
     async with state_manager_redis.modify_state(substate_token_redis):
         await asyncio.sleep(0.01)
 
-    with pytest.raises(LockExpiredError):
+    if environment.REFLEX_OPLOCK_ENABLED.get():
         async with state_manager_redis.modify_state(substate_token_redis):
             await asyncio.sleep(LOCK_EXPIRE_SLEEP)
+        await asyncio.sleep(LOCK_EXPIRE_SLEEP)
+        assert loop_exception is not None
+        with pytest.raises(LockExpiredError):
+            raise loop_exception
+    else:
+        with pytest.raises(LockExpiredError):
+            async with state_manager_redis.modify_state(substate_token_redis):
+                await asyncio.sleep(LOCK_EXPIRE_SLEEP)
+        assert loop_exception is None
 
 
 @pytest.mark.asyncio
@@ -1862,6 +1893,20 @@ async def test_state_manager_lock_expire_contend(
     state_manager_redis.lock_expiration = LOCK_EXPIRATION
     state_manager_redis.lock_warning_threshold = LOCK_WARNING_THRESHOLD
 
+    loop_exception = None
+
+    def loop_exception_handler(loop, context):
+        """Catch the LockExpiredError from the event loop.
+
+        Args:
+            loop: The event loop.
+            context: The exception context.
+        """
+        nonlocal loop_exception
+        loop_exception = context["exception"]
+
+    asyncio.get_event_loop().set_exception_handler(loop_exception_handler)
+
     order = []
     waiter_event = asyncio.Event()
 
@@ -1876,19 +1921,31 @@ async def _coro_waiter():
         await waiter_event.wait()
         async with state_manager_redis.modify_state(substate_token_redis) as state:
             order.append("waiter")
-            assert state.num1 != unexp_num1
             state.num1 = exp_num1
 
     tasks = [
         asyncio.create_task(_coro_blocker()),
         asyncio.create_task(_coro_waiter()),
     ]
-    with pytest.raises(LockExpiredError):
-        await tasks[0]
-    await tasks[1]
+    if environment.REFLEX_OPLOCK_ENABLED.get():
+        await tasks[0]  # Doesn't raise during `modify_state`, only on exit
+        await tasks[1]
+        await asyncio.sleep(LOCK_EXPIRE_SLEEP)
+        assert loop_exception is not None
+        with pytest.raises(LockExpiredError):
+            raise loop_exception
+        # In oplock mode, the blocker block's both updates
+        assert (await state_manager_redis.get_state(substate_token_redis)).num1 == 0
+    else:
+        with pytest.raises(LockExpiredError):
+            await tasks[0]
+        await tasks[1]
+        assert loop_exception is None
+        assert (
+            await state_manager_redis.get_state(substate_token_redis)
+        ).num1 == exp_num1
 
     assert order == ["blocker", "waiter"]
-    assert (await state_manager_redis.get_state(substate_token_redis)).num1 == exp_num1
 
 
 @pytest.mark.asyncio
@@ -1923,8 +1980,12 @@ async def _coro_blocker():
     ]
 
     await tasks[0]
-    console_warn.assert_called()
-    assert console_warn.call_count == 7
+    if environment.REFLEX_OPLOCK_ENABLED.get():
+        # When Oplock is enabled, we don't warn when lock is held too long.
+        console_warn.assert_not_called()
+    else:
+        console_warn.assert_called()
+        assert console_warn.call_count == 7
 
 
 class CopyingAsyncMock(AsyncMock):
@@ -2079,6 +2140,9 @@ async def test_state_proxy(
     assert sp._self_actx is None
     assert sp.value2 == "42"
 
+    if environment.REFLEX_OPLOCK_ENABLED.get():
+        await mock_app.state_manager.close()
+
     # Get the state from the state manager directly and check that the value is updated
     gotten_state = await mock_app.state_manager.get_state(
         _substate_key(grandchild_state.router.session.client_token, grandchild_state)
@@ -2289,6 +2353,9 @@ async def test_background_task_no_block(mock_app: rx.App, token: str):
         await task
     assert not mock_app._background_tasks
 
+    if environment.REFLEX_OPLOCK_ENABLED.get():
+        await mock_app.state_manager.close()
+
     exp_order = [
         "background_task:start",
         "other",
@@ -2385,6 +2452,9 @@ async def test_background_task_reset(mock_app: rx.App, token: str):
         await task
     assert not mock_app._background_tasks
 
+    if environment.REFLEX_OPLOCK_ENABLED.get():
+        await mock_app.state_manager.close()
+
     assert (
         await mock_app.state_manager.get_state(
             _substate_key(token, BackgroundTaskState)