src/third_party/chromite/lib/timeout_util.py

   1 # Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 """Functions for implementing timeouts."""
   6
   7 from __future__ import print_function
   8
   9 import contextlib
  10 import functools
  11 import signal
  12 import time
  13
  14 from chromite.lib import signals
  15
  16
  17 class TimeoutError(Exception):
  18   """Raises when code within Timeout has been run too long."""
  19
  20
  21 @contextlib.contextmanager
  22 def Timeout(max_run_time,
  23             error_message="Timeout occurred- waited %(time)s seconds."):
  24   """ContextManager that alarms if code is ran for too long.
  25
  26   Timeout can run nested and raises a TimeoutException if the timeout
  27   is reached. Timeout can also nest underneath FatalTimeout.
  28
  29   Args:
  30     max_run_time: Number (integer) of seconds to wait before sending SIGALRM.
  31     error_message: String to wrap in the TimeoutError exception on timeout.
  32   """
  33   max_run_time = int(max_run_time)
  34   if max_run_time <= 0:
  35     raise ValueError("max_run_time must be greater than zero")
  36
  37   # pylint: disable=W0613
  38   def kill_us(sig_num, frame):
  39     raise TimeoutError(error_message % {'time': max_run_time})
  40
  41   original_handler = signal.signal(signal.SIGALRM, kill_us)
  42   previous_time = int(time.time())
  43
  44   # Signal the min in case the leftover time was smaller than this timeout.
  45   remaining_timeout = signal.alarm(0)
  46   if remaining_timeout:
  47     signal.alarm(min(remaining_timeout, max_run_time))
  48   else:
  49     signal.alarm(max_run_time)
  50
  51   try:
  52     yield
  53   finally:
  54     # Cancel the alarm request and restore the original handler.
  55     signal.alarm(0)
  56     signal.signal(signal.SIGALRM, original_handler)
  57
  58     # Ensure the previous handler will fire if it was meant to.
  59     if remaining_timeout > 0:
  60       # Signal the previous handler if it would have already passed.
  61       time_left = remaining_timeout - (int(time.time()) - previous_time)
  62       if time_left <= 0:
  63         signals.RelaySignal(original_handler, signal.SIGALRM, None)
  64       else:
  65         signal.alarm(time_left)
  66
  67
  68 @contextlib.contextmanager
  69 def FatalTimeout(max_run_time):
  70   """ContextManager that exits the program if code is run for too long.
  71
  72   This implementation is fairly simple, thus multiple timeouts
  73   cannot be active at the same time.
  74
  75   Additionally, if the timeout has elapsed, it'll trigger a SystemExit
  76   exception within the invoking code, ultimately propagating that past
  77   itself.  If the underlying code tries to suppress the SystemExit, once
  78   a minute it'll retrigger SystemExit until control is returned to this
  79   manager.
  80
  81   Args:
  82     max_run_time: a positive integer.
  83   """
  84   max_run_time = int(max_run_time)
  85   if max_run_time <= 0:
  86     raise ValueError("max_run_time must be greater than zero")
  87
  88   # pylint: disable=W0613
  89   def kill_us(sig_num, frame):
  90     # While this SystemExit *should* crash it's way back up the
  91     # stack to our exit handler, we do have live/production code
  92     # that uses blanket except statements which could suppress this.
  93     # As such, keep scheduling alarms until our exit handler runs.
  94     # Note that there is a potential conflict via this code, and
  95     # RunCommand's kill_timeout; thus we set the alarming interval
  96     # fairly high.
  97     signal.alarm(60)
  98     raise SystemExit("Timeout occurred- waited %i seconds, failing."
  99                      % max_run_time)
 100
 101   original_handler = signal.signal(signal.SIGALRM, kill_us)
 102   remaining_timeout = signal.alarm(max_run_time)
 103   if remaining_timeout:
 104     # Restore things to the way they were.
 105     signal.signal(signal.SIGALRM, original_handler)
 106     signal.alarm(remaining_timeout)
 107     # ... and now complain.  Unfortunately we can't easily detect this
 108     # upfront, thus the reset dance above.
 109     raise Exception("_Timeout cannot be used in parallel to other alarm "
 110                     "handling code; failing")
 111   try:
 112     yield
 113   finally:
 114     # Cancel the alarm request and restore the original handler.
 115     signal.alarm(0)
 116     signal.signal(signal.SIGALRM, original_handler)
 117
 118
 119 def TimeoutDecorator(max_time):
 120   """Decorator used to ensure a func is interrupted if it's running too long."""
 121   # Save off the built-in versions of time.time, signal.signal, and
 122   # signal.alarm, in case they get mocked out later. We want to ensure that
 123   # tests don't accidentally mock out the functions used by Timeout.
 124   def _Save():
 125     return time.time, signal.signal, signal.alarm
 126   def _Restore(values):
 127     (time.time, signal.signal, signal.alarm) = values
 128   builtins = _Save()
 129
 130   def NestedTimeoutDecorator(func):
 131     @functools.wraps(func)
 132     def TimeoutWrapper(*args, **kwargs):
 133       new = _Save()
 134       try:
 135         _Restore(builtins)
 136         with Timeout(max_time):
 137           _Restore(new)
 138           try:
 139             func(*args, **kwargs)
 140           finally:
 141             _Restore(builtins)
 142       finally:
 143         _Restore(new)
 144
 145     return TimeoutWrapper
 146
 147   return NestedTimeoutDecorator
 148
 149
 150 def WaitForReturnTrue(*args, **kwargs):
 151   """Periodically run a function, waiting in between runs.
 152
 153   Continues to run until the function returns True.
 154
 155   Args:
 156     See WaitForReturnValue([True], ...)
 157
 158   Raises:
 159     TimeoutError when the timeout is exceeded.
 160   """
 161   WaitForReturnValue([True], *args, **kwargs)
 162
 163
 164 def WaitForReturnValue(values, *args, **kwargs):
 165   """Periodically run a function, waiting in between runs.
 166
 167   Continues to run until the function return value is in the list
 168   of accepted |values|.  See WaitForSuccess for more details.
 169
 170   Args:
 171     values: A list or set of acceptable return values.
 172     *args, **kwargs: See WaitForSuccess for remaining arguments.
 173
 174   Returns:
 175     The value most recently returned by |func|.
 176
 177   Raises:
 178     TimeoutError when the timeout is exceeded.
 179   """
 180   def _Retry(return_value):
 181     return return_value not in values
 182
 183   return WaitForSuccess(_Retry, *args, **kwargs)
 184
 185
 186 def WaitForSuccess(retry_check, func, timeout, period=1, side_effect_func=None,
 187                    func_args=None, func_kwargs=None, fallback_timeout=10):
 188   """Periodically run a function, waiting in between runs.
 189
 190   Continues to run given function until return value is accepted by retry check.
 191
 192   To retry based on raised exceptions see GenericRetry in retry_util.
 193
 194   Args:
 195     retry_check: A functor that will be passed the return value of |func| as
 196       the only argument.  If |func| should be retried |retry_check| should
 197       return True.
 198     func: The function to run to test for a value.
 199     timeout: The maximum amount of time to wait, in integer seconds.
 200     period: Integer number of seconds between calls to |func|.
 201     side_effect_func: Optional function to be called between polls of func,
 202                       typically to output logging messages. The remaining
 203                       time in minutes will be passed as the first arg to
 204                       |side_effect_func|.
 205     func_args: Optional list of positional arguments to be passed to |func|.
 206     func_kwargs: Optional dictionary of keyword arguments to be passed to
 207                  |func|.
 208     fallback_timeout: We set a secondary timeout based on sigalarm this many
 209                       seconds after the initial timeout. This should NOT be
 210                       considered robust, but can allow timeouts inside blocking
 211                       methods.
 212
 213   Returns:
 214     The value most recently returned by |func| that was not flagged for retry.
 215
 216   Raises:
 217     TimeoutError when the timeout is exceeded.
 218   """
 219   assert period >= 0
 220   func_args = func_args or []
 221   func_kwargs = func_kwargs or {}
 222
 223   timeout_end = time.time() + timeout
 224
 225   # Use a sigalarm after an extra delay, in case a function we call is
 226   # blocking for some reason. This should NOT be considered reliable.
 227   with Timeout(timeout + fallback_timeout):
 228     while True:
 229       period_start = time.time()
 230       period_end = period_start + period
 231
 232       value = func(*func_args, **func_kwargs)
 233       if not retry_check(value):
 234         return value
 235
 236       if side_effect_func:
 237         # The remaining time may be negative. Add 0.5 minutes to
 238         # offset it so we don't return a negative value.
 239         side_effect_func((timeout_end - time.time()) / 60 + 0.5)
 240
 241       time_remaining = min(timeout_end, period_end) - time.time()
 242       if time_remaining > 0:
 243         time.sleep(time_remaining)
 244
 245       if time.time() >= timeout_end:
 246         raise TimeoutError('Timed out after %d seconds' % timeout)