diff --git a/pathwaysutils/debug/__init__.py b/pathwaysutils/debug/__init__.py new file mode 100644 index 0000000..86550f5 --- /dev/null +++ b/pathwaysutils/debug/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. \ No newline at end of file diff --git a/pathwaysutils/debug/timing.py b/pathwaysutils/debug/timing.py new file mode 100644 index 0000000..9991c72 --- /dev/null +++ b/pathwaysutils/debug/timing.py @@ -0,0 +1,70 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Timing utilities. + +This module provides utilities for timing code with decorators and context +managers. +""" + +import functools +import logging +import time +from typing import Any, Callable + + +_logger = logging.getLogger(__name__) + + +class Timer: + """Timer context manager. + + Attributes: + name: The name of the timer. + start: The start time of the timer. + stop: The stop time of the timer. + duration: The elapsed time of the timer. + """ + + def __init__(self, name: str): + self.name = name + + def __enter__(self): + self.start = time.time() + return self + + def __exit__(self, exc_type, exc_value, tb): + self.stop = time.time() + self.duration = self.stop - self.start + _logger.debug(str(self)) + + def __str__(self): + return f"{self.name} elapsed {self.duration:.4f} seconds." + + +def timeit(func: Callable[..., Any]) -> Callable[..., Any]: + """Decorator to time a function. + + Args: + func: The function to time. + + Returns: + The decorated function. + """ + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any): + with Timer(getattr(func, "__name__", "Unknown")): + return func(*args, **kwargs) + + return wrapper diff --git a/pathwaysutils/debug/watchdog.py b/pathwaysutils/debug/watchdog.py new file mode 100644 index 0000000..5820e1e --- /dev/null +++ b/pathwaysutils/debug/watchdog.py @@ -0,0 +1,89 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Watchdog context manager. + +This context manager is used to monitor the progress of a long running process. +If the process takes longer than the specified timeout, it will print the stack +trace of all threads. +""" + +import contextlib +import logging +import os +import sys +import threading +import traceback + + +_logger = logging.getLogger(__name__) + + +def _log_thread_stack(thread: threading.Thread): + _logger.debug("Thread: %s", thread.ident) + _logger.debug( + "".join( + traceback.format_stack( + sys._current_frames().get( # pylint: disable=protected-access + thread.ident, [] + ) + ) + ) + ) + + +@contextlib.contextmanager +def watchdog(timeout: float, repeat: bool = True): + """Watchdog context manager. + + Prints the stack trace of all threads after `timeout` seconds. + + Args: + timeout: The timeout in seconds. If the timeout is reached, the stack trace + of all threads will be printed. + repeat: Whether to repeat the watchdog after the timeout. If False, the + process will be aborted after the first timeout. + + Yields: + None + """ + event = threading.Event() + + def handler(): + count = 0 + while not event.wait(timeout): + _logger.debug( + "Watchdog thread dump every %s seconds. Count: %s", timeout, count + ) + try: + for thread in threading.enumerate(): + try: + _log_thread_stack(thread) + except Exception: # pylint: disable=broad-exception-caught + _logger.debug("Error print traceback for thread: %s", thread.ident) + finally: + if not repeat: + _logger.critical("Timeout from watchdog!") + os.abort() + + count += 1 + + _logger.debug("Registering watchdog") + watchdog_thread = threading.Thread(target=handler, name="watchdog") + watchdog_thread.start() + try: + yield + finally: + event.set() + watchdog_thread.join() + _logger.debug("Deregistering watchdog") diff --git a/pathwaysutils/test/debug/__init__.py b/pathwaysutils/test/debug/__init__.py new file mode 100644 index 0000000..86550f5 --- /dev/null +++ b/pathwaysutils/test/debug/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. \ No newline at end of file diff --git a/pathwaysutils/test/debug/timing_test.py b/pathwaysutils/test/debug/timing_test.py new file mode 100644 index 0000000..a3278d7 --- /dev/null +++ b/pathwaysutils/test/debug/timing_test.py @@ -0,0 +1,86 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Timing tests.""" + +import logging +import time +from unittest import mock + +from pathwaysutils.debug import timing + +from absl.testing import absltest +from absl.testing import parameterized + + +class TimingTest(parameterized.TestCase): + + def test_timer_context_manager(self): + with mock.patch.object( + time, + "time", + side_effect=[1, 8.9], + autospec=True, + ): + with timing.Timer("test_timer") as timer: + pass + + self.assertEqual(timer.name, "test_timer") + self.assertEqual(timer.start, 1) + self.assertEqual(timer.stop, 8.9) + self.assertEqual(timer.duration, 7.9) + self.assertEqual(str(timer), "test_timer elapsed 7.9000 seconds.") + + def test_timeit_log(self): + + @timing.timeit + def my_function(): + pass + + with mock.patch.object( + time, + "time", + side_effect=[1, 8.9, 0], # Third time is used for logging. + autospec=True, + ): + with self.assertLogs(timing._logger, logging.DEBUG) as log_output: + my_function() + + self.assertEqual( + log_output.output, + [ + "DEBUG:pathwaysutils.debug.timing:my_function" + " elapsed 7.9000 seconds." + ], + ) + + def test_timeit_return_value(self): + + @timing.timeit + def my_function(): + return "test" + + self.assertEqual(my_function(), "test") + + def test_timeit_exception(self): + + @timing.timeit + def my_function(): + raise ValueError("test") + + with self.assertRaises(ValueError): + my_function() + + +if __name__ == "__main__": + absltest.main() diff --git a/pathwaysutils/test/debug/watchdog_test.py b/pathwaysutils/test/debug/watchdog_test.py new file mode 100644 index 0000000..8ad0352 --- /dev/null +++ b/pathwaysutils/test/debug/watchdog_test.py @@ -0,0 +1,96 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Watchdog tests.""" + +import logging +import sys +import threading +import traceback +from unittest import mock + +from pathwaysutils.debug import watchdog + +from absl.testing import absltest +from absl.testing import parameterized + + +class WatchdogTest(parameterized.TestCase): + def test_watchdog_start_join(self): + with ( + mock.patch.object( + threading.Thread, + "start", + autospec=True, + ) as mock_start, + mock.patch.object(threading.Thread, "join", autospec=True) as mock_join, + ): + with watchdog.watchdog(timeout=1): + mock_start.assert_called_once() + mock_join.assert_not_called() + + mock_start.assert_called_once() + mock_join.assert_called_once() + + @parameterized.named_parameters([ + ( + "thread 1", + 1, + [ + "DEBUG:pathwaysutils.debug.watchdog:Thread: 1", + "DEBUG:pathwaysutils.debug.watchdog:examplestack1", + ], + ), + ( + "thread 2", + 2, + [ + "DEBUG:pathwaysutils.debug.watchdog:Thread: 2", + "DEBUG:pathwaysutils.debug.watchdog:examplestack2", + ], + ), + ( + "thread 3", + 3, + [ + "DEBUG:pathwaysutils.debug.watchdog:Thread: 3", + "DEBUG:pathwaysutils.debug.watchdog:", + ], + ), + ]) + def test_log_thread_strack_succes(self, thread_ident, expected_log_output): + with ( + mock.patch.object( + sys, + "_current_frames", + return_value={1: ["example", "stack1"], 2: ["example", "stack2"]}, + autospec=True, + ), + mock.patch.object( + traceback, + "format_stack", + side_effect=lambda stack_str_list: stack_str_list, + autospec=True, + ), + ): + mock_thread = mock.create_autospec(threading.Thread, instance=True) + mock_thread.ident = thread_ident + + with self.assertLogs(watchdog._logger, logging.DEBUG) as log_output: + watchdog._log_thread_stack(mock_thread) + + self.assertEqual(log_output.output, expected_log_output) + + +if __name__ == "__main__": + absltest.main()