implement performance mode for existing state size check (#4392)

Author: benedikt-bartscher

reflex/config.py

@@ -454,6 +454,14 @@ class PathExistsFlag:
 ExistingPath = Annotated[Path, PathExistsFlag]
 
 
+class PerformanceMode(enum.Enum):
+    """Performance mode for the app."""
+
+    WARN = "warn"
+    RAISE = "raise"
+    OFF = "off"
+
+
 class EnvironmentVariables:
     """Environment variables class to instantiate environment variables."""
 
@@ -550,6 +558,12 @@ class EnvironmentVariables:
     # Whether to check for outdated package versions.
     REFLEX_CHECK_LATEST_VERSION: EnvVar[bool] = env_var(True)
 
+    # In which performance mode to run the app.
+    REFLEX_PERF_MODE: EnvVar[Optional[PerformanceMode]] = env_var(PerformanceMode.WARN)
+
+    # The maximum size of the reflex state in kilobytes.
+    REFLEX_STATE_SIZE_LIMIT: EnvVar[int] = env_var(1000)
+
 
 environment = EnvironmentVariables()
 

reflex/state.py

@@ -43,7 +43,7 @@
 from typing_extensions import Self
 
 from reflex import event
-from reflex.config import get_config
+from reflex.config import PerformanceMode, get_config
 from reflex.istate.data import RouterData
 from reflex.istate.storage import ClientStorageBase
 from reflex.model import Model
@@ -90,6 +90,7 @@
     ReflexRuntimeError,
     SetUndefinedStateVarError,
     StateSchemaMismatchError,
+    StateTooLargeError,
 )
 from reflex.utils.exec import is_testing_env
 from reflex.utils.serializers import serializer
@@ -110,10 +111,11 @@
 var = computed_var
 
 
-# If the state is this large, it's considered a performance issue.
-TOO_LARGE_SERIALIZED_STATE = 100 * 1024  # 100kb
-# Only warn about each state class size once.
-_WARNED_ABOUT_STATE_SIZE: Set[str] = set()
+if environment.REFLEX_PERF_MODE.get() != PerformanceMode.OFF:
+    # If the state is this large, it's considered a performance issue.
+    TOO_LARGE_SERIALIZED_STATE = environment.REFLEX_STATE_SIZE_LIMIT.get() * 1024
+    # Only warn about each state class size once.
+    _WARNED_ABOUT_STATE_SIZE: Set[str] = set()
 
 # Errors caught during pickling of state
 HANDLED_PICKLE_ERRORS = (
@@ -2097,25 +2099,32 @@ def __getstate__(self):
             state["__dict__"].pop(inherited_var_name, None)
         return state
 
-    def _warn_if_too_large(
+    def _check_state_size(
         self,
         pickle_state_size: int,
     ):
         """Print a warning when the state is too large.
 
         Args:
             pickle_state_size: The size of the pickled state.
+
+        Raises:
+            StateTooLargeError: If the state is too large.
         """
         state_full_name = self.get_full_name()
         if (
             state_full_name not in _WARNED_ABOUT_STATE_SIZE
             and pickle_state_size > TOO_LARGE_SERIALIZED_STATE
             and self.substates
         ):
-            console.warn(
+            msg = (
                 f"State {state_full_name} serializes to {pickle_state_size} bytes "
-                "which may present performance issues. Consider reducing the size of this state."
+                + "which may present performance issues. Consider reducing the size of this state."
             )
+            if environment.REFLEX_PERF_MODE.get() == PerformanceMode.WARN:
+                console.warn(msg)
+            elif environment.REFLEX_PERF_MODE.get() == PerformanceMode.RAISE:
+                raise StateTooLargeError(msg)
             _WARNED_ABOUT_STATE_SIZE.add(state_full_name)
 
     @classmethod
@@ -2157,7 +2166,8 @@ def _serialize(self) -> bytes:
         """
         try:
             pickle_state = pickle.dumps((self._to_schema(), self))
-            self._warn_if_too_large(len(pickle_state))
+            if environment.REFLEX_PERF_MODE.get() != PerformanceMode.OFF:
+                self._check_state_size(len(pickle_state))
             return pickle_state
         except HANDLED_PICKLE_ERRORS as og_pickle_error:
             error = (

reflex/utils/exceptions.py

@@ -151,6 +151,10 @@ class InvalidPropValueError(ReflexError):
     """Raised when a prop value is invalid."""
 
 
+class StateTooLargeError(ReflexError):
+    """Raised when the state is too large to be serialized."""
+
+
 class SystemPackageMissingError(ReflexError):
     """Raised when a system package is missing."""