feat: Queue render state updates on a thread (#182)

- Create a thread and queue for applying state updates
- Now callbacks and state updates are all queued on the same thread
- Cleaned up some typing in the `Renderer` class
- Added a `use_render_queue` hook and example for showing how to queue
state updates
- Fixes #116
- Tested using the snippet from #116:
```python
import deephaven.ui as ui
from deephaven.ui import use_state


@ui.component
def foo():
    a, set_a = use_state(0)
    b, set_b = use_state(0)

    print(f"Render with a {a} and b {b}")

    def handle_press():
        set_a(a + 1)
        set_b(b + 1)

    return ui.action_button(
        f"a is {a} and b is {b}", on_press=handle_press
    )


f = foo()
```
- Ensured that the print out was only printed once per press of the
button, and both values stayed the same
- Allow a callable to be passed into the use_state setter function that
takes the old value as a parameter. Tested using the following
component:
```python
import deephaven.ui as ui
from deephaven.ui import use_state


@ui.component
def bar():
    x, set_x = use_state(0)

    print(f"Render with x {x}")

    def handle_press():
        # Call set_x twice in the same method, using a callable. This should result in x increasing by 2 each time the button is pressed
        set_x(lambda old_x: old_x + 1)
        set_x(lambda old_x: old_x + 1)

    return ui.action_button(
        f"x is {x}", on_press=handle_press
    )


b = bar()
```
- Tested that trying to update state from the wrong thread throws an
error:
```python
import deephaven.ui as ui
import threading
import time
from deephaven.ui import use_state


@ui.component
def foo():
    a, set_a = use_state(0)
    b, set_b = use_state(0)

    print(f"Render with a {a} and b {b}")

    def handle_press():
        def update_state():
            time.sleep(1)
            set_a(a + 1)
            set_b(b + 1)
        # Not using the correct thread
        threading.Thread(target=update_state).start()

    return ui.action_button(
        f"a is {a} and b is {b}", on_press=handle_press
    )


f = foo()
```

plugins/ui/examples/README.md

@@ -813,3 +813,93 @@ watch = watch_lizards(stocks)
 ```
 
 ![Table Hooks](assets/table_hooks.png)
+
+## Multi-threading
+
+State updates must be called from the render thread. All callbacks are automatically called from the render thread, but sometimes you will need to do some long-running operations asynchronously. You can use the `use_render_queue` hook to run a callback on the render thread. In this example, we create a form that takes a URL as input, and loads the CSV file from another thread before updating the state on the current thread.
+
+```python
+import logging
+import threading
+import time
+from deephaven import read_csv, ui
+
+
+@ui.component
+def csv_loader():
+    # The render_queue we fetch using the `use_render_queue` hook at the top of the component
+    render_queue = ui.use_render_queue()
+    table, set_table = ui.use_state()
+    error, set_error = ui.use_state()
+
+    def handle_submit(data):
+        # We define a callable that we'll queue up on our own thread
+        def load_table():
+            try:
+                # Read the table from the URL
+                t = read_csv(data["url"])
+
+                # Define our state updates in another callable. We'll need to call this on the render thread
+                def update_state():
+                    set_error(None)
+                    set_table(t)
+
+                # Queue up the state update on the render thread
+                render_queue(update_state)
+            except Exception as e:
+                # In case we have any errors, we should show the error to the user. We still need to call this from the render thread,
+                # so we must assign the exception to a variable and call the render_queue with a callable that will set the error
+                error_message = e
+
+                def update_state():
+                    set_table(None)
+                    set_error(error_message)
+
+                # Queue up the state update on the render thread
+                render_queue(update_state)
+
+        # Start our own thread loading the table
+        threading.Thread(target=load_table).start()
+
+    return [
+        # Our form displaying input from the user
+        ui.form(
+            ui.flex(
+                ui.text_field(
+                    default_value="https://media.githubusercontent.com/media/deephaven/examples/main/DeNiro/csv/deniro.csv",
+                    label="Enter URL",
+                    label_position="side",
+                    name="url",
+                    flex_grow=1,
+                ),
+                ui.button(f"Load Table", type="submit"),
+                gap=10,
+            ),
+            on_submit=handle_submit,
+        ),
+        (
+            # Display a hint if the table is not loaded yet and we don't have an error
+            ui.illustrated_message(
+                ui.heading("Enter URL above"),
+                ui.content("Enter a URL of a CSV above and click 'Load' to load it"),
+            )
+            if error is None and table is None
+            else None
+        ),
+        # The loaded table. Doesn't show anything if it is not loaded yet
+        table,
+        # An error message if there is an error
+        (
+            ui.illustrated_message(
+                ui.icon("vsWarning"),
+                ui.heading("Error loading table"),
+                ui.content(f"{error}"),
+            )
+            if error != None
+            else None
+        ),
+    ]
+
+
+my_loader = csv_loader()
+```

plugins/ui/setup.cfg

@@ -25,7 +25,8 @@ package_dir=
     =src
 packages=find_namespace:
 install_requires =
-    deephaven-plugin
+    deephaven-core>=0.31.0
+    deephaven-plugin>=0.6.0
     json-rpc
 include_package_data = True
 

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

@@ -1,14 +1,43 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, Callable
+from types import TracebackType
+from typing import Any, Callable, Optional, TypeVar, Union
 from contextlib import AbstractContextManager
 
 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 = str
+"""
+The key for a state value. Should be the hook index.
+"""
+
+T = TypeVar("T")
+InitializerFunction = Callable[[], T]
+"""
+A function that returns the initial value for a state.
+"""
+
+UpdaterFunction = Callable[[T], T]
+"""
+A function that takes the old value and returns the new value for a state.
+"""
+
+ContextKey = Union[str, int]
+"""
+The key for a child context.
+"""
 
 
 class RenderContext(AbstractContextManager):
@@ -36,25 +65,40 @@ 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, on_queue_render: OnChangeCallable):
+        """
+        Create a new render context.
+
+        Args:
+            on_change: The callback to call when the state in the context has changes.
+            on_queue_render: The callback to call when work is being requested for the render loop.
+        """
+
         self._hook_index = -1
         self._hook_count = -1
         self._state = {}
         self._children_context = {}
-        self._on_change = lambda: None
+        self._on_change = on_change
+        self._on_queue_render = on_queue_render
 
     def __enter__(self) -> None:
         """
         Start rendering this component.
         """
         self._hook_index = -1
 
-    def __exit__(self, type, value, traceback) -> None:
+    def __exit__(
+        self,
+        type: Optional[type[BaseException]],
+        value: Optional[BaseException],
+        traceback: Optional[TracebackType],
+    ) -> None:
         """
         Finish rendering this component.
         """
@@ -68,47 +112,50 @@ 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.
         """
         return key in self._state
 
-    def get_state(self, key: StateKey, default: Any = None) -> None:
+    def get_state(self, key: StateKey) -> Any:
         """
         Get the state for the given key.
         """
-        if key not in self._state:
-            self._state[key] = default
         return self._state[key]
 
-    def set_state(self, key: StateKey, value: Any) -> None:
+    def init_state(self, key: StateKey, value: T | InitializerFunction[T]) -> None:
         """
-        Set the state for the given key.
+        Set the initial state for the given key. Will throw if the key has already been set.
         """
-        # TODO: Should we throw here if it's called when we're in the middle of a render?
-        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()
+            raise KeyError(f"Key {key} is already initialized")
+
+        # Just set the key value, we don't need to trigger an on_change or anything special on initialization
+        self._state[key] = value if not callable(value) else value()
+
+    def set_state(self, key: StateKey, value: T | UpdaterFunction[T]) -> None:
+        """
+        Update the state for the given key. If the key is not initialized via `init_state` yet, throw.
+
+        Args:
+            key: The key to set the state for.
+            value: The value to set the state to. Can be a callable that takes the old value and returns the new value.
+        """
+
+        if key not in self._state:
+            raise KeyError(f"Key {key} not initialized")
+
+        # We queue up the state change in a callable that will get called from the render loop
+        def update_state():
+            new_value = value
+            if callable(value):
+                old_value = self._state[key]
+                new_value = value(old_value)
+            self._state[key] = new_value
+
+        # 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":
         """
@@ -117,8 +164,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._on_queue_render)
             self._children_context[key] = child_context
         return self._children_context[key]
 
@@ -128,3 +174,12 @@ def next_hook_index(self) -> int:
         """
         self._hook_index += 1
         return self._hook_index
+
+    def queue_render(self, update: Callable[[], None]) -> None:
+        """
+        Queue up a state update. Needed in multi-threading scenarios.
+
+        Args:
+            update: The update to queue up.
+        """
+        self._on_queue_render(update)

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

@@ -1,5 +1,10 @@
-from .RenderContext import RenderContext
-from .shared import get_context, set_context
+from .RenderContext import (
+    RenderContext,
+    StateKey,
+    StateUpdateCallable,
+    OnChangeCallable,
+)
+from .shared import get_context, set_context, NoContextException
 from .utils import (
     get_component_name,
     get_component_qualname,

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

@@ -1,28 +1,27 @@
 from .RenderContext import RenderContext
 from typing import Optional
+import threading
 
 
-class UiSharedInternals:
-    """
-    Shared internal context for the deephaven.ui plugin to use when rendering.
-    Should be set at the start of a render call, and unset at the end.
-
-    TODO: Need to keep track of context for each given thread, in case we have multiple threads rendering at once.
-    """
-
-    _current_context: Optional[RenderContext] = None
+class NoContextException(Exception):
+    pass
 
-    @property
-    def current_context(self) -> RenderContext:
-        return self._current_context
 
-
-_deephaven_ui_shared_internals: UiSharedInternals = UiSharedInternals()
+_local_data = threading.local()
 
 
 def get_context() -> RenderContext:
-    return _deephaven_ui_shared_internals.current_context
+    try:
+        return _local_data.context
+    except AttributeError:
+        raise NoContextException("No context set")
 
 
-def set_context(context):
-    _deephaven_ui_shared_internals._current_context = context
+def set_context(context: Optional[RenderContext]):
+    """
+    Set the current context for the thread. Can be set to None to unset the context for a thread
+    """
+    if context is None:
+        del _local_data.context
+    else:
+        _local_data.context = context

plugins/ui/src/deephaven/ui/elements/FunctionElement.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 import logging
-from typing import Callable
-from .Element import Element
-from .._internal import RenderContext, get_context, set_context
+from typing import Callable, Optional
+from .Element import Element, PropsType
+from .._internal import RenderContext, get_context, set_context, NoContextException
 
 logger = logging.getLogger(__name__)
 
@@ -23,17 +23,21 @@ def __init__(self, name: str, render: Callable[[], list[Element]]):
     def name(self):
         return self._name
 
-    def render(self, context: RenderContext) -> list[Element]:
+    def render(self, context: RenderContext) -> PropsType:
         """
         Render the component. Should only be called when actually rendering the component, e.g. exporting it to the client.
 
         Args:
             context: Context to render the component in
 
         Returns:
-            The rendered component.
+            The props of this element.
         """
-        old_context = get_context()
+        old_context: Optional[RenderContext] = None
+        try:
+            old_context = get_context()
+        except NoContextException:
+            pass
         logger.debug("old context is %s and new context is %s", old_context, context)
 
         set_context(context)