Simplify StateManagerMemory expiration internals

Remove dead _token_last_touched dict, replace hand-rolled task
scheduling with ensure_task, move heap compaction off the hot path,
and fix touch ordering in get_state/set_state.

Author: FarhanAliRaza

reflex/istate/manager/_expiration.py

@@ -1,16 +1,19 @@
 """Internal helpers for in-memory state expiration."""
 
+from __future__ import annotations
+
 import asyncio
 import contextlib
 import dataclasses
 import heapq
 import time
-from typing import ClassVar
-
-from reflex.state import BaseState
+from typing import TYPE_CHECKING, ClassVar
 
 from . import _default_token_expiration
 
+if TYPE_CHECKING:
+    from reflex.state import BaseState
+
 
 @dataclasses.dataclass
 class StateManagerExpiration:
@@ -36,12 +39,6 @@ class StateManagerExpiration:
         init=False,
     )
 
-    # Last time a token was touched.
-    _token_last_touched: dict[str, float] = dataclasses.field(
-        default_factory=dict,
-        init=False,
-    )
-
     # Deadline-ordered token expiration heap.
     _token_expiration_heap: list[tuple[float, str]] = dataclasses.field(
         default_factory=list,
@@ -75,13 +72,10 @@ def _touch_token(self, token: str):
         Args:
             token: The token that was accessed.
         """
-        touched_at = time.time()
-        expires_at = touched_at + self.token_expiration
-        self._token_last_touched[token] = touched_at
+        expires_at = time.time() + self.token_expiration
         self._token_expires_at[token] = expires_at
         self._pending_locked_expirations.discard(token)
         heapq.heappush(self._token_expiration_heap, (expires_at, token))
-        self._maybe_compact_expiration_heap()
         if (
             self._scheduled_expiration_deadline is None
             or expires_at <= self._scheduled_expiration_deadline
@@ -124,7 +118,6 @@ def _purge_token(self, token: str):
         Args:
             token: The token to purge.
         """
-        self._token_last_touched.pop(token, None)
         self._token_expires_at.pop(token, None)
         self.states.pop(token, None)
         self._states_locks.pop(token, None)
@@ -133,20 +126,16 @@ def _purge_token(self, token: str):
     def _purge_expired_tokens(
         self,
         now: float | None = None,
-    ) -> list[str]:
+    ):
         """Purge expired in-memory state entries.
 
         If a token's state lock is currently held, defer cleanup until a later pass
         to avoid replacing the state while it is being modified.
 
         Args:
             now: The time to compare against.
-
-        Returns:
-            The list of purged tokens.
         """
         now = time.time() if now is None else now
-        expired_tokens = []
         while (
             next_expiration := self._next_expiration()
         ) is not None and next_expiration[0] <= now:
@@ -157,8 +146,7 @@ def _purge_expired_tokens(
                 self._pending_locked_expirations.add(token)
                 continue
             self._purge_token(token)
-            expired_tokens.append(token)
-        return expired_tokens
+        self._maybe_compact_expiration_heap()
 
     def _next_expiration_in(
         self,

reflex/istate/manager/memory.py

@@ -12,9 +12,7 @@
 from reflex.istate.manager import StateManager, StateModificationContext
 from reflex.istate.manager._expiration import StateManagerExpiration
 from reflex.state import BaseState, _split_substate_key
-from reflex.utils import console
-
-_EXPIRATION_ERROR_RETRY_SECONDS = 1.0
+from reflex.utils.tasks import ensure_task
 
 
 @dataclasses.dataclass
@@ -26,41 +24,24 @@ class StateManagerMemory(StateManagerExpiration, StateManager):
     # The mutex ensures the dict of mutexes is updated exclusively
     _state_manager_lock: asyncio.Lock = dataclasses.field(default=asyncio.Lock())
 
-    _expiration_task: asyncio.Task | None = None
+    _expiration_task: asyncio.Task | None = dataclasses.field(default=None, init=False)
 
     async def _expire_states_once(self):
         """Perform one expiration pass and wait for the next check."""
-        try:
-            now = time.time()
-            self._purge_expired_tokens(now=now)
-            await self._wait_for_token_activity(
-                self._prepare_expiration_wait(now=now),
-            )
-        except asyncio.CancelledError:
-            raise
-        except Exception as err:
-            console.error(f"Error expiring in-memory states: {err!r}")
-            await asyncio.sleep(_EXPIRATION_ERROR_RETRY_SECONDS)
-
-    async def _expire_states(self):
-        """Long running task that removes expired states from memory.
-
-        Raises:
-            asyncio.CancelledError: When the task is cancelled.
-        """
-        while True:
-            await self._expire_states_once()
-
-    async def _schedule_expiration_task(self):
-        """Schedule the expiration task if it is not already running."""
-        if self._expiration_task is None or self._expiration_task.done():
-            async with self._state_manager_lock:
-                if self._expiration_task is None or self._expiration_task.done():
-                    self._expiration_task = asyncio.create_task(
-                        self._expire_states(),
-                        name="StateManagerMemory|ExpirationProcessor",
-                    )
-                    await asyncio.sleep(0)
+        now = time.time()
+        self._purge_expired_tokens(now=now)
+        await self._wait_for_token_activity(
+            self._prepare_expiration_wait(now=now),
+        )
+
+    def _ensure_expiration_task(self):
+        """Ensure the expiration background task is running."""
+        ensure_task(
+            self,
+            "_expiration_task",
+            self._expire_states_once,
+            suppress_exceptions=[Exception],
+        )
 
     @override
     async def get_state(self, token: str) -> BaseState:
@@ -74,10 +55,10 @@ async def get_state(self, token: str) -> BaseState:
         """
         # Memory state manager ignores the substate suffix and always returns the top-level state.
         token = _split_substate_key(token)[0]
-        self._touch_token(token)
-        await self._schedule_expiration_task()
         if token not in self.states:
             self.states[token] = self.state(_reflex_internal_init=True)
+        self._touch_token(token)
+        self._ensure_expiration_task()
         return self.states[token]
 
     @override
@@ -95,9 +76,9 @@ async def set_state(
             context: The state modification context.
         """
         token = _split_substate_key(token)[0]
-        self._touch_token(token)
         self.states[token] = state
-        await self._schedule_expiration_task()
+        self._touch_token(token)
+        self._ensure_expiration_task()
 
     @override
     @contextlib.asynccontextmanager

tests/units/istate/manager/test_expiration.py

@@ -63,13 +63,13 @@ async def test_memory_state_manager_evicts_expired_state(
 
     assert token in state_manager_memory.states
     assert token in state_manager_memory._states_locks
-    assert token in state_manager_memory._token_last_touched
+    assert token in state_manager_memory._token_expires_at
 
     await _poll_until(
         lambda: (
             token not in state_manager_memory.states
             and token not in state_manager_memory._states_locks
-            and token not in state_manager_memory._token_last_touched
+            and token not in state_manager_memory._token_expires_at
         )
     )
 
@@ -84,13 +84,13 @@ async def test_memory_state_manager_get_state_refreshes_expiration(
     state = await state_manager_memory.get_state(state_token)
     assert isinstance(state, ExpiringState)
     state.value = 7
-    first_touch = state_manager_memory._token_last_touched[token]
+    first_expires_at = state_manager_memory._token_expires_at[token]
 
     await asyncio.sleep(0.6)
 
     same_state = await state_manager_memory.get_state(state_token)
     assert same_state is state
-    assert state_manager_memory._token_last_touched[token] > first_touch
+    assert state_manager_memory._token_expires_at[token] > first_expires_at
 
     await asyncio.sleep(0.6)
 
@@ -108,13 +108,13 @@ async def test_memory_state_manager_set_state_refreshes_expiration(
     state = await state_manager_memory.get_state(state_token)
     assert isinstance(state, ExpiringState)
     state.value = 17
-    first_touch = state_manager_memory._token_last_touched[token]
+    first_expires_at = state_manager_memory._token_expires_at[token]
 
     await asyncio.sleep(0.6)
 
     await state_manager_memory.set_state(state_token, state)
 
-    assert state_manager_memory._token_last_touched[token] > first_touch
+    assert state_manager_memory._token_expires_at[token] > first_expires_at
 
     await asyncio.sleep(0.6)