Enhance Python LivenessScope API (#4497)

The existing Python APIs for Liveness were difficult to use correctly
in cases more complex than a single script - this brings several parts
of the Java LivenessScope type to Python to enable more advanced use
cases. Rather than introduce a LivenessScopeStack, this adds an open()
method to LivenessScope itself. The liveness_scope() method will
continue to exist for simple use cases, where the scope will be closed
and released at the same time at the end of a block or function.

Also makes TableListenerHandle a JObjectWrapper so it can be more
easily unwrapped, can't accidentally be managed as a
LivePyObjectWrapper in Java.

Fixes #4471

py/server/deephaven/liveness_scope.py

@@ -3,76 +3,249 @@
 #
 
 """This module gives the users a finer degree of control over when to clean up unreferenced nodes in the query update
-graph instead of solely relying on garbage collection."""
+graph instead of solely relying on garbage collection.
+
+Examples:
+
+ Use the liveness_scope() function to produce a simple liveness scope that will be used only within a `with` expression
+ or as a decorator.
+
+ .. code-block:: python
+
+ with liveness_scope() as scope:
+ ticking_table = some_ticking_source()
+ table = ticking_table.snapshot().join(table=other_ticking_table, on=...)
+ scope.preserve(table)
+ return table
+
+ .. code-block:: python
+
+ @liveness_scope
+ def get_values():
+ ticking_table = some_ticking_source().last_by("Sym")
+ return dhnp.to_numpy(ticking_table)
+
+
+ Use the LivenessScope type for greater control, allowing a scope to be opened more than once, and to release the
+ resources that it manages when the scope will no longer be used.
+
+
+ .. code-block:: python
+
+ def make_table_and_scope(a: int):
+ scope = LivenessScope()
+ with scope.open():
+ ticking_table = some_ticking_source().where(f"A={a}")
+ return some_ticking_table, scope
+
+ t1, s1 = make_table_and_scope(1)
+ # .. wait for a while
+ s1.release()
+ t2, s2 = make_table_and_scope(2)
+ # etc
+
+ In both cases, the scope object has a few methods that can be used to more directly manage liveness referents:
+ * `scope.preserve(obj)` will preserve the given instance in the next scope outside the specified scope instance
+ * `scope.manange(obj)` will directly manage the given instance in the specified scope. Take care not to
+ double-manage objects when using in conjunction with a `with` block or function decorator.
+ * `scope.unmanage(obj)` will stop managing the given instance. This can be used regardless of how the instance
+ was managed to begin with.
+"""
 import contextlib
 
 import jpy
 
+from typing import Union, Iterator
+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")
 
 
-class LivenessScope(JObjectWrapper):
- """A LivenessScope automatically manages reference counting of tables and other query resources that are
- created in it. It implements the context manager protocol and thus can be used in the with statement.
+def _push(scope: _JLivenessScope) -> None:
+ _JLivenessScopeStack.push(scope)
+
+
+def _pop(scope: _JLivenessScope) -> None:
+ try:
+ _JLivenessScopeStack.pop(scope)
+ except Exception as e:
+ raise DHError(e, message="failed to pop the LivenessScope from the stack.")
 
- Note, LivenessScope should not be instantiated directly but rather through the 'liveness_scope' function.
- """
- j_object_type = _JLivenessScope
 
- def __init__(self, j_scope: jpy.JType):
- self.j_scope = j_scope
+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")
 
- def __enter__(self):
- return self
 
- def __exit__(self, exc_type, exc_val, exc_tb):
- self.close()
+class _BaseLivenessScope(JObjectWrapper):
+ """
+ Internal base type for Java LivenessScope types in python.
+ """
+ j_object_type = _JLivenessScope
+
+ def __init__(self):
+ self.j_scope = _JLivenessScope()
 
- def close(self):
+ def _release(self) -> None:
  """Closes the LivenessScope and releases all the query graph resources.
 
  Raises:
- DHError
+ DHError if this instance has been released too many times
  """
- 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.")
-
- @property
- def j_object(self) -> jpy.JType:
- return self.j_scope
+ try:
+ self.j_scope.release()
+ self.j_scope = None
+ except Exception as e:
+ raise DHError(e, message="failed to close the LivenessScope.")
 
- 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
+ DHError if the object isn't a liveness node, or this instance isn't currently at the stop of the stack
  """
+ 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 (Union[JObjectWrapper, jpy.JType]): 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 (Union[JObjectWrapper, jpy.JType]): 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")
+
+ @property
+ def j_object(self) -> jpy.JType:
+ return self.j_scope
+
+
+class SimpleLivenessScope(_BaseLivenessScope):
+ """
+ A SimpleLivenessScope automatically manages reference counting of tables and other query resources that
+ are created in it. Instances are created through calling `liveness_scope()` in a `with` block or as a
+ function decorator, and will be disposed of automatically.
+ """
+
+
+class LivenessScope(_BaseLivenessScope):
+ """A LivenessScope automatically manages reference counting of tables and other query resources that are
+ created in it. Any created instances must have `release()` called on them to correctly free resources
+ that have been managed.
+ """
 
+ def __init__(self):
+ super().__init__()
 
-def liveness_scope() -> LivenessScope:
- """Creates a LivenessScope for running a block of code.
+ def release(self) -> None:
+ """Closes the LivenessScope and releases all the managed resources.
+
+ Raises:
+ DHError if this instance has been released too many times
+ """
+ self._release()
+
+ @contextlib.contextmanager
+ def open(self) -> Iterator[None]:
+ """
+ Uses this scope for the duration of a `with` block, automatically managing all resources created in the block.
+
+ Returns: None, to allow changes in the future.
+ """
+ _push(self.j_scope)
+ try:
+ yield None
+ finally:
+ _pop(self.j_scope)
+
+
+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 may be a LivenessReferent
 
  Returns:
- a LivenessScope
+ 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
+
+
+@contextlib.contextmanager
+def liveness_scope() -> Iterator[SimpleLivenessScope]:
+ """Creates and opens a LivenessScope for running a block of code. Use
+ this function to wrap a block of code using a `with` statement.
+
+ For the duration of the `with` block, the liveness scope will be open
+ and any liveness referents created will be manged by it automatically.
+
+ Yields:
+ a SimpleLivenessScope
  """
- j_scope = _JLivenessScope()
- _JLivenessScopeStack.push(j_scope)
- return LivenessScope(j_scope=j_scope)
+ scope = SimpleLivenessScope()
+ _push(scope.j_scope)
+ try:
+ yield scope
+ finally:
+ _pop(scope.j_scope)
+ scope._release()

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