Skip to content

Commit ddc3080

Browse files
lukebaumanncopybara-github
lukebaumann
authored and
copybara-github
committed
Adding debugging utilities
PiperOrigin-RevId: 743287943
1 parent a22f7ce commit ddc3080

File tree

6 files changed

+367
-0
lines changed

6 files changed

+367
-0
lines changed

‎pathwaysutils/debug/__init__.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.

‎pathwaysutils/debug/timing.py‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
"""Timing utilities.
15+
16+
This module provides utilities for timing code with decorators and context
17+
managers.
18+
"""
19+
20+
import functools
21+
import logging
22+
import time
23+
from typing import Any, Callable
24+
25+
26+
_logger = logging.getLogger(__name__)
27+
28+
29+
class Timer:
30+
"""Timer context manager.
31+
32+
Attributes:
33+
name: The name of the timer.
34+
start: The start time of the timer.
35+
stop: The stop time of the timer.
36+
duration: The elapsed time of the timer.
37+
"""
38+
39+
def __init__(self, name: str):
40+
self.name = name
41+
42+
def __enter__(self):
43+
self.start = time.time()
44+
return self
45+
46+
def __exit__(self, exc_type, exc_value, tb):
47+
self.stop = time.time()
48+
self.duration = self.stop - self.start
49+
_logger.debug(str(self))
50+
51+
def __str__(self):
52+
return f"{self.name} elapsed {self.duration:.4f} seconds."
53+
54+
55+
def timeit(func: Callable[..., Any]) -> Callable[..., Any]:
56+
"""Decorator to time a function.
57+
58+
Args:
59+
func: The function to time.
60+
61+
Returns:
62+
The decorated function.
63+
"""
64+
65+
@functools.wraps(func)
66+
def wrapper(*args: Any, **kwargs: Any):
67+
with Timer(getattr(func, "__name__", "Unknown")):
68+
return func(*args, **kwargs)
69+
70+
return wrapper

‎pathwaysutils/debug/watchdog.py‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
"""Watchdog context manager.
15+
16+
This context manager is used to monitor the progress of a long running process.
17+
If the process takes longer than the specified timeout, it will print the stack
18+
trace of all threads.
19+
"""
20+
21+
import contextlib
22+
import logging
23+
import os
24+
import sys
25+
import threading
26+
import traceback
27+
28+
29+
_logger = logging.getLogger(__name__)
30+
31+
32+
def _log_thread_stack(thread: threading.Thread):
33+
_logger.debug("Thread: %s", thread.ident)
34+
_logger.debug(
35+
"".join(
36+
traceback.format_stack(
37+
sys._current_frames().get( # pylint: disable=protected-access
38+
thread.ident, []
39+
)
40+
)
41+
)
42+
)
43+
44+
45+
@contextlib.contextmanager
46+
def watchdog(timeout: float, repeat: bool = True):
47+
"""Watchdog context manager.
48+
49+
Prints the stack trace of all threads after `timeout` seconds.
50+
51+
Args:
52+
timeout: The timeout in seconds. If the timeout is reached, the stack trace
53+
of all threads will be printed.
54+
repeat: Whether to repeat the watchdog after the timeout. If False, the
55+
process will be aborted after the first timeout.
56+
57+
Yields:
58+
None
59+
"""
60+
event = threading.Event()
61+
62+
def handler():
63+
count = 0
64+
while not event.wait(timeout):
65+
_logger.debug(
66+
"Watchdog thread dump every %s seconds. Count: %s", timeout, count
67+
)
68+
try:
69+
for thread in threading.enumerate():
70+
try:
71+
_log_thread_stack(thread)
72+
except Exception: # pylint: disable=broad-exception-caught
73+
_logger.debug("Error print traceback for thread: %s", thread.ident)
74+
finally:
75+
if not repeat:
76+
_logger.critical("Timeout from watchdog!")
77+
os.abort()
78+
79+
count += 1
80+
81+
_logger.debug("Registering watchdog")
82+
watchdog_thread = threading.Thread(target=handler, name="watchdog")
83+
watchdog_thread.start()
84+
try:
85+
yield
86+
finally:
87+
event.set()
88+
watchdog_thread.join()
89+
_logger.debug("Deregistering watchdog")

‎pathwaysutils/test/debug/__init__.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.

‎pathwaysutils/test/debug/timing_test.py‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
"""Timing tests."""
15+
16+
import logging
17+
import time
18+
from unittest import mock
19+
20+
from pathwaysutils.debug import timing
21+
22+
from absl.testing import absltest
23+
from absl.testing import parameterized
24+
25+
26+
class TimingTest(parameterized.TestCase):
27+
28+
def test_timer_context_manager(self):
29+
with mock.patch.object(
30+
time,
31+
"time",
32+
side_effect=[1, 8.9],
33+
autospec=True,
34+
):
35+
with timing.Timer("test_timer") as timer:
36+
pass
37+
38+
self.assertEqual(timer.name, "test_timer")
39+
self.assertEqual(timer.start, 1)
40+
self.assertEqual(timer.stop, 8.9)
41+
self.assertEqual(timer.duration, 7.9)
42+
self.assertEqual(str(timer), "test_timer elapsed 7.9000 seconds.")
43+
44+
def test_timeit_log(self):
45+
46+
@timing.timeit
47+
def my_function():
48+
pass
49+
50+
with mock.patch.object(
51+
time,
52+
"time",
53+
side_effect=[1, 8.9, 0], # Third time is used for logging.
54+
autospec=True,
55+
):
56+
with self.assertLogs(timing._logger, logging.DEBUG) as log_output:
57+
my_function()
58+
59+
self.assertEqual(
60+
log_output.output,
61+
[
62+
"DEBUG:pathwaysutils.debug.timing:my_function"
63+
" elapsed 7.9000 seconds."
64+
],
65+
)
66+
67+
def test_timeit_return_value(self):
68+
69+
@timing.timeit
70+
def my_function():
71+
return "test"
72+
73+
self.assertEqual(my_function(), "test")
74+
75+
def test_timeit_exception(self):
76+
77+
@timing.timeit
78+
def my_function():
79+
raise ValueError("test")
80+
81+
with self.assertRaises(ValueError):
82+
my_function()
83+
84+
85+
if __name__ == "__main__":
86+
absltest.main()

‎pathwaysutils/test/debug/watchdog_test.py‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
"""Watchdog tests."""
15+
16+
import logging
17+
import sys
18+
import threading
19+
import traceback
20+
from unittest import mock
21+
22+
from pathwaysutils.debug import watchdog
23+
24+
from absl.testing import absltest
25+
from absl.testing import parameterized
26+
27+
28+
class WatchdogTest(parameterized.TestCase):
29+
def test_watchdog_start_join(self):
30+
with (
31+
mock.patch.object(
32+
threading.Thread,
33+
"start",
34+
autospec=True,
35+
) as mock_start,
36+
mock.patch.object(threading.Thread, "join", autospec=True) as mock_join,
37+
):
38+
with watchdog.watchdog(timeout=1):
39+
mock_start.assert_called_once()
40+
mock_join.assert_not_called()
41+
42+
mock_start.assert_called_once()
43+
mock_join.assert_called_once()
44+
45+
@parameterized.named_parameters([
46+
(
47+
"thread 1",
48+
1,
49+
[
50+
"DEBUG:pathwaysutils.debug.watchdog:Thread: 1",
51+
"DEBUG:pathwaysutils.debug.watchdog:examplestack1",
52+
],
53+
),
54+
(
55+
"thread 2",
56+
2,
57+
[
58+
"DEBUG:pathwaysutils.debug.watchdog:Thread: 2",
59+
"DEBUG:pathwaysutils.debug.watchdog:examplestack2",
60+
],
61+
),
62+
(
63+
"thread 3",
64+
3,
65+
[
66+
"DEBUG:pathwaysutils.debug.watchdog:Thread: 3",
67+
"DEBUG:pathwaysutils.debug.watchdog:",
68+
],
69+
),
70+
])
71+
def test_log_thread_strack_succes(self, thread_ident, expected_log_output):
72+
with (
73+
mock.patch.object(
74+
sys,
75+
"_current_frames",
76+
return_value={1: ["example", "stack1"], 2: ["example", "stack2"]},
77+
autospec=True,
78+
),
79+
mock.patch.object(
80+
traceback,
81+
"format_stack",
82+
side_effect=lambda stack_str_list: stack_str_list,
83+
autospec=True,
84+
),
85+
):
86+
mock_thread = mock.create_autospec(threading.Thread, instance=True)
87+
mock_thread.ident = thread_ident
88+
89+
with self.assertLogs(watchdog._logger, logging.DEBUG) as log_output:
90+
watchdog._log_thread_stack(mock_thread)
91+
92+
self.assertEqual(log_output.output, expected_log_output)
93+
94+
95+
if __name__ == "__main__":
96+
absltest.main()

0 commit comments

Comments
 (0)