fix: Move the renderer's LivenessScope into the context (#222)

Fixes #223

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

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
+import threading
 import logging
-from types import TracebackType
 from typing import Any, Callable, Optional, TypeVar, Union
-from contextlib import AbstractContextManager
+from deephaven.liveness_scope import LivenessScope
+from contextlib import AbstractContextManager, contextmanager
 
 logger = logging.getLogger(__name__)
 
@@ -40,7 +41,41 @@
 """
 
 
-class RenderContext(AbstractContextManager):
+_local_data = threading.local()
+
+
+class NoContextException(Exception):
+    pass
+
+
+def get_context() -> RenderContext:
+    """
+    Gets the currently active context, or throws NoContextException if none is set.
+
+    Returns:
+        The active RenderContext, or throws if none is present.
+    """
+    try:
+        return _local_data.context
+    except AttributeError:
+        raise NoContextException("No context set")
+
+
+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
+
+
+_READY_TO_OPEN: int = -2
+_OPENED_AND_UNUSED: int = -1
+
+
+class RenderContext:
     """
     Context for rendering a component. Keeps track of state and child contexts.
     Used by hooks to get and set state.
@@ -70,6 +105,10 @@ class RenderContext(AbstractContextManager):
     """
     The on_change callback to call when the context changes.
     """
+    _liveness_scope: LivenessScope
+    """
+    Liveness scope to create Deephaven items in. Need to retain the liveness scope so we don't release objects prematurely.
+    """
 
     def __init__(self, on_change: OnChangeCallable, on_queue_render: OnChangeCallable):
         """
@@ -80,28 +119,56 @@ def __init__(self, on_change: OnChangeCallable, on_queue_render: OnChangeCallabl
             on_queue_render: The callback to call when work is being requested for the render loop.
         """
 
-        self._hook_index = -1
+        self._hook_index = _READY_TO_OPEN
         self._hook_count = -1
         self._state = {}
         self._children_context = {}
         self._on_change = on_change
         self._on_queue_render = on_queue_render
+        self._liveness_scope = LivenessScope()
 
-    def __enter__(self) -> None:
+    @contextmanager
+    def open(self) -> AbstractContextManager:
         """
-        Start rendering this component.
-        """
-        self._hook_index = -1
+        Opens this context to track hook creation, sets this context as active on
+        this thread, and opens the liveness scope for user-created objects.
 
-    def __exit__(
-        self,
-        type: Optional[type[BaseException]],
-        value: Optional[BaseException],
-        traceback: Optional[TracebackType],
-    ) -> None:
-        """
-        Finish rendering this component.
+        This is not reentrant and not safe across threads, ensure it is only opened
+        once at a time. After it has been closed, it is safe to be opened again.
+
+        Returns:
+            A context manager to manage RenderContext resources.
         """
+        if self._hook_index != _READY_TO_OPEN:
+            raise RuntimeError(
+                "RenderContext.open() was already called, and is not reentrant"
+            )
+        self._hook_index = _OPENED_AND_UNUSED
+
+        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, self)
+        _set_context(self)
+
+        try:
+            new_liveness_scope = LivenessScope()
+            with new_liveness_scope.open():
+                yield self
+
+            # Following the "yield" so we don't do this if there was an error, we want to keep the old scope and kill
+            # the new one. We always release after creating the new one, so that each table/etc has its ref count go
+            # from 1 -> 2 -> 1, instead of 1 -> 0 -> 1 which would release the table prematurely
+            self._liveness_scope.release()
+            self._liveness_scope = new_liveness_scope
+        finally:
+            # Do this even if there was an error, old context must be restored
+            logger.debug("Resetting to old context %s", old_context)
+            _set_context(old_context)
+
+        # Outside the "finally" so we don't do this if there was an error, we don't want an incorrect hook count
         hook_count = self._hook_index + 1
         if self._hook_count < 0:
             self._hook_count = hook_count
@@ -112,6 +179,9 @@ def __exit__(
                 )
             )
 
+        # Reset count for next use to safeguard double-opening
+        self._hook_index = _READY_TO_OPEN
+
     def has_state(self, key: StateKey) -> bool:
         """
         Check if the given key is in the state.

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

@@ -3,8 +3,11 @@
     StateKey,
     StateUpdateCallable,
     OnChangeCallable,
+    InitializerFunction,
+    UpdaterFunction,
+    get_context,
+    NoContextException,
 )
-from .shared import get_context, set_context, NoContextException
 from .utils import (
     get_component_name,
     get_component_qualname,

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

@@ -1,8 +1,8 @@
 from __future__ import annotations
 import logging
-from typing import Callable, Optional
+from typing import Callable
 from .Element import Element, PropsType
-from .._internal import RenderContext, get_context, set_context, NoContextException
+from .._internal import RenderContext
 
 logger = logging.getLogger(__name__)
 
@@ -33,19 +33,7 @@ def render(self, context: RenderContext) -> PropsType:
         Returns:
             The props of this element.
         """
-        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)
-
-        with context:
+        with context.open():
             children = self._render()
 
-        logger.debug("Resetting to old context %s", old_context)
-        set_context(old_context)
-
         return {"children": children}

plugins/ui/src/deephaven/ui/hooks/use_render_queue.py

@@ -1,5 +1,4 @@
-from .._internal import OnChangeCallable
-from .._internal.shared import get_context
+from .._internal import OnChangeCallable, get_context
 
 
 def use_render_queue() -> OnChangeCallable:

plugins/ui/src/deephaven/ui/hooks/use_state.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 import logging
 from typing import Any, Callable, TypeVar, overload
-from .._internal.shared import get_context
-from .._internal.RenderContext import InitializerFunction, UpdaterFunction
+from .._internal import InitializerFunction, UpdaterFunction, get_context
 
 logger = logging.getLogger(__name__)
 

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

@@ -1,7 +1,6 @@
 from __future__ import annotations
 import logging
 from typing import Any
-from deephaven.liveness_scope import LivenessScope
 from .._internal import RenderContext
 from ..elements import Element, PropsType
 from .RenderedNode import RenderedNode
@@ -102,39 +101,11 @@ class Renderer:
     Context to render the element into. This is essentially the state of the element.
     """
 
-    _liveness_scope: LivenessScope
-    """
-    Liveness scope to create Deephaven items in. Need to retain the liveness scope so we don't release objects prematurely.
-    """
-
     def __init__(self, context: RenderContext):
         self._context = context
-        self._liveness_scope = LivenessScope()
-
-    def __del__(self):
-        self.release_liveness_scope()
-
-    def release_liveness_scope(self):
-        """
-        Release the liveness scope.
-        """
-        try:  # May not have an active liveness scope or already been released
-            self._liveness_scope.release()
-        except:
-            pass
 
     def render(self, element: Element) -> RenderedNode:
         """
         Render an element. Will update the liveness scope with the new objects from the render.
         """
-        new_liveness_scope = LivenessScope()
-        with new_liveness_scope.open():
-            render_res = _render_element(element, self._context)
-
-        # Release after in case of memoized tables
-        # Ref count goes 1 -> 2 -> 1 by releasing after
-        # instead of 1 -> 0 -> 1 which would release the table
-        self.release_liveness_scope()
-        self._liveness_scope = new_liveness_scope
-
-        return render_res
+        return _render_element(element, self._context)

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

@@ -1,7 +1,6 @@
 import threading
 import unittest
 from operator import itemgetter
-from time import sleep
 from typing import Callable
 from unittest.mock import Mock
 from .BaseTest import BaseTestCase
@@ -17,15 +16,13 @@ def render_hook(fn: Callable):
         Re-render will call the same function but with the new args passed in.
     """
     from deephaven.ui._internal.RenderContext import RenderContext
-    from deephaven.ui._internal.shared import get_context, set_context
 
     context = RenderContext(lambda x: x(), lambda x: x())
 
     return_dict = {"context": context, "result": None, "rerender": None}
 
     def _rerender(*args, **kwargs):
-        set_context(context)
-        with context:
+        with context.open():
             new_result = fn(*args, **kwargs)
             return_dict["result"] = new_result
         return new_result
@@ -156,7 +153,7 @@ def _test_table_listener(replayed_table_val=table, listener_val=listener):
             assert False, "listener was not called"
 
     def test_table_listener(self):
-        from deephaven import time_table, new_table, DynamicTableWriter
+        from deephaven import DynamicTableWriter
         import deephaven.dtypes as dht
 
         column_definitions = {"Numbers": dht.int32, "Words": dht.string}

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

@@ -2,7 +2,9 @@
 from .BaseTest import BaseTestCase
 
 
-def make_render_context(on_change=lambda x: x(), on_queue=lambda x: x()):
+def make_render_context(
+    on_change=lambda x: x(), on_queue=lambda x: x()
+) -> "deephaven.ui._internal.RenderContext":
     from deephaven.ui._internal.RenderContext import RenderContext
 
     return RenderContext(on_change, on_queue)
@@ -12,7 +14,7 @@ class RenderTestCase(BaseTestCase):
     def test_empty_render(self):
         on_change = Mock(side_effect=lambda x: x())
         rc = make_render_context(on_change)
-        self.assertEqual(rc._hook_index, -1)
+        self.assertEqual(rc._hook_index, -2)
         self.assertEqual(rc._state, {})
         self.assertEqual(rc._children_context, {})
         on_change.assert_not_called()
@@ -24,26 +26,26 @@ def test_hook_index(self):
         rc = make_render_context(on_change)
 
         # Set up the hooks used with initial render (3 hooks)
-        with rc:
+        with rc.open():
             self.assertEqual(rc.next_hook_index(), 0)
             self.assertEqual(rc.next_hook_index(), 1)
             self.assertEqual(rc.next_hook_index(), 2)
 
         # Verify it's the same on the next render
-        with rc:
+        with rc.open():
             self.assertEqual(rc.next_hook_index(), 0)
             self.assertEqual(rc.next_hook_index(), 1)
             self.assertEqual(rc.next_hook_index(), 2)
 
         # Check that an error is thrown if we don't use enough hooks
         with self.assertRaises(Exception):
-            with rc:
+            with rc.open():
                 self.assertEqual(rc.next_hook_index(), 0)
                 self.assertEqual(rc.next_hook_index(), 1)
 
         # Check that an error is thrown if we use too many hooks
         with self.assertRaises(Exception):
-            with rc:
+            with rc.open():
                 self.assertEqual(rc.next_hook_index(), 0)
                 self.assertEqual(rc.next_hook_index(), 1)
                 self.assertEqual(rc.next_hook_index(), 2)