diff --git a/news/+memory-expiration-index.performance.md b/news/+memory-expiration-index.performance.md new file mode 100644 index 00000000000..883021a1aae --- /dev/null +++ b/news/+memory-expiration-index.performance.md @@ -0,0 +1 @@ +Reduce in-memory session expiration overhead by checking the earliest deadlines instead of scanning every retained session. diff --git a/reflex/istate/manager/memory.py b/reflex/istate/manager/memory.py index 7f3d144ed39..ef629059d1c 100644 --- a/reflex/istate/manager/memory.py +++ b/reflex/istate/manager/memory.py @@ -3,6 +3,7 @@ import asyncio import contextlib import dataclasses +import heapq import time from collections.abc import AsyncIterator from typing import Any, cast @@ -42,6 +43,16 @@ class StateManagerMemory(StateManager): init=False, ) + # One heap entry per tracked token, updated in place on refresh. + _expiration_heap: list[tuple[float, str]] = dataclasses.field( + default_factory=list, + init=False, + ) + _expiration_indices: dict[str, int] = dataclasses.field( + default_factory=dict, + init=False, + ) + _expiration_task: asyncio.Task | None = dataclasses.field(default=None, init=False) def _get_or_create_state(self, token: StateToken[TOKEN_TYPE]) -> TOKEN_TYPE: @@ -65,21 +76,79 @@ def _get_or_create_state(self, token: StateToken[TOKEN_TYPE]) -> TOKEN_TYPE: def _track_token(self, token: StateToken): """Refresh the expiration deadline for an active token.""" - self._token_expires_at[token.cache_key] = ( - time.time() + self.token_expiration, - token, - ) + key = token.cache_key + expires_at = time.time() + self.token_expiration + self._token_expires_at[key] = (expires_at, token) + self._queue_expiration(key, expires_at) self._ensure_expiration_task() + def _queue_expiration(self, key: str, expires_at: float): + """Insert or update a token's deadline without retaining obsolete entries. + + Args: + key: The token's cache key. + expires_at: The expiration deadline. + """ + index = self._expiration_indices.get(key) + if index is None: + index = len(self._expiration_heap) + self._expiration_heap.append((expires_at, key)) + else: + self._expiration_heap[index] = (expires_at, key) + self._sift_expiration(index) + + def _remove_expiration(self, key: str): + """Remove a token's deadline from the index. + + Args: + key: The token's cache key. + """ + index = self._expiration_indices.pop(key, None) + if index is None: + return + last = self._expiration_heap.pop() + if index < len(self._expiration_heap): + self._expiration_heap[index] = last + self._sift_expiration(index) + + def _sift_expiration(self, index: int): + """Restore heap order and positions after a deadline changes. + + Args: + index: The changed entry's position. + """ + heap = self._expiration_heap + indices = self._expiration_indices + entry = heap[index] + while index > 0: + parent = (index - 1) // 2 + if heap[parent] <= entry: + break + heap[index] = heap[parent] + indices[heap[index][1]] = index + index = parent + while (child := 2 * index + 1) < len(heap): + if child + 1 < len(heap) and heap[child + 1] < heap[child]: + child += 1 + if entry <= heap[child]: + break + heap[index] = heap[child] + indices[heap[index][1]] = index + index = child + heap[index] = entry + indices[entry[1]] = index + def _purge_token(self, token: StateToken): """Remove a token from in-memory state bookkeeping. Args: token: The token to purge. """ - self._token_expires_at.pop(token.cache_key, None) + key = token.cache_key + self._remove_expiration(key) + self._token_expires_at.pop(key, None) self._states_locks.pop(token.lock_key, None) - self.states.pop(token.cache_key, None) + self.states.pop(key, None) def _purge_expired_tokens(self) -> float | None: """Purge expired in-memory state entries and return the next deadline. @@ -88,22 +157,71 @@ def _purge_expired_tokens(self) -> float | None: The next expiration deadline among unlocked tokens, if any. """ now = time.time() - next_expires_at = None - token_expires_at = self._token_expires_at - state_locks = self._states_locks - - for _cache_key, (expires_at, token) in list(token_expires_at.items()): - if ( - state_lock := state_locks.get(token.lock_key) - ) is not None and state_lock.locked(): - continue - if expires_at <= now: + count = len(self._expiration_heap) + if not count: + return None + # Bound heap repairs by the cost of a linear scan when many deadlines + # coincide or held locks must be skipped. + heap_budget = count // count.bit_length() + locked: list[tuple[float, str]] = [] + try: + while self._expiration_heap: + if not heap_budget: + # The scan rebuilds all remaining entries, including held + # tokens already removed from the heap by this pass. + locked.clear() + return self._purge_expired_tokens_scan(now) + expires_at, key = self._expiration_heap[0] + token = self._token_expires_at[key][1] + if ( + state_lock := self._states_locks.get(token.lock_key) + ) is not None and state_lock.locked(): + self._remove_expiration(key) + locked.append((expires_at, key)) + heap_budget -= 1 + continue + if expires_at > now: + return expires_at self._purge_token(token) - continue - if next_expires_at is None or expires_at < next_expires_at: - next_expires_at = expires_at + heap_budget -= 1 + return None + finally: + # Retain held tokens for the worker restarted after lock release. + for expires_at, key in locked: + self._queue_expiration(key, expires_at) + + def _purge_expired_tokens_scan(self, now: float) -> float | None: + """Expire a batch with a linear scan and rebuild the deadline index. + + Args: + now: The current time. - return next_expires_at + Returns: + The next expiration deadline among unlocked tokens, if any. + """ + self._expiration_heap.clear() + self._expiration_indices.clear() + next_expires_at = None + try: + for expires_at, token in list(self._token_expires_at.values()): + if ( + state_lock := self._states_locks.get(token.lock_key) + ) is not None and state_lock.locked(): + continue + if expires_at <= now: + self._purge_token(token) + elif next_expires_at is None or expires_at < next_expires_at: + next_expires_at = expires_at + return next_expires_at + finally: + self._expiration_heap.extend( + (expires_at, key) + for key, (expires_at, _) in self._token_expires_at.items() + ) + heapq.heapify(self._expiration_heap) + self._expiration_indices.update( + (key, index) for index, (_, key) in enumerate(self._expiration_heap) + ) async def _get_state_lock(self, token: StateToken) -> asyncio.Lock: """Get or create the lock for a token. diff --git a/tests/benchmarks/test_memory_expiration.py b/tests/benchmarks/test_memory_expiration.py new file mode 100644 index 00000000000..c465d99a075 --- /dev/null +++ b/tests/benchmarks/test_memory_expiration.py @@ -0,0 +1,115 @@ +"""Benchmarks for deadline lookup and refresh with many in-memory sessions.""" + +import asyncio +import itertools +from types import SimpleNamespace + +import pytest +from pytest_codspeed import BenchmarkFixture + +from reflex.istate.manager import memory +from reflex.istate.manager.memory import StateManagerMemory +from reflex.istate.manager.token import StateToken + + +@pytest.fixture(params=[1000, 10000]) +def expiration_manager( + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch +) -> tuple[StateManagerMemory, list[StateToken]]: + """Populate deadline bookkeeping without scheduling a background worker. + + Args: + request: The parametrized retained-session count. + monkeypatch: The pytest monkeypatch fixture. + + Returns: + The manager and its retained session tokens. + """ + manager = StateManagerMemory(token_expiration=3600) + monkeypatch.setattr(manager, "_ensure_expiration_task", lambda: None) + tokens = [StateToken(ident=str(i), cls=dict) for i in range(request.param)] + for token in tokens: + manager._track_token(token) + return manager, tokens + + +def test_next_memory_expiration( + expiration_manager: tuple[StateManagerMemory, list[StateToken]], + benchmark: BenchmarkFixture, +): + """Find the next deadline while no retained sessions have expired. + + Args: + expiration_manager: The manager and its retained session tokens. + benchmark: The CodSpeed benchmark fixture. + """ + manager, _ = expiration_manager + benchmark(manager._purge_expired_tokens) + + +def test_refresh_memory_expiration( + expiration_manager: tuple[StateManagerMemory, list[StateToken]], + benchmark: BenchmarkFixture, +): + """Refresh session deadlines in rotation to include index maintenance costs. + + Args: + expiration_manager: The manager and its retained session tokens. + benchmark: The CodSpeed benchmark fixture. + """ + manager, tokens = expiration_manager + cycle = itertools.cycle(tokens) + benchmark(lambda: manager._track_token(next(cycle))) + + +@pytest.mark.parametrize("count", [1000, 10000]) +@pytest.mark.parametrize( + "shape", ["all_due", "half_due", "all_held", "held_and_future"] +) +def test_memory_expiration_batches( + count: int, + shape: str, + monkeypatch: pytest.MonkeyPatch, + benchmark: BenchmarkFixture, +): + """Expire coalesced deadlines with setup excluded from measured work. + + Args: + count: The number of retained sessions. + shape: The distribution of due and held deadlines. + monkeypatch: The pytest monkeypatch fixture. + benchmark: The CodSpeed benchmark fixture. + """ + clock = SimpleNamespace(now=1000.0) + monkeypatch.setattr(memory, "time", SimpleNamespace(time=lambda: clock.now)) + + def setup(): + """Create a fresh manager before each timed expiration pass. + + Returns: + The manager argument and empty keyword arguments for the benchmark. + """ + manager = StateManagerMemory(token_expiration=10) + monkeypatch.setattr(manager, "_ensure_expiration_task", lambda: None) + for index in range(count): + future = shape in {"half_due", "held_and_future"} and index >= count // 2 + held = shape == "all_held" or ( + shape == "held_and_future" and index < count // 2 + ) + clock.now = 1020.0 if future else 1000.0 + token = StateToken(ident=str(index), cls=dict) + manager.states[token.cache_key] = None + manager._track_token(token) + lock = asyncio.Lock() + if held: + # Only locked() is read during expiration; no task owns these fixtures. + monkeypatch.setattr(lock, "_locked", True) + manager._states_locks[token.lock_key] = lock + clock.now = 1011.0 + return (manager,), {} + + benchmark.pedantic( + StateManagerMemory._purge_expired_tokens, + setup=setup, + rounds=5, + ) diff --git a/tests/units/istate/manager/test_expiration.py b/tests/units/istate/manager/test_expiration.py index d9a95227f21..68b31113ea5 100644 --- a/tests/units/istate/manager/test_expiration.py +++ b/tests/units/istate/manager/test_expiration.py @@ -1,14 +1,18 @@ """Tests for state manager token expiration.""" import asyncio +import random import time from collections.abc import AsyncGenerator, Callable +from types import SimpleNamespace +from unittest.mock import Mock import pytest import pytest_asyncio +from reflex.istate.manager import memory from reflex.istate.manager.memory import StateManagerMemory -from reflex.istate.manager.token import BaseStateToken +from reflex.istate.manager.token import BaseStateToken, StateToken from reflex.state import BaseState @@ -209,3 +213,231 @@ async def test_memory_state_manager_refreshes_expiration_after_locked_access( assert token in state_manager_memory.states await _poll_until(lambda: token not in state_manager_memory.states) + + +@pytest.mark.asyncio +async def test_memory_expiration_does_not_scan_future_tokens( + monkeypatch: pytest.MonkeyPatch, +): + """Finding the next deadline should not inspect every retained token. + + Args: + monkeypatch: The pytest monkeypatch fixture. + """ + manager = StateManagerMemory(token_expiration=3600) + monkeypatch.setattr(manager, "_ensure_expiration_task", lambda: None) + for i in range(1000): + await manager.get_state(StateToken(ident=str(i), cls=dict)) + locks = Mock(wraps=manager._states_locks) + monkeypatch.setattr(manager, "_states_locks", locks) + + assert manager._purge_expired_tokens() is not None + assert locks.get.call_count <= 1 + + +@pytest.fixture +def indexed_memory_manager( + monkeypatch: pytest.MonkeyPatch, +) -> tuple[StateManagerMemory, Mock]: + """Provide a manager whose deadlines are advanced without real sleeps. + + Args: + monkeypatch: The pytest monkeypatch fixture. + + Returns: + The manager and its controllable clock. + """ + clock = Mock(return_value=1000.0) + monkeypatch.setattr(memory, "time", SimpleNamespace(time=clock)) + manager = StateManagerMemory(token_expiration=10) + monkeypatch.setattr(manager, "_ensure_expiration_task", lambda: None) + return manager, clock + + +@pytest.mark.asyncio +async def test_memory_expiration_refresh_reorders_deadlines( + indexed_memory_manager: tuple[StateManagerMemory, Mock], +): + """Refreshes can move a deadline earlier or later without retaining old entries. + + Args: + indexed_memory_manager: The manager and its controllable clock. + """ + manager, clock = indexed_memory_manager + tokens = [StateToken(ident=str(i), cls=dict) for i in range(3)] + for i, token in enumerate(tokens): + clock.return_value = 1000.0 + i + await manager.get_state(token) + + clock.return_value = 1005.0 + await manager.get_state(tokens[0]) + manager.token_expiration = 1 + await manager.get_state(tokens[2]) + assert manager._purge_expired_tokens() == pytest.approx(1006.0) + assert len(manager._expiration_heap) == 3 + + for deadline, token, next_deadline in [ + (1006.0, tokens[2], 1011.0), + (1011.0, tokens[1], 1015.0), + (1015.0, tokens[0], None), + ]: + clock.return_value = deadline + assert manager._purge_expired_tokens() == next_deadline + assert token.cache_key not in manager.states + assert token.cache_key not in manager._expiration_indices + assert not manager._expiration_heap + assert not manager._token_expires_at + + +@pytest.mark.asyncio +async def test_memory_expiration_skips_and_restores_locked_tokens( + indexed_memory_manager: tuple[StateManagerMemory, Mock], +): + """Held tokens survive expired deadlines and remain indexed after skipped scans. + + Args: + indexed_memory_manager: The manager and its controllable clock. + """ + manager, clock = indexed_memory_manager + tokens = [StateToken(ident=str(i), cls=dict) for i in range(3)] + locks = [] + for i, token in enumerate(tokens): + clock.return_value = 1000.0 + i + await manager.get_state(token) + lock = await manager._get_state_lock(token) + await lock.acquire() + locks.append(lock) + + locks[2].release() + assert manager._purge_expired_tokens() == pytest.approx(1012.0) + assert len(manager._expiration_heap) == 3 + clock.return_value = 1020.0 + assert manager._purge_expired_tokens() is None + assert set(manager.states) == {token.cache_key for token in tokens[:2]} + assert len(manager._expiration_heap) == 2 + + for lock in locks[:2]: + lock.release() + assert manager._purge_expired_tokens() is None + assert not manager.states + assert not manager._states_locks + assert not manager._token_expires_at + assert not manager._expiration_heap + assert not manager._expiration_indices + + +@pytest.mark.asyncio +async def test_memory_expiration_index_stays_bounded_on_refresh( + indexed_memory_manager: tuple[StateManagerMemory, Mock], +): + """A frequently accessed session owns one deadline even at identical timestamps. + + Args: + indexed_memory_manager: The manager and its controllable clock. + """ + manager, clock = indexed_memory_manager + tokens = [StateToken(ident=str(i), cls=dict) for i in range(100)] + for token in tokens: + await manager.get_state(token) + for i in range(5000): + # Repeated equal times exercise refreshes within one clock tick. + clock.return_value = 1000.0 + i // 10 + await manager.get_state(tokens[i % len(tokens)]) + assert len(manager._expiration_heap) == len(tokens) + assert len(manager._expiration_indices) == len(tokens) + assert len(manager._token_expires_at) == len(tokens) + + +@pytest.mark.asyncio +async def test_memory_expiration_matches_scan_after_updates_and_removals( + indexed_memory_manager: tuple[StateManagerMemory, Mock], +): + """Indexed expiration matches a full-scan model across reordered deadlines. + + Args: + indexed_memory_manager: The manager and its controllable clock. + """ + manager, clock = indexed_memory_manager + rng = random.Random(0) + tokens = [StateToken(ident=str(i), cls=dict) for i in range(30)] + expected: dict[str, float] = {} + now = 1000.0 + for _ in range(500): + now += rng.randrange(3) + clock.return_value = now + manager.token_expiration = rng.randrange(20) + token = rng.choice(tokens) + await manager.get_state(token) + expected[token.cache_key] = now + manager.token_expiration + if rng.randrange(3) == 0: + expected = {key: expiry for key, expiry in expected.items() if expiry > now} + assert manager._purge_expired_tokens() == min( + expected.values(), default=None + ) + assert set(manager.states) == set(expected) + assert sorted(manager._expiration_heap) == sorted( + (expiry, key) for key, expiry in expected.items() + ) + assert manager._expiration_indices == { + key: index for index, (_, key) in enumerate(manager._expiration_heap) + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "shape", + ["bulk_due", "half_due", "all_held_due", "held_and_future", "all_held_future"], +) +async def test_memory_expiration_bounds_heap_work_for_batches( + indexed_memory_manager: tuple[StateManagerMemory, Mock], + monkeypatch: pytest.MonkeyPatch, + shape: str, +): + """Coalesced or held deadlines cannot require a heap repair per retained token. + + Args: + indexed_memory_manager: The manager and its controllable clock. + monkeypatch: The pytest monkeypatch fixture. + shape: The distribution of due and held deadlines. + """ + manager, clock = indexed_memory_manager + count = 256 + expected = set() + for i in range(count): + future = (shape in {"half_due", "held_and_future"} and i >= count // 2) or ( + shape == "all_held_future" + ) + held = shape in {"all_held_due", "all_held_future"} or ( + shape == "held_and_future" and i < count // 2 + ) + clock.return_value = 1000.0 + (20 if future else 0) + token = StateToken(ident=str(i), cls=dict) + await manager.get_state(token) + if held: + lock = await manager._get_state_lock(token) + await lock.acquire() + if future or held: + expected.add(token.cache_key) + clock.return_value = 1011.0 + repair = Mock(wraps=manager._sift_expiration) + monkeypatch.setattr(manager, "_sift_expiration", repair) + + deadline = manager._purge_expired_tokens() + + assert deadline == (1030 if shape in {"half_due", "held_and_future"} else None) + assert set(manager.states) == expected + assert repair.call_count <= count // count.bit_length() + assert len(manager._expiration_heap) == len(expected) + assert manager._expiration_indices == { + key: index for index, (_, key) in enumerate(manager._expiration_heap) + } + # A rebuilt heap must retain skipped held tokens for a subsequent unlock. + for lock in manager._states_locks.values(): + if lock.locked(): + lock.release() + clock.return_value = 1040.0 + assert manager._purge_expired_tokens() is None + assert not manager._expiration_heap + assert not manager._expiration_indices + assert not manager._token_expires_at + assert not manager.states