From e7f731bf322e2c61a1ed5f70fd34921bc599bf1c Mon Sep 17 00:00:00 2001 From: Jianfeng Mao <4297243+jmao-denver@users.noreply.github.com> Date: Mon, 20 Jan 2025 14:16:22 -0700 Subject: [PATCH] feat: DH-18423: add a systemic_obj_tracker module in the Py server API (#6577) Added a systemic_obj_tracker.py module which allows users to enable/disable systemic object marking on the current thread. --- .../SystemicObjectTracker.java | 3 +- py/server/deephaven/systemic_obj_tracker.py | 86 +++++++++++++++++++ py/server/test_helper/__init__.py | 1 + py/server/tests/test_systemic_obj_tracker.py | 58 +++++++++++++ 4 files changed, 147 insertions(+), 1 deletion(-) create mode 100644 py/server/deephaven/systemic_obj_tracker.py create mode 100644 py/server/tests/test_systemic_obj_tracker.py diff --git a/engine/updategraph/src/main/java/io/deephaven/engine/util/systemicmarking/SystemicObjectTracker.java b/engine/updategraph/src/main/java/io/deephaven/engine/util/systemicmarking/SystemicObjectTracker.java index cb351fe8e03..13ba20ae9ac 100644 --- a/engine/updategraph/src/main/java/io/deephaven/engine/util/systemicmarking/SystemicObjectTracker.java +++ b/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) { diff --git a/py/server/deephaven/systemic_obj_tracker.py b/py/server/deephaven/systemic_obj_tracker.py new file mode 100644 index 00000000000..27ac3521ac3 --- /dev/null +++ b/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.") diff --git a/py/server/test_helper/__init__.py b/py/server/test_helper/__init__.py index 434371c411e..84e4c8ad8fe 100644 --- a/py/server/test_helper/__init__.py +++ b/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: diff --git a/py/server/tests/test_systemic_obj_tracker.py b/py/server/tests/test_systemic_obj_tracker.py new file mode 100644 index 00000000000..cfc04f265ed --- /dev/null +++ b/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()