Associate threads in reraiser_thread.RunAsync with the calling ThreadGroup

build/android/devil/utils/timeout_retry.py

Index: build/android/devil/utils/timeout_retry.py
diff --git a/build/android/devil/utils/timeout_retry.py b/build/android/devil/utils/timeout_retry.py
index e6c77f129ca24ce89d51f9aad17de05117d036c8..7e3c7efeb6568c3770a5c50c82b14bf3de67a32b 100644
--- a/build/android/devil/utils/timeout_retry.py
+++ b/build/android/devil/utils/timeout_retry.py
@@ -14,11 +14,11 @@ from devil.utils import reraiser_thread
 from devil.utils import watchdog_timer
 
 
-class TimeoutRetryThread(reraiser_thread.ReraiserThread):
-  def __init__(self, func, timeout, name):
-    super(TimeoutRetryThread, self).__init__(func, name=name)
+
+class TimeoutRetryThreadGroup(reraiser_thread.ReraiserThreadGroup):
+  def __init__(self, timeout, threads=None):
+    super(TimeoutRetryThreadGroup, self).__init__(threads)
     self._watcher = watchdog_timer.WatchdogTimer(timeout)
-    self._expired = False
 
   def GetWatcher(self):
     """Returns the watchdog keeping track of this thread's time."""
@@ -52,34 +52,22 @@ class TimeoutRetryThread(reraiser_thread.ReraiserThread):
       if remaining > 0:
         msg += (', wait of %.1f secs required but only %.1f secs left'
                 % (required, remaining))
-      self._expired = True
       raise reraiser_thread.TimeoutError(msg)
     return remaining
 
-  def LogTimeoutException(self):
-    """Log the exception that terminated this thread."""
-    if not self._expired:
-      return
-    logging.critical('*' * 80)
-    logging.critical('%s on thread %r', self._exc_info[0].__name__, self.name)
-    logging.critical('*' * 80)
-    fmt_exc = ''.join(traceback.format_exception(*self._exc_info))
-    for line in fmt_exc.splitlines():
-      logging.critical(line.rstrip())
-    logging.critical('*' * 80)
-
 
-def CurrentTimeoutThread():
-  """Get the current thread if it is a TimeoutRetryThread.
+def CurrentTimeoutThreadGroup():
+  """Returns the thread group that owns or is blocked on the active thread.
 
   Returns:
-    The current thread if it is a TimeoutRetryThread, otherwise None.
+    Returns None if no TimeoutRetryThreadGroup is tracking the current thread.
   """
-  current_thread = threading.current_thread()
-  if isinstance(current_thread, TimeoutRetryThread):
-    return current_thread
-  else:
-    return None
+  thread_group = reraiser_thread.CurrentThreadGroup()
+  while thread_group:
+    if isinstance(thread_group, TimeoutRetryThreadGroup):
+      return thread_group
+    thread_group = thread_group.blocked_parent_thread_group
+  return None
 
 
 def WaitFor(condition, wait_period=5, max_tries=None):
@@ -88,45 +76,57 @@ def WaitFor(condition, wait_period=5, max_tries=None):
   Repeatedly call the function condition(), with no arguments, until it returns
   a true value.
 
-  If called within a TimeoutRetryThread, it cooperates nicely with it.
+  If called within a TimeoutRetryThreadGroup, it cooperates nicely with it.
 
   Args:
     condition: function with the condition to check
     wait_period: number of seconds to wait before retrying to check the
       condition
     max_tries: maximum number of checks to make, the default tries forever
-      or until the TimeoutRetryThread expires.
+      or until the TimeoutRetryThreadGroup expires.
 
   Returns:
     The true value returned by the condition, or None if the condition was
     not met after max_tries.
 
   Raises:
-    reraiser_thread.TimeoutError if the current thread is a TimeoutRetryThread
-      and the timeout expires.
+    reraiser_thread.TimeoutError: if the current thread is a
+      TimeoutRetryThreadGroup and the timeout expires.
   """
   condition_name = condition.__name__
-  timeout_thread = CurrentTimeoutThread()
+  timeout_thread_group = CurrentTimeoutThreadGroup()
   while max_tries is None or max_tries > 0:
     result = condition()
     if max_tries is not None:
       max_tries -= 1
     msg = ['condition', repr(condition_name), 'met' if result else 'not met']
-    if timeout_thread:
+    if timeout_thread_group:
       # pylint: disable=no-member
-      msg.append('(%.1fs)' % timeout_thread.GetElapsedTime())
+      msg.append('(%.1fs)' % timeout_thread_group.GetElapsedTime())
     logging.info(' '.join(msg))
     if result:
       return result
-    if timeout_thread:
+    if timeout_thread_group:
       # pylint: disable=no-member
-      timeout_thread.GetRemainingTime(wait_period,
+      timeout_thread_group.GetRemainingTime(wait_period,
           msg='Timed out waiting for %r' % condition_name)
     time.sleep(wait_period)
   return None
 
 
-def Run(func, timeout, retries, args=None, kwargs=None, desc=None):
+def _LogLastException(thread_name, attempt, max_attempts, log_func):
+  log_func('*' * 80)
+  log_func('Exception on thread %s (attempt %d of %d)', thread_name,
+                   attempt, max_attempts)
+  log_func('*' * 80)
+  fmt_exc = ''.join(traceback.format_exc())
+  for line in fmt_exc.splitlines():
+    log_func(line.rstrip())
+  log_func('*' * 80)
+
+
+def Run(func, timeout, retries, args=None, kwargs=None, desc=None,
+        error_log_func=logging.critical):
   """Runs the passed function in a separate thread with timeouts and retries.
 
   Args:
@@ -137,6 +137,7 @@ def Run(func, timeout, retries, args=None, kwargs=None, desc=None):
     kwargs: dictionary of keyword args to pass to |func|.
     desc: An optional description of |func| used in logging. If omitted,
       |func.__name__| will be used.
+    error_log_func: Logging function when logging errors.
 
   Returns:
     The return value of func(*args, **kwargs).
@@ -146,30 +147,28 @@ def Run(func, timeout, retries, args=None, kwargs=None, desc=None):
   if not kwargs:
     kwargs = {}
 
-  # The return value uses a list because Python variables are references, not
-  # values. Closures make a copy of the reference, so updating the closure's
-  # reference wouldn't update where the original reference pointed.
-  ret = [None]
-  def RunOnTimeoutThread():
-    ret[0] = func(*args, **kwargs)
-
   num_try = 1
   while True:
-    child_thread = TimeoutRetryThread(
-      RunOnTimeoutThread, timeout,
-      name='TimeoutThread-%d-for-%s' % (num_try,
-                                        threading.current_thread().name))
+    thread_name = 'TimeoutThread-%d-for-%s' % (num_try,
+                                               threading.current_thread().name)
+    child_thread = reraiser_thread.ReraiserThread(lambda: func(*args, **kwargs),
+                                                  name=thread_name)
     try:
-      thread_group = reraiser_thread.ReraiserThreadGroup([child_thread])
+      thread_group = TimeoutRetryThreadGroup(timeout, threads=[child_thread])
       thread_group.StartAll()

[jbudorick, 2015/10/01 19:28:44]

Should this set will_block=True? We join immediately after...

[agrieve, 2015/10/01 19:44:27]

Added.

I don't think the information will be used in this case though, since
CurrentTimeoutThreadGroup returns the first-hit group.

[jbudorick, 2015/10/01 20:37:12]

Ah. I wasn't sure.

       while True:
-        thread_group.JoinAll(watcher=child_thread.GetWatcher(), timeout=60)
+        thread_group.JoinAll(watcher=thread_group.GetWatcher(), timeout=60,
+                             error_log_func=error_log_func)
         if thread_group.IsAlive():
           logging.info('Still working on %s', desc if desc else func.__name__)
         else:
-          return ret[0]
+          return thread_group.GetAllReturnValues()[0]
+    except reraiser_thread.TimeoutError:
+      # Timeouts already get their stacks logged.

[jbudorick, 2015/10/01 19:28:44]

though the stack is only logged once, after all tries.

[agrieve, 2015/10/01 19:44:27]

The error_log_func I added just so that scary stacks wouldn't appear when
running unit tests.

This comment I added, but I've tried to match the existing behaviour (which was
accomplished by mutating self._expired from within a getter...). I don't think
it's a big deal to also re-log timeouts here, so if that's what you're getting
at, then I can delete this special case.

[jbudorick, 2015/10/01 20:37:12]

Definitely on board with not logging scary stacks when they aren't necessary.
I just meant that this was previously logging those timeouts on each try. It now
will only log a timeout at the end. I think that's ok, I just wanted to point it
out.

[agrieve, 2015/10/02 00:19:44]

Ahaaaa, okay.

This change I did on purpose, but it's not quite as you say. It's logging them
every time *except* the end.

The exception that is rethrown should be logged / handled by the calling
function.

+      if num_try > retries:
+        raise
     except:

[jbudorick, 2015/10/01 20:37:12]

nit: I think this should be Exception s.t. we at least don't catch and retry on
KeyboardInterrupt. (We haven't done this elsewhere and we weren't doing it
previously here, obviously, but I've had telemetry folks point it out to me on
other reviews and it seems like a good idea.)

[agrieve, 2015/10/02 00:19:45]

Agree! Done.

-      child_thread.LogTimeoutException()
       if num_try > retries:
         raise
-      num_try += 1
+      _LogLastException(thread_name, num_try, retries + 1, error_log_func)

[jbudorick, 2015/10/01 20:37:12]

optional nit: should we catch the exception with a variable and pass it in
rather than just pulling from traceback? If multiple threads throw exceptions in
separate calls to timeout_retry.Run, you could conceivably wind up logging a
different exception than the one you caught.

(i think)

[agrieve, 2015/10/02 00:19:45]

I would much prefer passing in the exception, but I don't actually see a python
API that lets you do that :(.

I suspect traceback.format_exc() is thread-local, so don't think there's any bug
at least.

+    num_try += 1