Only fetch required linked states (#6049)

Optimize the linked state loading path in redis mode by only linking states
which are already cached and augmenting `get_state` to link in dynamically
fetched states.

Adds SharedStateBaseInternal as an _always_dirty_substates from the root state.
This causes it to be always shallow-cached when loading any states from redis.
Therefore it will be available for use with `_modify_linked_states` without
explicitly getting it (and by consequence ALL SharedState instances). Only
required SharedState classes will be fetched now.

Author: masenf

reflex/istate/shared.py

@@ -338,10 +338,11 @@ async def _modify_linked_states(
                         linked_state_name
                     )
                 )
-                # TODO: Avoid always fetched linked states, it should be based on
-                # whether the state is accessed, however then `get_state` would need
-                # to know how to fetch in a linked state.
-                original_state = await self.get_state(linked_state_cls)
+                try:
+                    original_state = self._get_state_from_cache(linked_state_cls)
+                except ValueError:
+                    # This state wasn't required for processing the event.
+                    continue
                 linked_state = await original_state._internal_patch_linked_state(
                     linked_token
                 )
@@ -400,3 +401,9 @@ def __init_subclass__(cls, **kwargs):
         root_state = cls.get_root_state()
         if root_state.backend_vars["_reflex_internal_links"] is None:
             root_state.backend_vars["_reflex_internal_links"] = {}
+        if root_state is State:
+            # Always fetch SharedStateBaseInternal to access
+            # `_modify_linked_states` without having to use `.get_state()` which
+            # pulls in all linked states and substates which may not actually be
+            # accessed for this event.
+            root_state._always_dirty_substates.add(SharedStateBaseInternal.get_name())

reflex/state.py

@@ -2467,6 +2467,35 @@ class State(BaseState):
     # Maps the state full_name to an arbitrary token it is linked to for shared state.
     _reflex_internal_links: dict[str, str] | None = None
 
+    @_override_base_method
+    async def _get_state_from_redis(self, state_cls: type[T_STATE]) -> T_STATE:
+        """Get a state instance from redis with linking support.
+
+        Args:
+            state_cls: The class of the state.
+
+        Returns:
+            The instance of state_cls associated with this state's client_token.
+        """
+        state_instance = await super()._get_state_from_redis(state_cls)
+        if (
+            self._reflex_internal_links
+            and (
+                linked_token := self._reflex_internal_links.get(
+                    state_cls.get_full_name()
+                )
+            )
+            is not None
+            and (
+                internal_patch_linked_state := getattr(
+                    state_instance, "_internal_patch_linked_state", None
+                )
+            )
+            is not None
+        ):
+            return await internal_patch_linked_state(linked_token)
+        return state_instance
+
     @event
     def set_is_hydrated(self, value: bool) -> None:
         """Set the hydrated state.