Class StateManagerMemory

reflex/istate/manager/memory.py:21–221

A state manager that stores states in memory.

@dataclasses.dataclass
class StateManagerMemory(StateManager):
    """A state manager that stores states in memory."""

    # The token expiration time (s).
    token_expiration: int = dataclasses.field(default_factory=_default_token_expiration)

    # The mapping of client ids to states.
    states: dict[str, Any] = dataclasses.field(default_factory=dict)

    # The mutex ensures the dict of mutexes is updated exclusively
    _state_manager_lock: asyncio.Lock = dataclasses.field(default=asyncio.Lock())

    # The dict of mutexes for each client
    _states_locks: dict[str, asyncio.Lock] = dataclasses.field(
        default_factory=dict,
        init=False,
    )

    # The latest expiration deadline and token for each cache key.
    _token_expires_at: dict[str, tuple[float, StateToken]] = 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:
        """Get an existing state or create a fresh one for a token.

        Args:
            token: The normalized client token.

        Returns:
            The state for the token.
        """
        key = token.cache_key
        if key not in self.states:
            if isinstance(token, BaseStateToken):
                self.states[key] = token.cls.get_root_state()(
                    _reflex_internal_init=True
                )
            else:
                self.states[key] = token.cls()
        return cast(TOKEN_TYPE, self.states[key])

    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,
        )
        self._ensure_expiration_task()

    def _purge_token(self, token: StateToken):
        """Remove a token from in-memory state bookkeeping.

        Args:
            token: The token to purge.