feat: DH-18423: add a systemic_obj_tracker module in the Py server API (deephaven#6577)

Added a systemic_obj_tracker.py module which allows users to
enable/disable systemic object marking on the current thread.

Author: jmao-denver

engine/updategraph/src/main/java/io/deephaven/engine/util/systemicmarking/SystemicObjectTracker.java

@@ -43,7 +43,8 @@ public static void markThreadSystemic() {
     }
 
     /**
-     * Marks the current thread as systemically important. This can be changed with {@link #markThreadSystemic()} ()}
+     * Marks the current thread as not systemically important. This can be changed with {@link #markThreadSystemic()}
+     * ()}
      */
     public static void markThreadNotSystemic() {
         if (SYSTEMIC_OBJECT_MARKING_ENABLED) {

py/server/deephaven/systemic_obj_tracker.py

@@ -0,0 +1,86 @@
+#
+# Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending
+#
+"""This module allows user to enable/disable Deephaven systemic object marking. When enabled, Deephaven will mark
+all objects created in the current thread as systemic. These systemic objects will be tracked and if errors occur to
+them, the errors are deemed to be systemic and fatal.
+"""
+import contextlib
+
+import jpy
+from deephaven import DHError
+
+_JSystemicObjectTracker = jpy.get_type("io.deephaven.engine.util.systemicmarking.SystemicObjectTracker")
+
+
+def is_systemic_object_marking_enabled() -> bool:
+    """Returns True if the systemic object marking is enabled. When enabled, the current thread can be marked as
+    systemic or not systemic.
+    """
+    return _JSystemicObjectTracker.isSystemicObjectMarkingEnabled()
+
+
+def is_systemic() -> bool:
+    """Returns whether the current thread is systemic. If true, objects created on this thread are treated as
+    systemic."""
+    return _JSystemicObjectTracker.isSystemicThread()
+
+
+def set_systemic(systemic: bool) -> None:
+    """Sets whether the current thread is systemic. If true, objects created on this thread are treated as systemic.
+
+    Args:
+        systemic (bool): True to mark the current thread as systemic, False to mark it as not systemic.
+
+    Raises:
+        DHError: If the systemic object marking is not enabled.
+    """
+    if is_systemic_object_marking_enabled():
+        if systemic:
+            _JSystemicObjectTracker.markThreadSystemic()
+        else:
+            _JSystemicObjectTracker.markThreadNotSystemic()
+    else:
+        raise DHError(message="Systemic object marking is not enabled.")
+
+
+@contextlib.contextmanager
+def systemic_object_marking() -> None:
+    """A Context manager to ensure the current thread is marked as systemic for the execution of the enclosed code
+    block. On exit, the thread is restored to its previous systemic state.
+
+    Raises:
+        DHError: If the systemic object marking is not enabled.
+    """
+    if is_systemic_object_marking_enabled():
+        if not is_systemic():
+            try:
+                _JSystemicObjectTracker.markThreadSystemic()
+                yield
+            finally:
+                _JSystemicObjectTracker.markThreadNotSystemic()
+        else:
+            yield
+    else:
+        raise DHError(message="Systemic object marking is not enabled.")
+
+
+@contextlib.contextmanager
+def no_systemic_object_marking() -> None:
+    """A Context manager to ensure the current thread is marked as not systemic for the execution of the enclosed code
+    block. On exit, the thread is restored to its previous systemic state.
+
+    Raises:
+        DHError: If the systemic object marking is not enabled.
+    """
+    if is_systemic_object_marking_enabled():
+        if is_systemic():
+            try:
+                _JSystemicObjectTracker.markThreadNotSystemic()
+                yield
+            finally:
+                _JSystemicObjectTracker.markThreadSystemic()
+        else:
+            yield
+    else:
+        raise DHError(message="Systemic object marking is not enabled.")

py/server/test_helper/__init__.py

@@ -37,6 +37,7 @@ def start_jvm_for_tests(jvm_props: Dict[str, str] = None):
 
             'Calendar.default': 'USNYSE_EXAMPLE',
             'Calendar.importPath': '/test_calendar_imports.txt',
+            'SystemicObjectTracker.enabled': 'true',
         }
 
         if jvm_props:

py/server/tests/test_systemic_obj_tracker.py

@@ -0,0 +1,58 @@
+#
+# Copyright (c) 2016-2024 Deephaven Data Labs and Patent Pending
+#
+import unittest
+from unittest.mock import patch
+from deephaven import DHError, empty_table
+import deephaven.systemic_obj_tracker as tracker
+
+
+class TestSystemicObjectTracker(unittest.TestCase):
+
+    def test_is_systemic_object_marking_enabled(self):
+        self.assertTrue(tracker.is_systemic_object_marking_enabled())
+
+    def test_is_systemic(self):
+        # user thread is not systemic by default
+        self.assertFalse(tracker.is_systemic())
+
+    def test_set_systemic(self):
+        tracker.set_systemic(True)
+        self.assertTrue(tracker.is_systemic())
+        t = empty_table(10)
+        self.assertTrue(t.j_object.isSystemicObject())
+
+        tracker.set_systemic(False)
+        self.assertFalse(tracker.is_systemic())
+        t = empty_table(10)
+        self.assertFalse(t.j_object.isSystemicObject())
+
+        with patch("deephaven.systemic_obj_tracker._JSystemicObjectTracker") as mock_tracker:
+            mock_tracker.isSystemicObjectMarkingEnabled.return_value = False
+            with self.assertRaises(DHError) as cm:
+                tracker.set_systemic(True)
+            self.assertIn("Systemic object marking is not enabled", str(cm.exception))
+
+    def test_systemic_object_marking(self):
+        with tracker.systemic_object_marking():
+            self.assertTrue(tracker.is_systemic())
+            t = empty_table(10)
+        self.assertTrue(t.j_object.isSystemicObject())
+
+    def test_systemic_object_marking_error(self):
+        with patch("deephaven.systemic_obj_tracker._JSystemicObjectTracker") as mock_tracker:
+            mock_tracker.isSystemicObjectMarkingEnabled.return_value = False
+            with self.assertRaises(DHError) as cm:
+                with tracker.systemic_object_marking():
+                    pass
+            self.assertIn("Systemic object marking is not enabled", str(cm.exception))
+
+    def test_no_systemic_object_marking(self):
+        with tracker.no_systemic_object_marking():
+            self.assertFalse(tracker.is_systemic())
+            t = empty_table(10)
+        self.assertFalse(t.j_object.isSystemicObject())
+
+
+if __name__ == '__main__':
+    unittest.main()