feat: Queue render state updates on a thread

- Now updates are all queued on the same thread
- Fixes deephaven#116

plugins/ui/src/deephaven/ui/_internal/RenderContext.py

@@ -6,7 +6,17 @@
 
 logger = logging.getLogger(__name__)
 
-OnChangeCallable = Callable[[], None]
+StateUpdateCallable = Callable[[], None]
+"""
+A callable that updates the state. Used to queue up state changes.
+"""
+
+OnChangeCallable = Callable[[StateUpdateCallable], None]
+"""
+Callable that is called when there is a change in the context (setting the state).
+"""
+
+
 StateKey = int
 ContextKey = Union[str, int]
 
@@ -36,17 +46,25 @@ class RenderContext(AbstractContextManager):
     """
     The child contexts for this context. 
     """
+
     _on_change: OnChangeCallable
     """
     The on_change callback to call when the context changes.
     """
 
-    def __init__(self):
+    def __init__(self, on_change: OnChangeCallable):
+        """
+        Create a new render context.
+
+        Args:
+            on_change: The on_change callback to call when the state in the context has changes.
+        """
+
         self._hook_index = -1
         self._hook_count = -1
         self._state = {}
         self._children_context = {}
-        self._on_change = lambda: None
+        self._on_change = on_change
 
     def __enter__(self) -> None:
         """
@@ -68,21 +86,6 @@ def __exit__(self, type, value, traceback) -> None:
                 )
             )
 
-    def _notify_change(self) -> None:
-        """
-        Notify the parent context that this context has changed.
-        Note that we're just re-rendering the whole tree on change.
-        TODO: We should be able to do better than this, and only re-render the parts that have actually changed.
-        """
-        logger.debug("Notifying parent context that child context has changed")
-        self._on_change()
-
-    def set_on_change(self, on_change: OnChangeCallable) -> None:
-        """
-        Set the on_change callback.
-        """
-        self._on_change = on_change
-
     def has_state(self, key: StateKey) -> bool:
         """
         Check if the given key is in the state.
@@ -101,15 +104,17 @@ def set_state(self, key: StateKey, value: Any) -> None:
         """
         Set the state for the given key.
         """
-        # TODO: Should we throw here if it's called when we're in the middle of a render?
-        # TODO: How do we batch the state changes so they run on the next loop?
-        should_notify = False
-        if key in self._state:
-            # We only want to notify of a change when the value actually changes, not on the initial render
-            should_notify = True
-        self._state[key] = value
-        if should_notify:
-            self._notify_change()
+
+        # We queue up the state change in a callable that will get called from the render loop
+        def update_state():
+            self._state[key] = value
+
+        if key not in self._state:
+            # We haven't set the state for this key yet, this is the initial render. We can just set the state immediately, we don't need to queue it for notification
+            update_state()
+        else:
+            # This is not the initial state, queue up the state change on the render loop
+            self._on_change(update_state)
 
     def get_child_context(self, key: ContextKey) -> "RenderContext":
         """
@@ -118,8 +123,7 @@ def get_child_context(self, key: ContextKey) -> "RenderContext":
         logger.debug("Getting child context for key %s", key)
         if key not in self._children_context:
             logger.debug("Creating new child context for key %s", key)
-            child_context = RenderContext()
-            child_context.set_on_change(self._notify_change)
+            child_context = RenderContext(self._on_change)
             self._children_context[key] = child_context
         return self._children_context[key]
 

plugins/ui/src/deephaven/ui/_internal/__init__.py

@@ -1,4 +1,9 @@
-from .RenderContext import RenderContext
+from .RenderContext import (
+    RenderContext,
+    StateKey,
+    StateUpdateCallable,
+    OnChangeCallable,
+)
 from .shared import get_context, set_context
 from .utils import (
     get_component_name,

plugins/ui/src/deephaven/ui/object_types/ElementMessageStream.py

@@ -3,16 +3,39 @@
 import json
 from jsonrpc import JSONRPCResponseManager, Dispatcher
 import logging
-from typing import Any
+import threading
+from queue import Queue
+from typing import Any, Callable
 from deephaven.plugin.object_type import MessageStream
 from ..elements import Element
 from ..renderer import NodeEncoder, Renderer, RenderedNode
-from .._internal import RenderContext
+from .._internal import RenderContext, StateUpdateCallable
 from ..renderer.NodeEncoder import NodeEncoder
 
 logger = logging.getLogger(__name__)
 
-# TODO: Get a thread from a thread pool here for rendering and callbacks?
+_render_queue: Queue[Callable[[], None]] = Queue()
+
+_render_thread: threading.Thread | None = None
+
+
+def _render_loop():
+    global _render_queue
+    while True:
+        item = _render_queue.get()
+        try:
+            item()
+        except Exception as e:
+            logger.exception(e)
+
+
+def _start_render_loop():
+    global _render_thread
+    if _render_thread is None or _render_thread.is_alive() is False:
+        _render_thread = threading.Thread(
+            target=_render_loop, name="deephaven.ui render loop"
+        )
+        _render_thread.start()
 
 
 class ElementMessageStream(MessageStream):
@@ -46,6 +69,26 @@ class ElementMessageStream(MessageStream):
     The connection to send the rendered element to.
     """
 
+    _context: RenderContext
+    """
+    Render context for this element
+    """
+
+    _renderer: Renderer
+    """
+    Renderer for this element
+    """
+
+    _queued_updates: Queue[StateUpdateCallable]
+    """
+    State updates that need to be applied on the next render.
+    """
+
+    _is_render_queued: bool
+    """
+    Whether or not a render is queued.
+    """
+
     def __init__(self, element: Element, connection: MessageStream):
         """
         Create a new ElementMessageStream. Renders the element in a render context, and sends the rendered result to the
@@ -61,18 +104,42 @@ def __init__(self, element: Element, connection: MessageStream):
         self._manager = JSONRPCResponseManager()
         self._dispatcher = Dispatcher()
         self._encoder = NodeEncoder(separators=(",", ":"))
+        self._context = RenderContext(self._queue_state_update)
+        self._renderer = Renderer(self._context)
+        self._queued_updates = Queue()
+        self._is_render_queued = False
 
-    def start(self) -> None:
-        context = RenderContext()
-        renderer = Renderer(context)
+    def _render(self) -> None:
+        logger.debug("ElementMessageStream._render")
+
+        # Resolve any pending state updates first
+        while not self._queued_updates.empty():
+            state_update = self._queued_updates.get()
+            state_update()
+
+        node = self._renderer.render(self._element)
+        self._send_document_update(node)
+        self._is_render_queued = False
+
+    def _queue_render(self) -> None:
+        """
+        Queue a render to be resolved on the next render.
+        """
+        if self._is_render_queued:
+            return
+        self._is_render_queued = True
+        _render_queue.put(self._render)
 
-        def update():
-            logger.debug("ElementMessageStream update")
-            node = renderer.render(self._element)
-            self._send_document_update(node)
+    def _queue_state_update(self, state_update: StateUpdateCallable) -> None:
+        """
+        Queue a state update to be resolved on the next render.
+        """
+        self._queued_updates.put(state_update)
+        self._queue_render()
 
-        context.set_on_change(update)
-        update()
+    def start(self) -> None:
+        _start_render_loop()
+        self._queue_render()
 
     def on_close(self) -> None:
         pass
@@ -144,6 +211,12 @@ def _send_document_update(self, root: RenderedNode) -> None:
         dispatcher = Dispatcher()
         for callable, callable_id in callable_id_dict.items():
             logger.debug("Registering callable %s", callable_id)
-            dispatcher[callable_id] = callable
+
+            def wrapped_callable(*args, **kwargs):
+                # We want all callables to be triggered from the deephaven.ui render thread
+                logger.debug("Calling callable %s", callable_id)
+                _render_queue.put(lambda: callable(*args, **kwargs))
+
+            dispatcher[callable_id] = wrapped_callable
         self._dispatcher = dispatcher
         self._connection.on_data(payload.encode(), new_objects)

plugins/ui/src/deephaven/ui/renderer/Renderer.py

@@ -76,7 +76,7 @@ def _render_element(element: Element, context: RenderContext) -> RenderedNode:
 class Renderer:
     _liveness_scope: LivenessScope
 
-    def __init__(self, context: RenderContext = RenderContext()):
+    def __init__(self, context: RenderContext):
         self._context = context
         self._liveness_scope = LivenessScope()
 

plugins/ui/test/deephaven/ui/test_hooks.py

@@ -19,7 +19,7 @@ def render_hook(fn: Callable):
     from deephaven.ui._internal.RenderContext import RenderContext
     from deephaven.ui._internal.shared import get_context, set_context
 
-    context = RenderContext()
+    context = RenderContext(lambda x: None)
 
     return_dict = {"context": context, "result": None, "rerender": None}
 

plugins/ui/test/deephaven/ui/test_render.py

@@ -10,7 +10,6 @@ def test_empty_render(self):
         self.assertEqual(rc._hook_index, -1)
         self.assertEqual(rc._state, {})
         self.assertEqual(rc._children_context, {})
-        self.assertEqual(rc._on_change(), None)
 
     def test_hook_index(self):
         from deephaven.ui._internal.RenderContext import RenderContext
@@ -46,10 +45,8 @@ def test_hook_index(self):
     def test_state(self):
         from deephaven.ui._internal.RenderContext import RenderContext
 
-        rc = RenderContext()
-
         on_change = Mock()
-        rc.set_on_change(on_change)
+        rc = RenderContext(on_change)
 
         self.assertEqual(rc.has_state(0), False)
         self.assertEqual(rc.get_state(0), None)
@@ -68,10 +65,8 @@ def test_state(self):
     def test_context(self):
         from deephaven.ui._internal.RenderContext import RenderContext
 
-        rc = RenderContext()
-
         on_change = Mock()
-        rc.set_on_change(on_change)
+        rc = RenderContext(on_change)
 
         child_context0 = rc.get_child_context(0)
         child_context1 = rc.get_child_context(1)