diff --git a/news/+idle-disk-state-cleanup.performance.md b/news/+idle-disk-state-cleanup.performance.md new file mode 100644 index 00000000000..e56f498d765 --- /dev/null +++ b/news/+idle-disk-state-cleanup.performance.md @@ -0,0 +1 @@ +Avoid rescanning unexpired disk state files on every idle write-queue tick. diff --git a/reflex/istate/manager/disk.py b/reflex/istate/manager/disk.py index b1da40ae209..529a0f325cd 100644 --- a/reflex/istate/manager/disk.py +++ b/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() diff --git a/tests/benchmarks/test_disk_expiration.py b/tests/benchmarks/test_disk_expiration.py new file mode 100644 index 00000000000..7dcb1cb5291 --- /dev/null +++ b/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() diff --git a/tests/units/istate/manager/test_disk.py b/tests/units/istate/manager/test_disk.py index ac5f9fd29d6..e1bd97d5a2c 100644 --- a/tests/units/istate/manager/test_disk.py +++ b/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()