Enhance Python LivenessScope API

py/server/deephaven/liveness_scope.py

@@ -8,11 +8,34 @@
 
 import jpy
 
+from typing import Union
+from warnings import warn
+
 from deephaven import DHError
 from deephaven._wrapper import JObjectWrapper
 
 _JLivenessScopeStack = jpy.get_type("io.deephaven.engine.liveness.LivenessScopeStack")
 _JLivenessScope = jpy.get_type("io.deephaven.engine.liveness.LivenessScope")
+_JLivenessReferent = jpy.get_type("io.deephaven.engine.liveness.LivenessReferent")
+
+
+def _push(scope: _JLivenessScope):
+    _JLivenessScopeStack.push(scope)
+
+
+def _pop(scope: _JLivenessScope):
+    try:
+        _JLivenessScopeStack.pop(scope)
+    except Exception as e:
+        raise DHError(e, message="failed to pop the LivenessScope from the stack.")
+
+
+def _unwrap_to_liveness_referent(referent: Union[JObjectWrapper, jpy.JType]) -> jpy.JType:
+    if isinstance(referent, jpy.JType) and _JLivenessReferent.jclass.isInstance(referent):
+        return referent
+    if isinstance(referent, JObjectWrapper):
+        return _unwrap_to_liveness_referent(referent.j_object)
+    raise DHError("Provided referent isn't a LivenessReferent or a JObjectWrapper around one")
 
 
 class LivenessScope(JObjectWrapper):
@@ -27,44 +50,143 @@ def __init__(self, j_scope: jpy.JType):
         self.j_scope = j_scope
 
     def __enter__(self):
+        warn('Instead of passing liveness_scope() to with, call open() on it first',
+             DeprecationWarning, stacklevel=2)
+        _push(self.j_scope)
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
+        _pop(self.j_scope)
+        self.release()
+
+    def release(self):
+        """Closes the LivenessScope and releases all the query graph resources.
+
+        Raises:
+            DHError
+        """
+        try:
+            self.j_scope.release()
+            self.j_scope = None
+        except Exception as e:
+            raise DHError(e, message="failed to close the LivenessScope.")
 
     def close(self):
         """Closes the LivenessScope and releases all the query graph resources.
 
         Raises:
             DHError
         """
-        if self.j_scope:
-            try:
-                _JLivenessScopeStack.pop(self.j_scope)
-                self.j_scope.release()
-                self.j_scope = None
-            except Exception as e:
-                raise DHError(e, message="failed to close the LivenessScope.")
+        warn('This function is deprecated, prefer release(). Use cases that rely on this are likely to now fail.',
+             DeprecationWarning, stacklevel=2)
+        _pop(self.j_scope)
+        self.release()
+
+    @contextlib.contextmanager
+    def open(self, release_after_block: bool = True) -> "LivenessScope":
+        """
+        Uses this scope for the duration of the `with` block. The scope defaults to being closed
+        when the block ends, disable by passing release_after_block=False
+
+        Args:
+            release_after_block: True to release the scope when the block ends, False to leave the
+            scope open, allowing it to be reused and keeping collected referents live.
+
+        Returns: None, to allow changes in the future.
+
+        """
+        _push(self.j_scope)
+        try:
+            yield self
+        finally:
+            _pop(self.j_scope)
+            if release_after_block:
+                self.release()
 
     @property
     def j_object(self) -> jpy.JType:
         return self.j_scope
 
-    def preserve(self, wrapper: JObjectWrapper) -> None:
+    def preserve(self, referent: Union[JObjectWrapper, jpy.JType]) -> None:
         """Preserves a query graph node (usually a Table) to keep it live for the outer scope.
 
         Args:
-            wrapper (JObjectWrapper): a wrapped Java object such as a Table
+            referent (Union[JObjectWrapper, jpy.JType]): an object to preserve in the next outer liveness scope
 
         Raises:
             DHError
         """
+        referent = _unwrap_to_liveness_referent(referent)
         try:
+            # Ensure we are the current scope, throw DHError if we aren't
             _JLivenessScopeStack.pop(self.j_scope)
-            _JLivenessScopeStack.peek().manage(wrapper.j_object)
-            _JLivenessScopeStack.push(self.j_scope)
+        except Exception as e:
+            raise DHError(e, message="failed to pop the current scope - is preserve() being called on the right scope?")
+
+        try:
+            # Manage the object in the next outer scope on this thread.
+            _JLivenessScopeStack.peek().manage(_unwrap_to_liveness_referent(referent))
         except Exception as e:
             raise DHError(e, message="failed to preserve a wrapped object in this LivenessScope.")
+        finally:
+            # Success or failure, restore the scope that was successfully popped
+            _JLivenessScopeStack.push(self.j_scope)
+
+    def manage(self, referent: Union[JObjectWrapper, jpy.JType]) -> None:
+        """
+        Explicitly manage the given java object in this scope. Must only be passed a Java LivenessReferent, or
+        a Python wrapper around a LivenessReferent
+
+        Args:
+            referent: the object to manage by this scope
+
+        Returns: None
+
+        Raises:
+            DHError if the referent isn't a LivenessReferent, or if it is no longer live
+
+        """
+        referent = _unwrap_to_liveness_referent(referent)
+        try:
+            self.j_scope.manage(referent)
+        except Exception as e:
+            raise DHError(e, message="failed to manage object")
+
+    def unmanage(self, referent: Union[JObjectWrapper, jpy.JType]) -> None:
+        """
+        Explicitly unmanage the given java object from this scope. Must only be passed a Java LivenessReferent, or
+        a Python wrapper around a LivenessReferent
+
+        Args:
+            referent: the object to unmanage from this scope
+
+        Returns: None
+
+        Raises:
+            DHError if the referent isn't a LivenessReferent, or if it is no longer live
+
+        """
+        referent = _unwrap_to_liveness_referent(referent)
+        try:
+            self.j_scope.unmanage(referent)
+        except Exception as e:
+            raise DHError(e, message="failed to unmanage object")
+
+
+def is_liveness_referent(referent: Union[JObjectWrapper, jpy.JType]) -> bool:
+    """
+    Returns True if the provided object is a LivenessReferent, and so can be managed by a LivenessScope.
+    Args:
+        referent: the object that maybe a LivenessReferent
+
+    Returns:
+        True if the object is a LivenessReferent, False otherwise.
+    """
+    if isinstance(referent, jpy.JType) and _JLivenessReferent.jclass.isInstance(referent):
+        return True
+    if isinstance(referent, JObjectWrapper):
+        return is_liveness_referent(referent.j_object)
+    return False
 
 
 def liveness_scope() -> LivenessScope:
@@ -74,5 +196,4 @@ def liveness_scope() -> LivenessScope:
         a LivenessScope
     """
     j_scope = _JLivenessScope()
-    _JLivenessScopeStack.push(j_scope)
     return LivenessScope(j_scope=j_scope)

py/server/deephaven/table_listener.py

@@ -22,7 +22,6 @@
 from deephaven.table import Table
 from deephaven.update_graph import UpdateGraph
 
-_JPythonListenerAdapter = jpy.get_type("io.deephaven.integrations.python.PythonListenerAdapter")
 _JPythonReplayListenerAdapter = jpy.get_type("io.deephaven.integrations.python.PythonReplayListenerAdapter")
 _JTableUpdate = jpy.get_type("io.deephaven.engine.table.TableUpdate")
 _JTableUpdateDataReader = jpy.get_type("io.deephaven.integrations.python.PythonListenerTableUpdateDataReader")
@@ -335,8 +334,9 @@ def listen(t: Table, listener: Union[Callable, TableListener], description: str
     return table_listener_handle
 
 
-class TableListenerHandle:
+class TableListenerHandle(JObjectWrapper):
     """A handle to manage a table listener's lifecycle."""
+    j_object_type = _JPythonReplayListenerAdapter
 
     def __init__(self, t: Table, listener: Union[Callable, TableListener], description: str = None):
         """Creates a new table listener handle.
@@ -408,3 +408,7 @@ def stop(self) -> None:
             return
         self.t.j_table.removeUpdateListener(self.listener)
         self.started = False
+
+    @property
+    def j_object(self) -> jpy.JType:
+        return self.listener

py/server/tests/test_liveness.py

@@ -24,7 +24,10 @@ def create_table(self):
         with exclusive_lock(self.test_update_graph):
             return time_table("PT00:00:00.001").update(["X=i%11"]).sort("X").tail(16)
 
-    def test_liveness(self):
+    def test_deprecated_liveness(self):
+        # Old test code, to make sure the basic cases still behave. Complex cases didn't make
+        # sense before, so aren't heavily tested. Note that these will emit warnings in the
+        # test output.
         not_managed = self.create_table()
         with liveness_scope() as l_scope:
             to_discard = self.create_table()
@@ -43,28 +46,51 @@ def test_liveness(self):
             must_keep = self.create_table()
             df = to_pandas(must_keep)
 
-        with self.assertRaises(DHError):
-            l_scope = liveness_scope()
+    def test_liveness(self):
+        not_managed = self.create_table()
+        with liveness_scope().open() as l_scope:
             to_discard = self.create_table()
             df = to_pandas(to_discard)
-            l_scope_2 = liveness_scope()
             must_keep = self.create_table()
             df = to_pandas(must_keep)
             l_scope.preserve(must_keep)
-            l_scope.close()
-            l_scope_2.close()
-        l_scope_2.close()
-        l_scope.close()
+
+        self.assertTrue(not_managed.j_table.tryRetainReference())
+        self.assertTrue(must_keep.j_table.tryRetainReference())
+        self.assertFalse(to_discard.j_table.tryRetainReference())
+
+        with liveness_scope().open():
+            to_discard = self.create_table()
+            df = to_pandas(to_discard)
+            must_keep = self.create_table()
+            df = to_pandas(must_keep)
+
+        with self.assertRaises(DHError):
+            l_scope = liveness_scope()
+            with l_scope.open(False):
+                to_discard = self.create_table()
+                df = to_pandas(to_discard)
+                l_scope_2 = liveness_scope()
+                with l_scope_2.open(False):
+                    must_keep = self.create_table()
+                    df = to_pandas(must_keep)
+                    l_scope.preserve(must_keep) # throws DHError
+                l_scope.release() # will never run
+            l_scope_2.release() # will never run
+        l_scope_2.release()
+        l_scope.release()
 
     def test_liveness_nested(self):
-        with liveness_scope() as l_scope:
+        l_scope = liveness_scope()
+        with l_scope.open():
             to_discard = self.create_table()
             df = to_pandas(to_discard)
             must_keep = self.create_table()
             df = to_pandas(must_keep)
             l_scope.preserve(must_keep)
 
-            with liveness_scope() as nested_l_scope:
+            nested_l_scope = liveness_scope()
+            with nested_l_scope.open():
                 nested_to_discard = self.create_table()
                 df = to_pandas(nested_to_discard)
                 nested_must_keep = self.create_table()
@@ -80,6 +106,35 @@ def test_liveness_nested(self):
         self.assertFalse(nested_must_keep.j_table.tryRetainReference())
         self.assertFalse(nested_to_discard.j_table.tryRetainReference())
 
+    def test_liveness_siblings(self):
+        l_scope = liveness_scope()
+        with l_scope.open():
+            to_discard = self.create_table()
+            df = to_pandas(to_discard)
+            must_keep = self.create_table()
+            df = to_pandas(must_keep)
+            l_scope.preserve(must_keep)
+
+        # Create a second scope and interact with it
+        other_l_scope = liveness_scope()
+        with other_l_scope.open():
+            nested_to_discard = self.create_table()
+            df = to_pandas(nested_to_discard)
+            nested_must_keep = self.create_table()
+            df = to_pandas(nested_must_keep)
+            other_l_scope.preserve(nested_must_keep)
+        self.assertTrue(nested_must_keep.j_table.tryRetainReference())
+        # drop the extra reference obtained by the tryRetainReference() call in the above assert
+        nested_must_keep.j_table.dropReference()
+        self.assertFalse(nested_to_discard.j_table.tryRetainReference())
+
+
+
+        self.assertTrue(must_keep.j_table.tryRetainReference())
+        self.assertFalse(to_discard.j_table.tryRetainReference())
+        self.assertFalse(nested_must_keep.j_table.tryRetainReference())
+        self.assertFalse(nested_to_discard.j_table.tryRetainReference())
+
 
 if __name__ == '__main__':
     unittest.main()

py/server/tests/testbase.py

@@ -35,9 +35,11 @@ def tearDownClass(cls) -> None:
 
     def setUp(self) -> None:
         self._liveness_scope = liveness_scope()
+        self.opened_scope = liveness_scope().open()
 
     def tearDown(self) -> None:
-        self._liveness_scope.close()
+        with self.opened_scope:
+            ...
 
     def wait_ticking_table_update(self, table: Table, row_count: int, timeout: int):
         """Waits for a ticking table to grow to the specified size or times out.