Add proper concurrency support in websockets mode (#1036)

Author: wwwillchen

docs/api/config.md

@@ -12,6 +12,14 @@ Allows concurrent updates to state in the same session. If this is not updated,
 
 By default, this is not enabled. You can enable this by setting it to `true`.
 
+### MESOP_WEB_SOCKETS_ENABLED
+
+!!! warning "Experimental feature"
+
+    This is an experimental feature and is subject to breaking change. Please follow [https://github.com/google/mesop/issues/1028](https://github.com/google/mesop/issues/1028) for updates.
+
+This uses WebSockets instead of HTTP Server-Sent Events (SSE) as the transport protocol for UI updates. If you set this environment variable to `true`, then [`MESOP_CONCURRENT_UPDATES_ENABLED`](#MESOP_CONCURRENT_UPDATES_ENABLED) will automatically be enabled as well.
+
 ### MESOP_STATE_SESSION_BACKEND
 
 Sets the backend to use for caching state data server-side. This makes it so state does

mesop/cli/cli.py

@@ -11,6 +11,7 @@
   execute_module,
   get_module_name_from_runfile_path,
 )
+from mesop.env.env import MESOP_WEBSOCKETS_ENABLED
 from mesop.exceptions import format_traceback
 from mesop.runtime import (
   enable_debug_mode,
@@ -22,7 +23,6 @@
 from mesop.server.flags import port
 from mesop.server.logging import log_startup
 from mesop.server.server import configure_flask_app
-from mesop.server.server_utils import MESOP_WEBSOCKETS_ENABLED
 from mesop.server.static_file_serving import configure_static_file_serving
 from mesop.utils.host_util import get_public_host
 from mesop.utils.runfiles import get_runfile_location

mesop/env/BUILD

@@ -0,0 +1,10 @@
+load("//build_defs:defaults.bzl", "py_library")
+
+package(
+    default_visibility = ["//build_defs:mesop_internal"],
+)
+
+py_library(
+    name = "env",
+    srcs = glob(["*.py"]),
+)

mesop/env/__init__.py

@@ -0,0 +1,27 @@
+import os
+
+AI_SERVICE_BASE_URL = os.environ.get(
+  "MESOP_AI_SERVICE_BASE_URL", "http://localhost:43234"
+)
+
+MESOP_WEBSOCKETS_ENABLED = (
+  os.environ.get("MESOP_WEBSOCKETS_ENABLED", "false").lower() == "true"
+)
+
+MESOP_CONCURRENT_UPDATES_ENABLED = (
+  os.environ.get("MESOP_CONCURRENT_UPDATES_ENABLED", "false").lower() == "true"
+)
+
+if MESOP_WEBSOCKETS_ENABLED:
+  print("Experiment enabled: MESOP_WEBSOCKETS_ENABLED")
+  print("Auto-enabling MESOP_CONCURRENT_UPDATES_ENABLED")
+  MESOP_CONCURRENT_UPDATES_ENABLED = True
+elif MESOP_CONCURRENT_UPDATES_ENABLED:
+  print("Experiment enabled: MESOP_CONCURRENT_UPDATES_ENABLED")
+
+EXPERIMENTAL_EDITOR_TOOLBAR_ENABLED = (
+  os.environ.get("MESOP_EXPERIMENTAL_EDITOR_TOOLBAR", "false").lower() == "true"
+)
+
+if EXPERIMENTAL_EDITOR_TOOLBAR_ENABLED:
+  print("Experiment enabled: EXPERIMENTAL_EDITOR_TOOLBAR_ENABLED")

mesop/env/env.py

@@ -12,6 +12,9 @@
 from mesop.examples import composite as composite
 from mesop.examples import concurrency_state as concurrency_state
 from mesop.examples import concurrent_updates as concurrent_updates
+from mesop.examples import (
+  concurrent_updates_websockets as concurrent_updates_websockets,
+)
 from mesop.examples import custom_font as custom_font
 from mesop.examples import dict_state as dict_state
 from mesop.examples import docs as docs

mesop/examples/__init__.py

@@ -0,0 +1,35 @@
+import time
+
+import mesop as me
+
+
+@me.page(path="/concurrent_updates_websockets")
+def page():
+  state = me.state(State)
+  me.text("concurrent_updates_websockets")
+  me.button(label="Slow state update", on_click=slow_state_update)
+  me.button(label="Fast state update", on_click=fast_state_update)
+  me.text("Slow state: " + str(state.slow_state))
+  me.text("Fast state: " + str(state.fast_state))
+  if state.show_box:
+    with me.box():
+      me.text("Box!")
+
+
+@me.stateclass
+class State:
+  show_box: bool
+  slow_state: bool
+  fast_state: bool
+
+
+def slow_state_update(e: me.ClickEvent):
+  time.sleep(3)
+  me.state(State).show_box = True
+  me.state(State).slow_state = True
+  yield
+
+
+def fast_state_update(e: me.ClickEvent):
+  me.state(State).show_box = True
+  me.state(State).fast_state = True

mesop/examples/concurrent_updates_websockets.py

@@ -9,11 +9,13 @@ py_library(
     srcs = glob(["*.py"]),
     deps = [
         "//mesop/dataclass_utils",
+        "//mesop/env",
         "//mesop/events",
         "//mesop/exceptions",
         "//mesop/protos:ui_py_pb2",
         "//mesop/security",
         "//mesop/server:state_sessions",
         "//mesop/utils",
+        "//mesop/warn",
     ] + THIRD_PARTY_PY_FLASK,
 )

mesop/runtime/BUILD

@@ -1,5 +1,6 @@
 import asyncio
 import copy
+import threading
 import types
 import urllib.parse as urlparse
 from typing import Any, Callable, Generator, Sequence, TypeVar, cast
@@ -10,6 +11,7 @@
   serialize_dataclass,
   update_dataclass_from_json,
 )
+from mesop.env.env import MESOP_WEBSOCKETS_ENABLED
 from mesop.exceptions import (
   MesopDeveloperException,
   MesopException,
@@ -42,6 +44,22 @@ def __init__(
     self._theme_settings: pb.ThemeSettings | None = None
     self._js_modules: set[str] = set()
     self._query_params: dict[str, list[str]] = {}
+    if MESOP_WEBSOCKETS_ENABLED:
+      self._lock = threading.Lock()
+
+  def acquire_lock(self) -> None:
+    # No-op if websockets is not enabled because
+    # there shouldn't be concurrent updates to the same
+    # context instance.
+    if MESOP_WEBSOCKETS_ENABLED:
+      self._lock.acquire()
+
+  def release_lock(self) -> None:
+    # No-op if websockets is not enabled because
+    # there shouldn't be concurrent updates to the same
+    # context instance.
+    if MESOP_WEBSOCKETS_ENABLED:
+      self._lock.release()
 
   def register_js_module(self, js_module_path: str) -> None:
     self._js_modules.add(js_module_path)

mesop/runtime/context.py

@@ -2,14 +2,16 @@
 from dataclasses import dataclass
 from typing import Any, Callable, Generator, Type, TypeVar, cast
 
-from flask import g
+from flask import g, request
 
 import mesop.protos.ui_pb2 as pb
+from mesop.env.env import MESOP_WEBSOCKETS_ENABLED
 from mesop.events import LoadEvent, MesopEvent
 from mesop.exceptions import MesopDeveloperException, MesopUserException
 from mesop.key import Key
 from mesop.security.security_policy import SecurityPolicy
 from mesop.utils.backoff import exponential_backoff
+from mesop.warn import warn
 
 from .context import Context
 
@@ -54,12 +56,25 @@ def __init__(self):
     self._state_classes: list[type[Any]] = []
     self._loading_errors: list[pb.ServerError] = []
     self._has_served_traffic = False
+    self._contexts = {}
 
   def context(self) -> Context:
+    if MESOP_WEBSOCKETS_ENABLED and hasattr(request, "sid"):
+      # flask-socketio adds sid (session id) to the request object.
+      sid = request.sid  # type: ignore
+      if sid not in self._contexts:
+        self._contexts[sid] = self.create_context()
+      return self._contexts[sid]
     if "_mesop_context" not in g:
       g._mesop_context = self.create_context()
     return g._mesop_context
 
+  def delete_context(self, sid: str) -> None:
+    if sid in self._contexts:
+      del self._contexts[sid]
+    else:
+      warn(f"Tried to delete context with sid={sid} that doesn't exist.")
+
   def create_context(self) -> Context:
     # If running in prod mode, *always* enable the has served traffic safety check.
     # If running in debug mode, *disable* the has served traffic safety check.