Home | History | Annotate | Download | only in cros
      1 # Copyright (c) 2012 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 import logging
      6 import random
      7 import signal
      8 import sys
      9 import threading
     10 import time
     11 
     12 from autotest_lib.client.common_lib import error
     13 
     14 
     15 def install_sigalarm_handler(new_handler):
     16     """
     17     Try installing a sigalarm handler.
     18 
     19     In order to protect apache, wsgi intercepts any attempt to install a
     20     sigalarm handler, so our function will feel the full force of a sigalarm
     21     even if we try to install a pacifying signal handler. To avoid this we
     22     need to confirm that the handler we tried to install really was installed.
     23 
     24     @param new_handler: The new handler to install. This must be a callable
     25                         object, or signal.SIG_IGN/SIG_DFL which correspond to
     26                         the numbers 1,0 respectively.
     27     @return: True if the installation of new_handler succeeded, False otherwise.
     28     """
     29     if (new_handler is None or
     30         (not callable(new_handler) and
     31          new_handler != signal.SIG_IGN and
     32          new_handler != signal.SIG_DFL)):
     33         logging.warning('Trying to install an invalid sigalarm handler.')
     34         return False
     35 
     36     signal.signal(signal.SIGALRM, new_handler)
     37     installed_handler = signal.getsignal(signal.SIGALRM)
     38     return installed_handler == new_handler
     39 
     40 
     41 def set_sigalarm_timeout(timeout_secs, default_timeout=60):
     42     """
     43     Set the sigalarm timeout.
     44 
     45     This methods treats any timeout <= 0 as a possible error and falls back to
     46     using it's default timeout, since negative timeouts can have 'alarming'
     47     effects. Though 0 is a valid timeout, it is often used to cancel signals; in
     48     order to set a sigalarm of 0 please call signal.alarm directly as there are
     49     many situations where a 0 timeout is considered invalid.
     50 
     51     @param timeout_secs: The new timeout, in seconds.
     52     @param default_timeout: The default timeout to use, if timeout <= 0.
     53     @return: The old sigalarm timeout
     54     """
     55     timeout_sec_n = int(timeout_secs)
     56     if timeout_sec_n <= 0:
     57         timeout_sec_n = int(default_timeout)
     58     return signal.alarm(timeout_sec_n)
     59 
     60 
     61 def sigalarm_wrapper(message):
     62     """
     63     Raise a TimeoutException with the given message.  Needed because the body
     64     of a closure (lambda) can only be an expression, not a statement (such
     65     as "raise") :P :P :P
     66 
     67     @param message: the exception message.
     68     """
     69     raise error.TimeoutException(message)
     70 
     71 
     72 def custom_sigalarm_handler(func, timeout_sec):
     73     """
     74     Returns a sigalarm handler which produces an exception with a custom
     75     error message (function name and timeout length) instead of a generic
     76     one.
     77 
     78     @param func: the function that may time out
     79     @param timeout_sec: timeout length in seconds
     80     """
     81     try:
     82         name = str(func.__name__)
     83     except Exception as e:
     84         name = '(unavailable function name: exception: %s)' % e
     85     message = "sigalarm timeout (%d seconds) in %s" % (timeout_sec, name)
     86     return lambda signum, frame: sigalarm_wrapper(message)
     87 
     88 
     89 def timeout(func, args=(), kwargs={}, timeout_sec=60.0, default_result=None):
     90     """
     91     This function run the given function using the args, kwargs and
     92     return the given default value if the timeout_sec is exceeded.
     93 
     94     @param func: function to be called.
     95     @param args: arguments for function to be called.
     96     @param kwargs: keyword arguments for function to be called.
     97     @param timeout_sec: timeout setting for call to exit, in seconds.
     98     @param default_result: default return value for the function call.
     99 
    100     @return 1: is_timeout 2: result of the function call. If
    101             is_timeout is True, the call is timed out. If the
    102             value is False, the call is finished on time.
    103     """
    104     old_alarm_sec = 0
    105     old_handler = signal.getsignal(signal.SIGALRM)
    106     handler = custom_sigalarm_handler(func, timeout_sec)
    107     installed_handler = install_sigalarm_handler(handler)
    108     if installed_handler:
    109         old_alarm_sec = set_sigalarm_timeout(timeout_sec, default_timeout=60)
    110 
    111     # If old_timeout_time = 0 we either didn't install a handler, or sigalrm
    112     # had a signal.SIG_DFL handler with 0 timeout. In the latter case we still
    113     # need to restore the handler/timeout.
    114     old_timeout_time = (time.time() + old_alarm_sec) if old_alarm_sec > 0 else 0
    115 
    116     try:
    117         default_result = func(*args, **kwargs)
    118         return False, default_result
    119     except error.TimeoutException:
    120         return True, default_result
    121     finally:
    122         # If we installed a sigalarm handler, cancel it since our function
    123         # returned on time. If we can successfully restore the old handler,
    124         # reset the old timeout, or, if the old timeout's deadline has passed,
    125         # set the sigalarm to fire in one second. If the old_timeout_time is 0
    126         # we don't need to set the sigalarm timeout since we have already set it
    127         # as a byproduct of cancelling the current signal.
    128         if installed_handler:
    129             signal.alarm(0)
    130             if install_sigalarm_handler(old_handler) and old_timeout_time:
    131                 set_sigalarm_timeout(int(old_timeout_time - time.time()),
    132                                      default_timeout=1)
    133 
    134 
    135 
    136 def retry(ExceptionToCheck, timeout_min=1.0, delay_sec=3, blacklist=None,
    137           exception_to_raise=None, label=None):
    138     """Retry calling the decorated function using a delay with jitter.
    139 
    140     Will raise RPC ValidationError exceptions from the decorated
    141     function without retrying; a malformed RPC isn't going to
    142     magically become good. Will raise exceptions in blacklist as well.
    143 
    144     If the retry is done in a child thread, timeout may not be enforced as
    145     signal only works in main thread. Therefore, the retry inside a child
    146     thread may run longer than timeout or even hang.
    147 
    148     original from:
    149       http://www.saltycrane.com/blog/2009/11/trying-out-retry-decorator-python/
    150 
    151     @param ExceptionToCheck: the exception to check.  May be a tuple of
    152                              exceptions to check.
    153     @param timeout_min: timeout in minutes until giving up.
    154     @param delay_sec: pre-jittered delay between retries in seconds.  Actual
    155                       delays will be centered around this value, ranging up to
    156                       50% off this midpoint.
    157     @param blacklist: a list of exceptions that will be raised without retrying.
    158     @param exception_to_raise: the exception to raise. Callers can specify the
    159                                exception they want to raise.
    160     @param label: a label added to the exception message to help debug.
    161     """
    162     def deco_retry(func):
    163         """
    164         Decorator wrapper.
    165 
    166         @param func: the function to be retried and timed-out.
    167         """
    168         random.seed()
    169 
    170 
    171         def delay():
    172             """
    173             'Jitter' the delay, up to 50% in either direction.
    174             """
    175             random_delay = random.uniform(.5 * delay_sec, 1.5 * delay_sec)
    176             logging.warning('Retrying in %f seconds...', random_delay)
    177             time.sleep(random_delay)
    178 
    179 
    180         def func_retry(*args, **kwargs):
    181             """
    182             Used to cache exception to be raised later.
    183             """
    184             exc_info = None
    185             delayed_enabled = False
    186             exception_tuple = () if blacklist is None else tuple(blacklist)
    187             start_time = time.time()
    188             remaining_time = timeout_min * 60
    189             is_main_thread = isinstance(threading.current_thread(),
    190                                         threading._MainThread)
    191             if label:
    192                 details = 'label="%s"' % label
    193             elif hasattr(func, '__name__'):
    194                 details = 'function="%s()"' % func.__name__
    195             else:
    196                 details = 'unknown function'
    197 
    198             exception_message = ('retry exception (%s), timeout = %ds' %
    199                                  (details, timeout_min * 60))
    200 
    201             while remaining_time > 0:
    202                 if delayed_enabled:
    203                     delay()
    204                 else:
    205                     delayed_enabled = True
    206                 try:
    207                     # Clear the cache
    208                     exc_info = None
    209                     if is_main_thread:
    210                         is_timeout, result = timeout(func, args, kwargs,
    211                                                      remaining_time)
    212                         if not is_timeout:
    213                             return result
    214                     else:
    215                         return func(*args, **kwargs)
    216                 except exception_tuple:
    217                     raise
    218                 except error.CrosDynamicSuiteException:
    219                     raise
    220                 except ExceptionToCheck as e:
    221                     logging.warning('%s(%s)', e.__class__, e)
    222                     # Cache the exception to be raised later.
    223                     exc_info = sys.exc_info()
    224 
    225                 remaining_time = int(timeout_min * 60 -
    226                                      (time.time() - start_time))
    227 
    228             # The call must have timed out or raised ExceptionToCheck.
    229             if not exc_info:
    230                 if exception_to_raise:
    231                     raise exception_to_raise(exception_message)
    232                 else:
    233                     raise error.TimeoutException(exception_message)
    234             # Raise the cached exception with original backtrace.
    235             if exception_to_raise:
    236                 raise exception_to_raise('%s: %s' % (exc_info[0], exc_info[1]))
    237             raise exc_info[0], exc_info[1], exc_info[2]
    238 
    239 
    240         return func_retry  # true decorator
    241     return deco_retry
    242