Add a simple signal-based Python callstack sampler for debugging

official/utils/misc/callstack_sampler.py

+"""A simple Python callstack sampler."""
+
+import contextlib
+import datetime
+import signal
+import traceback
+
+
+class CallstackSampler(object):
+  """A simple signal-based Python callstack sampler.
+  """
+
+  def __init__(self, interval=None):
+    self.stacks = []
+    self.interval = 0.001 if interval is None else interval
+
+  def _sample(self, signum, frame):
+    """Samples the current stack."""
+    del signum
+    stack = traceback.extract_stack(frame)
+    formatted_stack = []
+    formatted_stack.append(datetime.datetime.utcnow())
+    for filename, lineno, function_name, text in stack:
+      formatted_frame = '{}:{}({})({})'.format(filename, lineno, function_name,
+                                               text)
+      formatted_stack.append(formatted_frame)
+    self.stacks.append(formatted_stack)
+    signal.setitimer(signal.ITIMER_VIRTUAL, self.interval, 0)
+
+  @contextlib.contextmanager
+  def profile(self):
+    signal.signal(signal.SIGVTALRM, self._sample)
+    signal.setitimer(signal.ITIMER_VIRTUAL, self.interval, 0)
+    try:
+      yield
+    finally:
+      signal.setitimer(signal.ITIMER_VIRTUAL, 0)
+
+  def save(self, fname):
+    with open(fname, 'w') as f:
+      for s in self.stacks:
+        for l in s:
+          f.write('%s\n' % l)
+        f.write('\n')
+
+
+@contextlib.contextmanager
+def callstack_sampling(filename, interval=None):
+  """Periodically samples the Python callstack.
+
+  Args:
+    filename: the filename
+    interval: the sampling interval, in seconds. Defaults to 0.001.
+
+  Yields:
+   nothing
+  """
+  sampler = CallstackSampler(interval=interval)
+  with sampler.profile():
+    yield
+  sampler.save(filename)
+