Browse Source

[REF-1982] state: Warn if redis state is "too big" (#2868)

If the state serializes to over 100kb and has substates, then print a warning
suggesting the developer reduce the size of the state.
Masen Furer 1 year ago
parent
commit
58f706ac7a
1 changed files with 33 additions and 1 deletions
  1. 33 1
      reflex/state.py

+ 33 - 1
reflex/state.py

@@ -52,6 +52,10 @@ Delta = Dict[str, Any]
 var = computed_var
 
 
+# If the state is this large, it's considered a performance issue.
+TOO_LARGE_SERIALIZED_STATE = 100 * 1024  # 100kb
+
+
 class HeaderData(Base):
     """An object containing headers data."""
 
@@ -2183,6 +2187,9 @@ class StateManagerRedis(StateManager):
         b"evicted",
     }
 
+    # Only warn about each state class size once.
+    _warned_about_state_size: ClassVar[Set[str]] = set()
+
     def _get_root_state(self, state: BaseState) -> BaseState:
         """Chase parent_state pointers to find an instance of the top-level state.
 
@@ -2334,6 +2341,29 @@ class StateManagerRedis(StateManager):
             return self._get_root_state(state)
         return state
 
+    def _warn_if_too_large(
+        self,
+        state: BaseState,
+        pickle_state_size: int,
+    ):
+        """Print a warning when the state is too large.
+
+        Args:
+            state: The state to check.
+            pickle_state_size: The size of the pickled state.
+        """
+        state_full_name = state.get_full_name()
+        if (
+            state_full_name not in self._warned_about_state_size
+            and pickle_state_size > TOO_LARGE_SERIALIZED_STATE
+            and state.substates
+        ):
+            console.warn(
+                f"State {state_full_name} serializes to {pickle_state_size} bytes "
+                "which may present performance issues. Consider reducing the size of this state."
+            )
+            self._warned_about_state_size.add(state_full_name)
+
     async def set_state(
         self,
         token: str,
@@ -2382,9 +2412,11 @@ class StateManagerRedis(StateManager):
             )
         # Persist only the given state (parents or substates are excluded by BaseState.__getstate__).
         if state._get_was_touched():
+            pickle_state = cloudpickle.dumps(state)
+            self._warn_if_too_large(state, len(pickle_state))
             await self.redis.set(
                 _substate_key(client_token, state),
-                cloudpickle.dumps(state),
+                pickle_state,
                 ex=self.token_expiration,
             )