add context for lock expiry

Author: adhami3310

reflex/app.py

@@ -1758,7 +1758,9 @@ async def process(
             constants.RouteVar.CLIENT_IP: client_ip,
         })
         # Get the state for the session exclusively.
-        async with app.state_manager.modify_state(event.substate_token) as state:
+        async with app.state_manager.modify_state(
+            event.substate_token, context=event
+        ) as state:
             # When this is a brand new instance of the state, signal the
             # frontend to reload before processing it.
             if (

reflex/istate/manager/__init__.py

@@ -7,6 +7,7 @@
 
 from reflex import constants
 from reflex.config import get_config
+from reflex.event import Event
 from reflex.state import BaseState
 from reflex.utils import console, prerequisites
 from reflex.utils.exceptions import InvalidStateManagerModeError
@@ -71,21 +72,27 @@ async def get_state(self, token: str) -> BaseState:
         """
 
     @abstractmethod
-    async def set_state(self, token: str, state: BaseState):
+    async def set_state(
+        self, token: str, state: BaseState, *, context: Event | None = None
+    ):
         """Set the state for a token.
 
         Args:
             token: The token to set the state for.
             state: The state to set.
+            context: The event context.
         """
 
     @abstractmethod
     @contextlib.asynccontextmanager
-    async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
+    async def modify_state(
+        self, token: str, *, context: Event | None = None
+    ) -> AsyncIterator[BaseState]:
         """Modify the state for a token while holding exclusive lock.
 
         Args:
             token: The token to modify the state for.
+            context: The event context.
 
         Yields:
             The state for the token.

reflex/istate/manager/disk.py

@@ -12,6 +12,7 @@
 from typing_extensions import override
 
 from reflex.environment import environment
+from reflex.event import Event
 from reflex.istate.manager import StateManager, _default_token_expiration
 from reflex.state import BaseState, _split_substate_key, _substate_key
 from reflex.utils import console, path_ops, prerequisites
@@ -301,12 +302,15 @@ async def _schedule_process_write_queue(self):
                     await asyncio.sleep(0)  # Yield to allow the task to start.
 
     @override
-    async def set_state(self, token: str, state: BaseState):
+    async def set_state(
+        self, token: str, state: BaseState, *, context: Event | None = None
+    ):
         """Set the state for a token.
 
         Args:
             token: The token to set the state for.
             state: The state to set.
+            context: The event context.
         """
         client_token, _ = _split_substate_key(token)
         if self._write_debounce_seconds > 0:
@@ -325,11 +329,14 @@ async def set_state(self, token: str, state: BaseState):
 
     @override
     @contextlib.asynccontextmanager
-    async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
+    async def modify_state(
+        self, token: str, *, context: Event | None = None
+    ) -> AsyncIterator[BaseState]:
         """Modify the state for a token while holding exclusive lock.
 
         Args:
             token: The token to modify the state for.
+            context: The event context.
 
         Yields:
             The state for the token.
@@ -344,7 +351,7 @@ async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
         async with self._states_locks[client_token]:
             state = await self.get_state(token)
             yield state
-            await self.set_state(token, state)
+            await self.set_state(token, state, context=context)
 
     async def close(self):
         """Close the state manager, flushing any pending writes to disk."""

reflex/istate/manager/memory.py

@@ -7,6 +7,7 @@
 
 from typing_extensions import override
 
+from reflex.event import Event
 from reflex.istate.manager import StateManager
 from reflex.state import BaseState, _split_substate_key
 
@@ -43,23 +44,29 @@ async def get_state(self, token: str) -> BaseState:
         return self.states[token]
 
     @override
-    async def set_state(self, token: str, state: BaseState):
+    async def set_state(
+        self, token: str, state: BaseState, *, context: Event | None = None
+    ):
         """Set the state for a token.
 
         Args:
             token: The token to set the state for.
             state: The state to set.
+            context: The event context.
         """
         token = _split_substate_key(token)[0]
         self.states[token] = state
 
     @override
     @contextlib.asynccontextmanager
-    async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
+    async def modify_state(
+        self, token: str, *, context: Event | None = None
+    ) -> AsyncIterator[BaseState]:
         """Modify the state for a token while holding exclusive lock.
 
         Args:
             token: The token to modify the state for.
+            context: The event context.
 
         Yields:
             The state for the token.

reflex/istate/manager/redis.py

@@ -14,6 +14,7 @@
 
 from reflex.config import get_config
 from reflex.environment import environment
+from reflex.event import Event
 from reflex.istate.manager import StateManager, _default_token_expiration
 from reflex.state import BaseState, _split_substate_key, _substate_key
 from reflex.utils import console
@@ -261,14 +262,17 @@ async def set_state(
         self,
         token: str,
         state: BaseState,
+        *,
         lock_id: bytes | None = None,
+        context: Event | None = None,
     ):
         """Set the state for a token.
 
         Args:
             token: The token to set the state for.
             state: The state to set.
             lock_id: If provided, the lock_key must be set to this value to set the state.
+            context: The event context.
 
         Raises:
             LockExpiredError: If lock_id is provided and the lock for the token is not held by that ID.
@@ -283,6 +287,7 @@ async def set_state(
                 f"Lock expired for token {token} while processing. Consider increasing "
                 f"`app.state_manager.lock_expiration` (currently {self.lock_expiration}) "
                 "or use `@rx.event(background=True)` decorator for long-running tasks."
+                + (f" Happened in event: {context.name}" if context is not None else "")
             )
             raise LockExpiredError(msg)
         if lock_id is not None:
@@ -292,7 +297,12 @@ async def set_state(
             if time_taken > self.lock_warning_threshold / 1000:
                 console.warn(
                     f"Lock for token {token} was held too long {time_taken=}s, "
-                    f"use `@rx.event(background=True)` decorator for long-running tasks.",
+                    f"use `@rx.event(background=True)` decorator for long-running tasks."
+                    + (
+                        f" Happened in event: {context.name}"
+                        if context is not None
+                        else ""
+                    ),
                     dedupe=True,
                 )
 
@@ -308,7 +318,8 @@ async def set_state(
                 self.set_state(
                     _substate_key(client_token, substate),
                     substate,
-                    lock_id,
+                    lock_id=lock_id,
+                    context=context,
                 ),
                 name=f"reflex_set_state|{client_token}|{substate.get_full_name()}",
             )
@@ -330,19 +341,22 @@ async def set_state(
 
     @override
     @contextlib.asynccontextmanager
-    async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
+    async def modify_state(
+        self, token: str, *, context: Event | None = None
+    ) -> AsyncIterator[BaseState]:
         """Modify the state for a token while holding exclusive lock.
 
         Args:
             token: The token to modify the state for.
+            context: The event context.
 
         Yields:
             The state for the token.
         """
         async with self._lock(token) as lock_id:
             state = await self.get_state(token)
             yield state
-            await self.set_state(token, state, lock_id)
+            await self.set_state(token, state, lock_id=lock_id, context=context)
 
     @staticmethod
     def _lock_key(token: str) -> bytes: