add context for lock expiry (#5918)

* add context for lock expiry

* unpack

---------

Co-authored-by: Masen Furer <[email protected]>

Author: adhami3310

README.md

@@ -98,7 +98,6 @@ Now you can modify the source code in `my_app_name/my_app_name.py`. Reflex has f
 
 If you installed Reflex without a virtual environment and the `reflex` command is not found, you can run commands using: `python3 -m reflex init` and `python3 -m reflex run`
 
-
 ## 🫧 Example App
 
 Let's go over an example: creating an image generation UI around [DALL·E](https://platform.openai.com/docs/guides/images/image-generation?context=node). For simplicity, we just call the [OpenAI API](https://platform.openai.com/docs/api-reference/authentication), but you could replace this with an ML model run locally.

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, event=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

@@ -4,14 +4,27 @@
 import dataclasses
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator
+from typing import TypedDict
+
+from typing_extensions import ReadOnly, Unpack
 
 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
 
 
+class StateModificationContext(TypedDict, total=False):
+    """The context for modifying state."""
+
+    event: ReadOnly[Event]
+
+
+EmptyContext = StateModificationContext()
+
+
 @dataclasses.dataclass
 class StateManager(ABC):
     """A class to manage many client states."""
@@ -71,21 +84,30 @@ 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: Unpack[StateModificationContext],
+    ):
         """Set the state for a token.
 
         Args:
             token: The token to set the state for.
             state: The state to set.
+            context: The state modification context.
         """
 
     @abstractmethod
     @contextlib.asynccontextmanager
-    async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
+    async def modify_state(
+        self, token: str, **context: Unpack[StateModificationContext]
+    ) -> AsyncIterator[BaseState]:
         """Modify the state for a token while holding exclusive lock.
 
         Args:
             token: The token to modify the state for.
+            context: The state modification context.
 
         Yields:
             The state for the token.

reflex/istate/manager/disk.py

@@ -9,10 +9,14 @@
 from hashlib import md5
 from pathlib import Path
 
-from typing_extensions import override
+from typing_extensions import Unpack, override
 
 from reflex.environment import environment
-from reflex.istate.manager import StateManager, _default_token_expiration
+from reflex.istate.manager import (
+    StateManager,
+    StateModificationContext,
+    _default_token_expiration,
+)
 from reflex.state import BaseState, _split_substate_key, _substate_key
 from reflex.utils import console, path_ops, prerequisites
 from reflex.utils.misc import run_in_thread
@@ -301,12 +305,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: Unpack[StateModificationContext]
+    ):
         """Set the state for a token.
 
         Args:
             token: The token to set the state for.
             state: The state to set.
+            context: The state modification context.
         """
         client_token, _ = _split_substate_key(token)
         if self._write_debounce_seconds > 0:
@@ -325,11 +332,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: Unpack[StateModificationContext]
+    ) -> AsyncIterator[BaseState]:
         """Modify the state for a token while holding exclusive lock.
 
         Args:
             token: The token to modify the state for.
+            context: The state modification context.
 
         Yields:
             The state for the token.
@@ -344,7 +354,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)
 
     async def close(self):
         """Close the state manager, flushing any pending writes to disk."""

reflex/istate/manager/memory.py

@@ -5,9 +5,9 @@
 import dataclasses
 from collections.abc import AsyncIterator
 
-from typing_extensions import override
+from typing_extensions import Unpack, override
 
-from reflex.istate.manager import StateManager
+from reflex.istate.manager import StateManager, StateModificationContext
 from reflex.state import BaseState, _split_substate_key
 
 
@@ -43,23 +43,32 @@ 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: Unpack[StateModificationContext],
+    ):
         """Set the state for a token.
 
         Args:
             token: The token to set the state for.
             state: The state to set.
+            context: The state modification 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: Unpack[StateModificationContext]
+    ) -> AsyncIterator[BaseState]:
         """Modify the state for a token while holding exclusive lock.
 
         Args:
             token: The token to modify the state for.
+            context: The state modification context.
 
         Yields:
             The state for the token.

reflex/istate/manager/redis.py

@@ -10,11 +10,15 @@
 from redis import ResponseError
 from redis.asyncio import Redis
 from redis.asyncio.client import PubSub
-from typing_extensions import override
+from typing_extensions import Unpack, override
 
 from reflex.config import get_config
 from reflex.environment import environment
-from reflex.istate.manager import StateManager, _default_token_expiration
+from reflex.istate.manager import (
+    StateManager,
+    StateModificationContext,
+    _default_token_expiration,
+)
 from reflex.state import BaseState, _split_substate_key, _substate_key
 from reflex.utils import console
 from reflex.utils.exceptions import (
@@ -261,14 +265,17 @@ async def set_state(
         self,
         token: str,
         state: BaseState,
+        *,
         lock_id: bytes | None = None,
+        **context: Unpack[StateModificationContext],
     ):
         """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 +290,11 @@ 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: {event.name}"
+                    if (event := context.get("event")) is not None
+                    else ""
+                )
             )
             raise LockExpiredError(msg)
         if lock_id is not None:
@@ -292,7 +304,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: {event.name}"
+                        if (event := context.get("event")) is not None
+                        else ""
+                    ),
                     dedupe=True,
                 )
 
@@ -308,7 +325,8 @@ async def set_state(
                 self.set_state(
                     _substate_key(client_token, substate),
                     substate,
-                    lock_id,
+                    lock_id=lock_id,
+                    **context,
                 ),
                 name=f"reflex_set_state|{client_token}|{substate.get_full_name()}",
             )
@@ -330,19 +348,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: Unpack[StateModificationContext]
+    ) -> AsyncIterator[BaseState]:
         """Modify the state for a token while holding exclusive lock.
 
         Args:
             token: The token to modify the state for.
+            context: The state modification 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)
 
     @staticmethod
     def _lock_key(token: str) -> bytes: