Skip idle disk-state scans until expiration

news/+idle-disk-state-cleanup.performance.md

@@ -0,0 +1 @@
+Avoid rescanning unexpired disk state files on every idle write-queue tick.

reflex/istate/manager/disk.py

@@ -71,11 +71,19 @@ class StateManagerDisk(StateManager):
         default=environment.REFLEX_STATE_MANAGER_DISK_DEBOUNCE_SECONDS.get()
     )
 
+    _next_disk_purge: float = dataclasses.field(default=0.0, init=False)
+    _disk_purge_token_expiration: int | None = dataclasses.field(
+        default=None, init=False
+    )
+    _disk_write_generation: int = dataclasses.field(default=0, init=False)
+
     def __post_init__(self):
         """Create a new state manager."""
         path_ops.mkdir(self.states_directory)
 
-        self._purge_expired_states()
+        self._disk_purge_token_expiration, self._next_disk_purge = (
+            self._purge_expired_states()
+        )
 
     @functools.cached_property
     def states_directory(self) -> Path:
@@ -89,8 +97,15 @@ def states_directory(self) -> Path:
         """
         return prerequisites.get_states_dir().absolute()
 
-    def _purge_expired_states(self):
-        """Purge expired states from the disk."""
+    def _purge_expired_states(self) -> tuple[int, float]:
+        """Purge expired files and find the next possible disk expiration.
+
+        Returns:
+            The lifetime used and next expiration deadline, capped at one token
+            lifetime from the scan's start so newly created files are covered too.
+        """
+        token_expiration = self.token_expiration
+        next_expiration = time.time() + token_expiration
         for path in path_ops.ls(self.states_directory):
             # check path is a pickle file
             if path.suffix != ".pkl":
@@ -100,9 +115,27 @@ def _purge_expired_states(self):
             last_edited = path.stat().st_mtime
 
             # check if the file is older than the token expiration time
-            if time.time() - last_edited > self.token_expiration:
+            if time.time() - last_edited > token_expiration:
                 # remove the file
                 path.unlink()
+            else:
+                next_expiration = min(next_expiration, last_edited + token_expiration)
+        return token_expiration, next_expiration
+
+    async def _maybe_purge_expired_states(self):
+        """Scan disk only when a file may expire or a write invalidates the scan."""
+        if (
+            time.time() < self._next_disk_purge
+            and self.token_expiration == self._disk_purge_token_expiration
+        ):
+            return
+        generation = self._disk_write_generation
+        token_expiration, next_expiration = await run_in_thread(
+            self._purge_expired_states
+        )
+        if generation == self._disk_write_generation:
+            self._next_disk_purge = next_expiration
+            self._disk_purge_token_expiration = token_expiration
 
     def token_path(self, token: StateToken) -> Path:
         """Get the path for a token.
@@ -225,6 +258,8 @@ async def set_state_for_substate(
                 await run_in_thread(
                     lambda: self.token_path(substate_token).write_bytes(pickle_state),
                 )
+                self._disk_write_generation += 1
+                self._disk_purge_token_expiration = None
 
         if isinstance(token, BaseStateToken) and isinstance(substate, BaseState):
             for substate_substate in substate.substates.values():
@@ -282,7 +317,7 @@ async def _process_write_queue(self):
                     if now - last_touched > self.token_expiration:
                         self._token_last_touched.pop(cache_key)
                         self.states.pop(cache_key, None)
-                await run_in_thread(self._purge_expired_states)
+                await self._maybe_purge_expired_states()
                 await self._process_write_queue_delay()
             except asyncio.CancelledError:  # noqa: PERF203
                 await self._flush_write_queue()

tests/benchmarks/test_disk_expiration.py

@@ -0,0 +1,36 @@
+"""Benchmarks for idle disk-state expiration checks."""
+
+import asyncio
+import time
+
+import pytest
+from pytest_codspeed import BenchmarkFixture
+
+from reflex.istate.manager import disk
+from reflex.istate.manager.disk import StateManagerDisk
+
+
+@pytest.mark.parametrize("files", [100, 10_000])
+def test_idle_disk_expiration_check(
+    benchmark: BenchmarkFixture, tmp_path, monkeypatch, files
+):
+    """Check an idle manager without scanning its unexpired files every time.
+
+    Args:
+        benchmark: The benchmark fixture.
+        tmp_path: The temporary state directory.
+        monkeypatch: The monkeypatch fixture.
+        files: The number of unexpired state files.
+    """
+    monkeypatch.setattr(disk.prerequisites, "get_states_dir", lambda: tmp_path)
+    for index in range(files):
+        (tmp_path / f"{index}.pkl").touch()
+    manager = StateManagerDisk(token_expiration=3600)
+    assert manager._next_disk_purge > time.time()
+    loop = asyncio.new_event_loop()
+    try:
+        benchmark(
+            lambda: loop.run_until_complete(manager._maybe_purge_expired_states())
+        )
+    finally:
+        loop.close()

tests/units/istate/manager/test_disk.py

@@ -1,9 +1,16 @@
 """Tests for the disk state manager."""
 
+import asyncio
 import os
 from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock
 
+import pytest
+
+from reflex.istate.manager import disk
 from reflex.istate.manager.disk import StateManagerDisk
+from reflex.istate.manager.token import StateToken
 
 
 def test_states_directory_survives_chdir(tmp_path: Path, monkeypatch):
@@ -25,3 +32,145 @@ def test_states_directory_survives_chdir(tmp_path: Path, monkeypatch):
     assert manager.states_directory == states_dir
     # Purge resolves against the original directory, not the new cwd.
     manager._purge_expired_states()
+
+
+async def test_idle_write_queue_scans_only_when_files_can_expire(tmp_path, monkeypatch):
+    """Idle queue ticks skip unexpired files but still purge at their deadline."""
+    now = [1000.0]
+    monkeypatch.setattr(disk, "time", SimpleNamespace(time=lambda: now[0]))
+    monkeypatch.setattr(disk.prerequisites, "get_states_dir", lambda: tmp_path)
+    state_file = tmp_path / "session.pkl"
+    state_file.touch()
+    os.utime(state_file, (950, 950))
+    manager = StateManagerDisk(token_expiration=100)
+    listing = Mock(wraps=disk.path_ops.ls)
+    monkeypatch.setattr(disk.path_ops, "ls", listing)
+    ticks = iter([1002, 1049, 1051])
+
+    async def advance():
+        await asyncio.sleep(0)
+        assert state_file.exists() == (now[0] <= 1050)
+        if (next_tick := next(ticks, None)) is None:
+            raise asyncio.CancelledError
+        now[0] = next_tick
+
+    monkeypatch.setattr(manager, "_process_write_queue_delay", advance)
+    with pytest.raises(asyncio.CancelledError):
+        await manager._process_write_queue()
+    assert listing.call_count == 1
+
+
+@pytest.fixture
+def disk_clock(tmp_path, monkeypatch):
+    """Control disk deadlines without changing the process-wide clock.
+
+    Returns:
+        A list containing the current disk-manager timestamp.
+    """
+    now = [1000.0]
+    monkeypatch.setattr(disk, "time", SimpleNamespace(time=lambda: now[0]))
+    monkeypatch.setattr(disk.prerequisites, "get_states_dir", lambda: tmp_path)
+    return now
+
+
+async def test_new_files_after_empty_scan_expire(tmp_path, monkeypatch, disk_clock):
+    """An empty scan cannot postpone discovery beyond one token lifetime."""
+    manager = StateManagerDisk(token_expiration=100)
+    state_file = tmp_path / "later.pkl"
+    state_file.touch()
+    os.utime(state_file, (1001, 1001))
+    listing = Mock(wraps=disk.path_ops.ls)
+    monkeypatch.setattr(disk.path_ops, "ls", listing)
+
+    disk_clock[0] = 1002
+    await manager._maybe_purge_expired_states()
+    listing.assert_not_called()
+    disk_clock[0] = 1100
+    await manager._maybe_purge_expired_states()
+    assert state_file.exists()
+    disk_clock[0] = 1102
+    await manager._maybe_purge_expired_states()
+    assert not state_file.exists()
+    assert listing.call_count == 2
+
+
+@pytest.mark.parametrize("lifetime", [5, 200])
+async def test_disk_expiration_lifetime_changes(tmp_path, disk_clock, lifetime):
+    """Changing the lifetime invalidates deadlines from the previous setting."""
+    state_file = tmp_path / "existing.pkl"
+    state_file.touch()
+    os.utime(state_file, (990, 990))
+    manager = StateManagerDisk(token_expiration=100)
+    manager.token_expiration = lifetime
+
+    await manager._maybe_purge_expired_states()
+    assert state_file.exists() == (lifetime == 200)
+    disk_clock[0] = 1191
+    await manager._maybe_purge_expired_states()
+    assert not state_file.exists()
+
+
+async def test_disk_write_invalidates_expiration_scan(tmp_path, disk_clock):
+    """A persisted file is discovered even if the clock moved back during writing."""
+    manager = StateManagerDisk(token_expiration=100)
+    token = StateToken(ident="written", cls=dict)
+    disk_clock[0] = 900
+    await manager.set_state_for_substate(token, {"count": 1})
+    state_file = manager.token_path(token)
+    os.utime(state_file, (900, 900))
+
+    disk_clock[0] = 1001
+    await manager._maybe_purge_expired_states()
+    assert not state_file.exists()
+
+
+async def test_disk_write_during_scan_keeps_invalidation(monkeypatch, disk_clock):
+    """A worker scan cannot replace invalidation from a concurrent file write."""
+    manager = StateManagerDisk(token_expiration=100)
+    token = StateToken(ident="concurrent", cls=dict)
+    state_file = manager.token_path(token)
+    disk_clock[0] = 1101
+    wrote = False
+
+    async def run_with_concurrent_write(fn):
+        nonlocal wrote
+        result = fn()
+        if fn == manager._purge_expired_states and not wrote:
+            wrote = True
+            disk_clock[0] = 1000
+            await manager.set_state_for_substate(token, {"count": 1})
+            os.utime(state_file, (1000, 1000))
+            disk_clock[0] = 1101
+        return result
+
+    monkeypatch.setattr(disk, "run_in_thread", run_with_concurrent_write)
+    await manager._maybe_purge_expired_states()
+    assert state_file.exists()
+    await manager._maybe_purge_expired_states()
+    assert not state_file.exists()
+
+
+async def test_disk_expiration_lifetime_changes_while_scan_dispatched(
+    tmp_path, monkeypatch, disk_clock
+):
+    """A deadline is associated with the lifetime actually used by the worker."""
+    state_file = tmp_path / "changed-lifetime.pkl"
+    state_file.touch()
+    os.utime(state_file, (990, 990))
+    manager = StateManagerDisk(token_expiration=100)
+    manager._next_disk_purge = 0
+
+    async def scan_with_temporary_lifetime(fn):
+        await asyncio.sleep(0)
+        manager.token_expiration = 200
+        result = fn()
+        manager.token_expiration = 100
+        return result
+
+    with monkeypatch.context() as patch:
+        patch.setattr(disk, "run_in_thread", scan_with_temporary_lifetime)
+        await manager._maybe_purge_expired_states()
+
+    disk_clock[0] = 1091
+    await manager._maybe_purge_expired_states()
+    assert not state_file.exists()