From 09da1cb6e21901e852ac5b0eaaa9dc0a135100de Mon Sep 17 00:00:00 2001
From: Benedikt Bartscher <bb@senbax.de>
Date: Mon, 18 Nov 2024 22:43:37 +0100
Subject: [PATCH] implement performance mode for existing state size check

---
 reflex/config.py           | 14 ++++++++++++++
 reflex/state.py            | 28 +++++++++++++++++++---------
 reflex/utils/exceptions.py |  4 ++++
 3 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/reflex/config.py b/reflex/config.py
index 049cc2e83..2495a434c 100644
--- a/reflex/config.py
+++ b/reflex/config.py
@@ -452,6 +452,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."""
 
@@ -545,6 +553,12 @@ class EnvironmentVariables:
     # Where to save screenshots when tests fail.
     SCREENSHOT_DIR: EnvVar[Optional[Path]] = env_var(None)
 
+    # 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()
 
diff --git a/reflex/state.py b/reflex/state.py
index 719ff43b3..54f0d2c78 100644
--- a/reflex/state.py
+++ b/reflex/state.py
@@ -43,7 +43,7 @@ from sqlalchemy.orm import DeclarativeBase
 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
@@ -89,6 +89,7 @@ from reflex.utils.exceptions import (
     LockExpiredError,
     SetUndefinedStateVarError,
     StateSchemaMismatchError,
+    StateTooLargeError,
 )
 from reflex.utils.exec import is_testing_env
 from reflex.utils.serializers import serializer
@@ -109,10 +110,11 @@ 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
-# 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 = (
@@ -2092,7 +2094,7 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
             state["__dict__"].pop(inherited_var_name, None)
         return state
 
-    def _warn_if_too_large(
+    def _check_state_size(
         self,
         pickle_state_size: int,
     ):
@@ -2100,6 +2102,9 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
 
         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 (
@@ -2107,10 +2112,14 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
             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
@@ -2152,7 +2161,8 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
         """
         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 = (
diff --git a/reflex/utils/exceptions.py b/reflex/utils/exceptions.py
index 7611e5a4d..714dc912c 100644
--- a/reflex/utils/exceptions.py
+++ b/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."""