Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
86 changes: 86 additions & 0 deletions py/server/deephaven/systemic_obj_tracker.py
Original file line number Diff line number Diff line change
@@ -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.")
1 change: 1 addition & 0 deletions py/server/test_helper/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
58 changes: 58 additions & 0 deletions py/server/tests/test_systemtic_obj_tracker.py
Original file line number Diff line number Diff line change
@@ -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):
jmao-denver marked this conversation as resolved.
Show resolved Hide resolved

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()