Index in-memory session expiration deadlines

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.

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.

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,
+    )