diff --git a/news/6740.performance.md b/news/6740.performance.md new file mode 100644 index 00000000000..3512eb02e06 --- /dev/null +++ b/news/6740.performance.md @@ -0,0 +1 @@ +Speed up MutableProxy: immutable elements retrieved through a proxy skip the dataclasses frame-walk check, and the proxy for each mutable state var is cached per instance instead of rebuilt on every attribute read. diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py index 776a519fe2c..7622e24694b 100644 --- a/reflex/istate/proxy.py +++ b/reflex/istate/proxy.py @@ -522,22 +522,24 @@ def _wrap_recursive(self, value: Any) -> Any: Returns: The wrapped value. """ - # When called from dataclasses internal code, return the unwrapped value - if self._is_called_from_dataclasses_internal(): - return value # If we already have a proxy, unwrap and rewrap to make sure the state # reference is up to date. if isinstance(value, MutableProxy): value = value.__wrapped__ + # Immutable values (the common case when iterating a container of + # scalars) never need wrapping nor the frame inspection below. + if not is_mutable_type(type(value)): + return value + # When called from dataclasses internal code, return the unwrapped value + if self._is_called_from_dataclasses_internal(): + return value # Recursively wrap mutable types. - if is_mutable_type(type(value)): - base_cls = globals()[self.__base_proxy__] - return base_cls( - wrapped=value, - state=self._self_state, - field_name=self._self_field_name, - ) - return value + base_cls = globals()[self.__base_proxy__] + return base_cls( + wrapped=value, + state=self._self_state, + field_name=self._self_field_name, + ) def _wrap_recursive_decorator( self, wrapped: Callable, instance: BaseState, args: list, kwargs: dict diff --git a/reflex/state.py b/reflex/state.py index 3762ad0ac2a..bed2085a1e6 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -1478,7 +1478,17 @@ def _get_attribute(self, name: str) -> Any: name in super().__getattribute__("base_vars") or name in backend_vars ): # track changes in mutable containers (list, dict, set, etc) - return MutableProxy(wrapped=value, state=self, field_name=name) + cache = super().__getattribute__("__dict__").get("_mutable_proxy_cache") + if cache is None: + cache = {} + object.__setattr__(self, "_mutable_proxy_cache", cache) + proxy = cache.get(name) + # isinstance also rejects entries degraded by deepcopy, which + # copies a MutableProxy as its unwrapped value. + if not isinstance(proxy, MutableProxy) or proxy.__wrapped__ is not value: + proxy = MutableProxy(wrapped=value, state=self, field_name=name) + cache[name] = proxy + return proxy return value @@ -1513,6 +1523,9 @@ def __setattr__(self, name: str, value: Any): if name in self.backend_vars: self._backend_vars.__setitem__(name, value) + if (cache := self.__dict__.get("_mutable_proxy_cache")) is not None: + # Drop the proxy wrapping the replaced value. + cache.pop(name, None) self.dirty_vars.add(name) self._mark_dirty() return @@ -1544,6 +1557,10 @@ def __setattr__(self, name: str, value: Any): # Set the attribute. object.__setattr__(self, name, value) + if (cache := self.__dict__.get("_mutable_proxy_cache")) is not None: + # Drop the proxy wrapping the replaced value. + cache.pop(name, None) + # Add the var to the dirty list. if name in self.base_vars: self.dirty_vars.add(name) @@ -2069,6 +2086,8 @@ def __getstate__(self): state.pop("parent_state", None) state.pop("substates", None) state.pop("_was_touched", None) + # Proxies wrap live state references and are rebuilt on access. + state.pop("_mutable_proxy_cache", None) # Remove all inherited vars. for inherited_var_name in self.inherited_vars: state.pop(inherited_var_name, None) diff --git a/tests/units/istate/test_proxy.py b/tests/units/istate/test_proxy.py index 2b909720147..9d13f736f17 100644 --- a/tests/units/istate/test_proxy.py +++ b/tests/units/istate/test_proxy.py @@ -71,3 +71,50 @@ async def mock_modify_state_context(*args, **kwargs): # After the exception, we should be able to enter the context again without issues async with state_proxy: pass + + +def test_mutable_proxy_cached_per_field(): + """Repeated reads of a mutable var reuse the proxy until reassignment.""" + state = ProxyTestState() + first = state.items + assert isinstance(first, MutableProxy) + assert state.items is first + # Reassignment invalidates the cached proxy. + state.items = [Item(2)] + second = state.items + assert isinstance(second, MutableProxy) + assert second is not first + assert second[0].id == 2 + # In-place mutation keeps the same wrapped object, so the proxy is reused. + second.append(Item(3)) + assert state.items is second + # Reassignment immediately evicts the cache entry, so no strong reference + # to the replaced value lingers until the next read. + state.items = [Item(4)] + assert "items" not in state.__dict__["_mutable_proxy_cache"] + + +def test_mutable_proxy_cache_not_serialized(): + """The per-instance proxy cache never leaks into pickles or copies.""" + state = ProxyTestState() + state.items.append(Item(1)) # populate the proxy cache + assert "_mutable_proxy_cache" in state.__dict__ + assert "_mutable_proxy_cache" not in state.__getstate__() + + restored = pickle.loads(pickle.dumps(state)) + assert "_mutable_proxy_cache" not in restored.__dict__ + restored_items = restored.items + assert isinstance(restored_items, MutableProxy) + # The restored proxy tracks the restored state, not the original. + assert restored_items._self_state is restored + + +def test_mutable_proxy_iteration_yields_plain_immutables(): + """Iterating a proxied container returns immutable elements unwrapped.""" + state = ProxyTestState() + state.items = [Item(1), Item(2)] + numbers = [item.id for item in state.items] + assert numbers == [1, 2] + assert all(type(n) is int for n in numbers) + # Mutable elements remain wrapped so nested mutations mark the state dirty. + assert all(isinstance(item, MutableProxy) for item in state.items) diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 076024c0363..c0a1a892a24 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -3243,11 +3243,11 @@ class BaseFieldSetterState(BaseState): bfss.dirty_vars.clear() assert "c1" not in bfss.dirty_vars - # Assert identity of MutableProxy + # Repeated reads reuse the cached MutableProxy for the same field. mp = bfss.c1 assert isinstance(mp, MutableProxy) mp3 = bfss.c1 - assert mp is not mp3 + assert mp is mp3 # Since none of these set calls had values, the state should not be dirty assert not bfss.dirty_vars