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
   8 import contextlib
   9 import functools
  10 import logging
  11 import json
  12 import signal
  13 import time
  14 import urllib
  15
  16 from chromite.buildbot import constants
  17 from chromite.lib import signals
  18
  19
  20 class TimeoutError(Exception):
  21   """Raises when code within Timeout has been run too long."""
  22
  23
  24 @contextlib.contextmanager
  25 def Timeout(max_run_time):
  26   """ContextManager that alarms if code is ran for too long.
  27
  28   Timeout can run nested and raises a TimeoutException if the timeout
  29   is reached. Timeout can also nest underneath FatalTimeout.
  30
  31   Args:
  32     max_run_time: Number (integer) of seconds to wait before sending SIGALRM.
  33   """
  34   max_run_time = int(max_run_time)
  35   if max_run_time <= 0:
  36     raise ValueError("max_run_time must be greater than zero")
  37
  38   # pylint: disable=W0613
  39   def kill_us(sig_num, frame):
  40     raise TimeoutError("Timeout occurred- waited %s seconds." % max_run_time)
  41
  42   original_handler = signal.signal(signal.SIGALRM, kill_us)
  43   previous_time = int(time.time())
  44
  45   # Signal the min in case the leftover time was smaller than this timeout.
  46   remaining_timeout = signal.alarm(0)
  47   if remaining_timeout:
  48     signal.alarm(min(remaining_timeout, max_run_time))
  49   else:
  50     signal.alarm(max_run_time)
  51
  52   try:
  53     yield
  54   finally:
  55     # Cancel the alarm request and restore the original handler.
  56     signal.alarm(0)
  57     signal.signal(signal.SIGALRM, original_handler)
  58
  59     # Ensure the previous handler will fire if it was meant to.
  60     if remaining_timeout > 0:
  61       # Signal the previous handler if it would have already passed.
  62       time_left = remaining_timeout - (int(time.time()) - previous_time)
  63       if time_left <= 0:
  64         signals.RelaySignal(original_handler, signal.SIGALRM, None)
  65       else:
  66         signal.alarm(time_left)
  67
  68
  69 @contextlib.contextmanager
  70 def FatalTimeout(max_run_time):
  71   """ContextManager that exits the program if code is run for too long.
  72
  73   This implementation is fairly simple, thus multiple timeouts
  74   cannot be active at the same time.
  75
  76   Additionally, if the timeout has elapsed, it'll trigger a SystemExit
  77   exception within the invoking code, ultimately propagating that past
  78   itself.  If the underlying code tries to suppress the SystemExit, once
  79   a minute it'll retrigger SystemExit until control is returned to this
  80   manager.
  81
  82   Args:
  83     max_run_time: a positive integer.
  84   """
  85   max_run_time = int(max_run_time)
  86   if max_run_time <= 0:
  87     raise ValueError("max_run_time must be greater than zero")
  88
  89   # pylint: disable=W0613
  90   def kill_us(sig_num, frame):
  91     # While this SystemExit *should* crash it's way back up the
  92     # stack to our exit handler, we do have live/production code
  93     # that uses blanket except statements which could suppress this.
  94     # As such, keep scheduling alarms until our exit handler runs.
  95     # Note that there is a potential conflict via this code, and
  96     # RunCommand's kill_timeout; thus we set the alarming interval
  97     # fairly high.
  98     signal.alarm(60)
  99     raise SystemExit("Timeout occurred- waited %i seconds, failing."
 100                      % max_run_time)
 101
 102   original_handler = signal.signal(signal.SIGALRM, kill_us)
 103   remaining_timeout = signal.alarm(max_run_time)
 104   if remaining_timeout:
 105     # Restore things to the way they were.
 106     signal.signal(signal.SIGALRM, original_handler)
 107     signal.alarm(remaining_timeout)
 108     # ... and now complain.  Unfortunately we can't easily detect this
 109     # upfront, thus the reset dance above.
 110     raise Exception("_Timeout cannot be used in parallel to other alarm "
 111                     "handling code; failing")
 112   try:
 113     yield
 114   finally:
 115     # Cancel the alarm request and restore the original handler.
 116     signal.alarm(0)
 117     signal.signal(signal.SIGALRM, original_handler)
 118
 119
 120 def TimeoutDecorator(max_time):
 121   """Decorator used to ensure a func is interrupted if it's running too long."""
 122   # Save off the built-in versions of time.time, signal.signal, and
 123   # signal.alarm, in case they get mocked out later. We want to ensure that
 124   # tests don't accidentally mock out the functions used by Timeout.
 125   def _Save():
 126     return time.time, signal.signal, signal.alarm
 127   def _Restore(values):
 128     (time.time, signal.signal, signal.alarm) = values
 129   builtins = _Save()
 130
 131   def NestedTimeoutDecorator(func):
 132     @functools.wraps(func)
 133     def TimeoutWrapper(*args, **kwargs):
 134       new = _Save()
 135       try:
 136         _Restore(builtins)
 137         with Timeout(max_time):
 138           _Restore(new)
 139           try:
 140             func(*args, **kwargs)
 141           finally:
 142             _Restore(builtins)
 143       finally:
 144         _Restore(new)
 145
 146     return TimeoutWrapper
 147
 148   return NestedTimeoutDecorator
 149
 150
 151 def WaitForReturnTrue(*args, **kwargs):
 152   """Periodically run a function, waiting in between runs.
 153
 154   Continues to run until the function returns True.
 155
 156   Args:
 157     See WaitForReturnValue([True], ...)
 158
 159   Raises:
 160     TimeoutError when the timeout is exceeded.
 161   """
 162   WaitForReturnValue([True], *args, **kwargs)
 163
 164
 165 def WaitForReturnValue(values, *args, **kwargs):
 166   """Periodically run a function, waiting in between runs.
 167
 168   Continues to run until the function return value is in the list
 169   of accepted |values|.  See WaitForSuccess for more details.
 170
 171   Args:
 172     values: A list or set of acceptable return values.
 173     *args, **kwargs: See WaitForSuccess for remaining arguments.
 174
 175   Returns:
 176     The value most recently returned by |func|.
 177
 178   Raises:
 179     TimeoutError when the timeout is exceeded.
 180   """
 181   def _Retry(return_value):
 182     return return_value not in values
 183
 184   return WaitForSuccess(_Retry, *args, **kwargs)
 185
 186
 187 def WaitForSuccess(retry_check, func, timeout, period=1, side_effect_func=None,
 188                    func_args=None, func_kwargs=None,
 189                    fallback_timeout=10):
 190   """Periodically run a function, waiting in between runs.
 191
 192   Continues to run given function until return value is accepted by retry check.
 193
 194   To retry based on raised exceptions see GenericRetry in retry_util.
 195
 196   Args:
 197     retry_check: A functor that will be passed the return value of |func| as
 198       the only argument.  If |func| should be retried |retry_check| should
 199       return True.
 200     func: The function to run to test for a value.
 201     timeout: The maximum amount of time to wait, in integer seconds.
 202     period: Integer number of seconds between calls to |func|.
 203     side_effect_func: Optional function to be called between polls of func,
 204                       typically to output logging messages.
 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         side_effect_func()
 238
 239       time_remaining = min(timeout_end, period_end) - time.time()
 240       if time_remaining > 0:
 241         time.sleep(time_remaining)
 242
 243       if time.time() >= timeout_end:
 244         raise TimeoutError('Timed out after %d seconds' % timeout)
 245
 246
 247 def _GetStatus(status_url):
 248   """Polls |status_url| and returns the retrieved tree status.
 249
 250   This function gets a JSON response from |status_url|, and returns the
 251   value associated with the 'general_state' key, if one exists and the
 252   http request was successful.
 253
 254   Returns:
 255     The tree status, as a string, if it was successfully retrieved. Otherwise
 256     None.
 257   """
 258   try:
 259     # Check for successful response code.
 260     response = urllib.urlopen(status_url)
 261     if response.getcode() == 200:
 262       data = json.load(response)
 263       if data.has_key('general_state'):
 264         return data['general_state']
 265   # We remain robust against IOError's.
 266   except IOError as e:
 267     logging.error('Could not reach %s: %r', status_url, e)
 268
 269
 270 def WaitForTreeStatus(status_url, period=1, timeout=1, throttled_ok=False):
 271   """Wait for tree status to be open (or throttled, if |throttled_ok|).
 272
 273   Args:
 274     status_url: The status url to check i.e.
 275                 'https://status.appspot.com/current?format=json'
 276     period: How often to poll for status updates.
 277     timeout: How long to wait until a tree status is discovered.
 278     throttled_ok: is TREE_THROTTLED an acceptable status?
 279
 280   Returns:
 281     The most recent tree status, either constants.TREE_OPEN or
 282     constants.TREE_THROTTLED (if |throttled_ok|)
 283
 284   Raises:
 285     TimeoutError if timeout expired before tree reached acceptable status.
 286   """
 287   acceptable_states = set([constants.TREE_OPEN])
 288   verb = 'open'
 289   if throttled_ok:
 290     acceptable_states.add(constants.TREE_THROTTLED)
 291     verb = 'not be closed'
 292
 293   timeout = max(timeout, 1)
 294
 295   end_time = time.time() + timeout
 296
 297   def _LogMessage():
 298     time_left = end_time - time.time()
 299     logging.info('Waiting for the tree to %s (%d minutes left)...', verb,
 300                  time_left / 60)
 301
 302   def _get_status():
 303     return _GetStatus(status_url)
 304
 305   return WaitForReturnValue(acceptable_states, _get_status, timeout=timeout,
 306                             period=period, side_effect_func=_LogMessage)
 307
 308
 309
 310 def IsTreeOpen(status_url, period=1, timeout=1, throttled_ok=False):
 311   """Wait for tree status to be open (or throttled, if |throttled_ok|).
 312
 313   Args:
 314     status_url: The status url to check i.e.
 315                 'https://status.appspot.com/current?format=json'
 316     period: How often to poll for status updates.
 317     timeout: How long to wait until a tree status is discovered.
 318     throttled_ok: Does TREE_THROTTLED count as open?
 319
 320   Returns:
 321     True if the tree is open (or throttled, if |throttled_ok|). False if
 322     timeout expired before tree reached acceptable status.
 323   """
 324   try:
 325     WaitForTreeStatus(status_url, period, timeout, throttled_ok)
 326   except TimeoutError:
 327     return False
 328   return True
 329
 330
 331 def GetTreeStatus(status_url, polling_period=0, timeout=0):
 332   """Returns the current tree status as fetched from |status_url|.
 333
 334   This function returns the tree status as a string, either
 335   constants.TREE_OPEN, constants.TREE_THROTTLED, or constants.TREE_CLOSED.
 336
 337   Args:
 338     status_url: The status url to check i.e.
 339       'https://status.appspot.com/current?format=json'
 340     polling_period: Time to wait in seconds between polling attempts.
 341     timeout: Maximum time in seconds to wait for status.
 342
 343   Returns:
 344     constants.TREE_OPEN, constants.TREE_THROTTLED, or constants.TREE_CLOSED
 345
 346   Raises:
 347     TimeoutError if the timeout expired before the status could be successfully
 348     fetched.
 349   """
 350   acceptable_states = set([constants.TREE_OPEN, constants.TREE_THROTTLED,
 351                            constants.TREE_CLOSED])
 352
 353   def _get_status():
 354     return _GetStatus(status_url)
 355
 356   return WaitForReturnValue(acceptable_states, _get_status, timeout,
 357                             polling_period)