ENG-7948: StateManagerDisk deferred write queue

masenf · masenf · commit 96c0b451fb8f · 2025-10-14T13:45:09.000-07:00
* New env var: REFLEX_STATE_MANAGER_DISK_DEBOUNCE_SECONDS (default 2.0)
* If the debounce is non-zero, then state manager will queue the disk write
* Queued writes will be processed in order of set time after they exceed the debounce timeout
* New StateManager.close method standardized in base class
* Close app.state_manager when the server is going down
* Flush all queued writes when the StateManagerDisk closes
* Update test cases to always call `state_manager.close()`
diff --git a/reflex/app_mixins/lifespan.py b/reflex/app_mixins/lifespan.py
@@ -71,6 +71,13 @@ async def _run_lifespan_tasks(self, app: Starlette):
                     await event_namespace._token_manager.disconnect_all()
             except Exception as e:
                 console.error(f"Error during lifespan cleanup: {e}")
+        # Flush any pending writes from the state manager.
+        try:
+            state_manager = self.state_manager  # pyright: ignore[reportAttributeAccessIssue]
+        except AttributeError:
+            pass
+        else:
+            await state_manager.close()
 
     def register_lifespan_task(self, task: Callable | asyncio.Task, **task_kwargs):
         """Register a task to run during the lifespan of the app.
diff --git a/reflex/environment.py b/reflex/environment.py
@@ -95,6 +95,26 @@ def interpret_int_env(value: str, field_name: str) -> int:
         raise EnvironmentVarValueError(msg) from ve
 
 
+def interpret_float_env(value: str, field_name: str) -> float:
+    """Interpret a float environment variable value.
+
+    Args:
+        value: The environment variable value.
+        field_name: The field name.
+
+    Returns:
+        The interpreted value.
+
+    Raises:
+        EnvironmentVarValueError: If the value is invalid.
+    """
+    try:
+        return float(value)
+    except ValueError as ve:
+        msg = f"Invalid float value: {value!r} for {field_name}"
+        raise EnvironmentVarValueError(msg) from ve
+
+
 def interpret_existing_path_env(value: str, field_name: str) -> ExistingPath:
     """Interpret a path environment variable value as an existing path.
 
@@ -228,6 +248,8 @@ def interpret_env_var_value(
         return loglevel
     if field_type is int:
         return interpret_int_env(value, field_name)
+    if field_type is float:
+        return interpret_float_env(value, field_name)
     if field_type is Path:
         return interpret_path_env(value, field_name)
     if field_type is ExistingPath:
@@ -671,6 +693,9 @@ class EnvironmentVariables:
     # Whether to mount the compiled frontend app in the backend server in production.
     REFLEX_MOUNT_FRONTEND_COMPILED_APP: EnvVar[bool] = env_var(False, internal=True)
 
+    # How long to delay writing updated states to disk. (Higher values mean less writes, but more chance of lost data.)
+    REFLEX_STATE_MANAGER_DISK_DEBOUNCE_SECONDS: EnvVar[float] = env_var(2.0)
+
 
 environment = EnvironmentVariables()
 
diff --git a/reflex/istate/manager/__init__.py b/reflex/istate/manager/__init__.py
@@ -92,6 +92,9 @@ async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
         """
         yield self.state()
 
+    async def close(self):  # noqa: B027
+        """Close the state manager."""
+
 
 def _default_token_expiration() -> int:
     """Get the default token expiration time.
diff --git a/reflex/istate/manager/disk.py b/reflex/istate/manager/disk.py
@@ -4,15 +4,27 @@
 import contextlib
 import dataclasses
 import functools
+import time
 from collections.abc import AsyncIterator
 from hashlib import md5
 from pathlib import Path
 
 from typing_extensions import override
 
+from reflex.environment import environment
 from reflex.istate.manager import StateManager, _default_token_expiration
 from reflex.state import BaseState, _split_substate_key, _substate_key
-from reflex.utils import path_ops, prerequisites
+from reflex.utils import console, path_ops, prerequisites
+from reflex.utils.misc import run_in_thread
+
+
+@dataclasses.dataclass(frozen=True)
+class QueueItem:
+    """An item in the write queue."""
+
+    token: str
+    state: BaseState
+    timestamp: float
 
 
 @dataclasses.dataclass
@@ -34,6 +46,22 @@ class StateManagerDisk(StateManager):
     # The token expiration time (s).
     token_expiration: int = dataclasses.field(default_factory=_default_token_expiration)
 
+    # Last time a token was touched.
+    _token_last_touched: dict[str, float] = dataclasses.field(
+        default_factory=dict,
+        init=False,
+    )
+
+    # Pending writes
+    _write_queue: dict[str, QueueItem] = dataclasses.field(
+        default_factory=dict,
+        init=False,
+    )
+    _write_queue_task: asyncio.Task | None = None
+    _write_debounce_seconds: float = dataclasses.field(
+        default=environment.REFLEX_STATE_MANAGER_DISK_DEBOUNCE_SECONDS.get()
+    )
+
     def __post_init__(self):
         """Create a new state manager."""
         path_ops.mkdir(self.states_directory)
@@ -51,8 +79,6 @@ def states_directory(self) -> Path:
 
     def _purge_expired_states(self):
         """Purge expired states from the disk."""
-        import time
-
         for path in path_ops.ls(self.states_directory):
             # check path is a pickle file
             if path.suffix != ".pkl":
@@ -137,6 +163,7 @@ async def get_state(
             The state for the token.
         """
         client_token = _split_substate_key(token)[0]
+        self._token_last_touched[client_token] = time.time()
         root_state = self.states.get(client_token)
         if root_state is not None:
             # Retrieved state from memory.
@@ -170,11 +197,90 @@ async def set_state_for_substate(self, client_token: str, substate: BaseState):
             if pickle_state:
                 if not self.states_directory.exists():
                     self.states_directory.mkdir(parents=True, exist_ok=True)
-                self.token_path(substate_token).write_bytes(pickle_state)
+                await run_in_thread(
+                    lambda: self.token_path(substate_token).write_bytes(pickle_state),
+                )
 
         for substate_substate in substate.substates.values():
             await self.set_state_for_substate(client_token, substate_substate)
 
+    async def _process_write_queue_delay(self):
+        """Wait for the debounce period before processing the write queue again."""
+        if self._write_queue:
+            # There are still items in the queue, schedule another run.
+            now = time.time()
+            next_write_in = min(
+                self._write_debounce_seconds - (now - item.timestamp)
+                for item in self._write_queue.values()
+            )
+            await asyncio.sleep(next_write_in)
+        elif self._write_debounce_seconds > 0:
+            # No items left, wait a bit before checking again.
+            await asyncio.sleep(self._write_debounce_seconds)
+        else:
+            # No debounce, wait a minute before processing expirations.
+            await asyncio.sleep(60)
+
+    async def _process_write_queue(self):
+        """Long running task that checks for states to write to disk.
+
+        Raises:
+            asyncio.CancelledError: When the task is cancelled.
+        """
+        while True:
+            try:
+                now = time.time()
+                # sort the _write_queue by oldest timestamp and exclude items younger than debounce time
+                items_to_write = sorted(
+                    (
+                        item
+                        for item in self._write_queue.values()
+                        if now - item.timestamp >= self._write_debounce_seconds
+                    ),
+                    key=lambda item: item.timestamp,
+                )
+                for item in items_to_write:
+                    token = item.token
+                    client_token, _ = _split_substate_key(token)
+                    await self.set_state_for_substate(
+                        client_token, self._write_queue.pop(token).state
+                    )
+                # Check for expired states to purge.
+                for token, last_touched in list(self._token_last_touched.items()):
+                    if now - last_touched > self.token_expiration:
+                        self._token_last_touched.pop(token)
+                        self.states.pop(token, None)
+                await run_in_thread(self._purge_expired_states)
+                await self._process_write_queue_delay()
+            except asyncio.CancelledError:  # noqa: PERF203
+                n_outstanding_items = len(self._write_queue)
+                # When the task is cancelled, write all remaining items to disk.
+                console.debug(
+                    f"Closing StateManagerDisk, writing {n_outstanding_items} remaining items to disk"
+                )
+                for item in self._write_queue.values():
+                    token = item.token
+                    client_token, _ = _split_substate_key(token)
+                    await self.set_state_for_substate(
+                        client_token,
+                        item.state,
+                    )
+                console.debug(f"Finished writing {n_outstanding_items} items to disk")
+                raise
+            except Exception as e:
+                console.error(f"Error processing write queue: {e!r}")
+                await self._process_write_queue_delay()
+
+    async def _schedule_process_write_queue(self):
+        """Schedule the write queue processing task if not already running."""
+        if self._write_queue_task is None or self._write_queue_task.done():
+            async with self._state_manager_lock:
+                if self._write_queue_task is None or self._write_queue_task.done():
+                    self._write_queue_task = asyncio.create_task(
+                        self._process_write_queue()
+                    )
+                    await asyncio.sleep(0)  # Yield to allow the task to start.
+
     @override
     async def set_state(self, token: str, state: BaseState):
         """Set the state for a token.
@@ -184,7 +290,19 @@ async def set_state(self, token: str, state: BaseState):
             state: The state to set.
         """
         client_token, _ = _split_substate_key(token)
-        await self.set_state_for_substate(client_token, state)
+        if self._write_debounce_seconds > 0:
+            # Deferred write to reduce disk IO overhead.
+            if client_token not in self._write_queue:
+                self._write_queue[client_token] = QueueItem(
+                    token=client_token,
+                    state=state,
+                    timestamp=time.time(),
+                )
+        else:
+            # Immediate write to disk.
+            await self.set_state_for_substate(client_token, state)
+        # Ensure the processing task is scheduled to handle expirations and any deferred writes.
+        await self._schedule_process_write_queue()
 
     @override
     @contextlib.asynccontextmanager
@@ -208,3 +326,11 @@ async def modify_state(self, token: str) -> AsyncIterator[BaseState]:
             state = await self.get_state(token)
             yield state
             await self.set_state(token, state)
+
+    async def close(self):
+        """Close the state manager, flushing any pending writes to disk."""
+        if self._write_queue_task:
+            self._write_queue_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._write_queue_task
+            self._write_queue_task = None
diff --git a/tests/units/test_app.py b/tests/units/test_app.py
@@ -451,8 +451,7 @@ async def test_initialize_with_state(test_state: type[ATestState], token: str):
     assert isinstance(state, test_state)
     assert state.var == 0
 
-    if isinstance(app.state_manager, StateManagerRedis):
-        await app.state_manager.close()
+    await app.state_manager.close()
 
 
 @pytest.mark.asyncio
@@ -486,8 +485,7 @@ async def test_set_and_get_state(test_state):
     assert state1.var == 1
     assert state2.var == 2
 
-    if isinstance(app.state_manager, StateManagerRedis):
-        await app.state_manager.close()
+    await app.state_manager.close()
 
 
 @pytest.mark.asyncio
@@ -999,8 +997,7 @@ def getlist(key: str):
         "image2.jpg",
     ]
 
-    if isinstance(app.state_manager, StateManagerRedis):
-        await app.state_manager.close()
+    await app.state_manager.close()
 
 
 @pytest.mark.asyncio
@@ -1046,8 +1043,7 @@ def getlist(key: str):
         == f"`{state.get_full_name()}.handle_upload2` handler should have a parameter annotated as list[rx.UploadFile]"
     )
 
-    if isinstance(app.state_manager, StateManagerRedis):
-        await app.state_manager.close()
+    await app.state_manager.close()
 
 
 @pytest.mark.asyncio
@@ -1093,8 +1089,7 @@ def getlist(key: str):
         == f"@rx.event(background=True) is not supported for upload handler `{state.get_full_name()}.bg_upload`."
     )
 
-    if isinstance(app.state_manager, StateManagerRedis):
-        await app.state_manager.close()
+    await app.state_manager.close()
 
 
 class DynamicState(BaseState):
@@ -1372,8 +1367,7 @@ def _dynamic_state_event(name, val, **kwargs):
     assert state.loaded == len(exp_vals)
     assert state.counter == len(exp_vals)
 
-    if isinstance(app.state_manager, StateManagerRedis):
-        await app.state_manager.close()
+    await app.state_manager.close()
 
 
 @pytest.mark.asyncio
@@ -1412,8 +1406,7 @@ async def test_process_events(mocker: MockerFixture, token: str):
     assert (await app.state_manager.get_state(event.substate_token)).value == 5
     assert app._postprocess.call_count == 6  # pyright: ignore [reportAttributeAccessIssue]
 
-    if isinstance(app.state_manager, StateManagerRedis):
-        await app.state_manager.close()
+    await app.state_manager.close()
 
 
 @pytest.mark.parametrize(
diff --git a/tests/units/test_state.py b/tests/units/test_state.py
@@ -1685,8 +1685,7 @@ async def state_manager(request) -> AsyncGenerator[StateManager, None]:
 
     yield state_manager
 
-    if isinstance(state_manager, StateManagerRedis):
-        await state_manager.close()
+    await state_manager.close()
 
 
 @pytest.fixture
@@ -1737,6 +1736,8 @@ async def test_state_manager_modify_state(
         if state_manager._states_locks:
             assert sm2._states_locks != state_manager._states_locks
 
+        await sm2.close()
+
 
 @pytest.mark.asyncio
 async def test_state_manager_contend(
@@ -2960,8 +2961,7 @@ def index():
     async for update in state._process(events[1]):
         assert update.delta == exp_is_hydrated(state)
 
-    if isinstance(app.state_manager, StateManagerRedis):
-        await app.state_manager.close()
+    await app.state_manager.close()
 
 
 @pytest.mark.asyncio
@@ -3016,8 +3016,7 @@ def index():
     async for update in state._process(events[2]):
         assert update.delta == exp_is_hydrated(state)
 
-    if isinstance(app.state_manager, StateManagerRedis):
-        await app.state_manager.close()
+    await app.state_manager.close()
 
 
 @pytest.mark.asyncio
@@ -3658,13 +3657,15 @@ class Child(State):
         c = await root.get_state(Child)
         assert s._get_was_touched()
         assert not c._get_was_touched()
+    await dsm.close()
 
     dsm2 = StateManagerDisk(state=Root)
     root = await dsm2.get_state(token)
     s = await root.get_state(State)
     assert s.num == 43
     c = await root.get_state(Child)
     assert c.foo == "bar"
+    await dsm2.close()
 
 
 class Obj(Base):
