1# Copyright 2014 The Chromium 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""" 5Function/method decorators that provide timeout and retry logic. 6""" 7 8import functools 9import itertools 10import sys 11 12import six 13 14from devil.android import device_errors 15from devil.utils import cmd_helper 16from devil.utils import reraiser_thread 17from devil.utils import timeout_retry 18 19DEFAULT_TIMEOUT_ATTR = '_default_timeout' 20DEFAULT_RETRIES_ATTR = '_default_retries' 21 22 23def _TimeoutRetryWrapper(f, 24 timeout_func, 25 retries_func, 26 retry_if_func=timeout_retry.AlwaysRetry, 27 pass_values=False): 28 """ Wraps a funcion with timeout and retry handling logic. 29 30 Args: 31 f: The function to wrap. 32 timeout_func: A callable that returns the timeout value. 33 retries_func: A callable that returns the retries value. 34 pass_values: If True, passes the values returned by |timeout_func| and 35 |retries_func| to the wrapped function as 'timeout' and 36 'retries' kwargs, respectively. 37 Returns: 38 The wrapped function. 39 """ 40 41 @functools.wraps(f) 42 def timeout_retry_wrapper(*args, **kwargs): 43 timeout = timeout_func(*args, **kwargs) 44 retries = retries_func(*args, **kwargs) 45 if pass_values: 46 kwargs['timeout'] = timeout 47 kwargs['retries'] = retries 48 49 @functools.wraps(f) 50 def impl(): 51 return f(*args, **kwargs) 52 53 try: 54 if timeout_retry.CurrentTimeoutThreadGroup(): 55 # Don't wrap if there's already an outer timeout thread. 56 return impl() 57 else: 58 desc = '%s(%s)' % (f.__name__, ', '.join( 59 itertools.chain( 60 (str(a) for a in args), 61 ('%s=%s' % (k, str(v)) for k, v in six.iteritems(kwargs))))) 62 return timeout_retry.Run( 63 impl, timeout, retries, desc=desc, retry_if_func=retry_if_func) 64 except reraiser_thread.TimeoutError as e: 65 six.reraise( 66 device_errors.CommandTimeoutError, 67 device_errors.CommandTimeoutError(str(e)), 68 sys.exc_info()[2]) 69 except cmd_helper.TimeoutError as e: 70 six.reraise( 71 device_errors.CommandTimeoutError, 72 device_errors.CommandTimeoutError(str(e), output=e.output), 73 sys.exc_info()[2]) 74 75 return timeout_retry_wrapper 76 77 78def WithTimeoutAndRetries(f): 79 """A decorator that handles timeouts and retries. 80 81 'timeout' and 'retries' kwargs must be passed to the function. 82 83 Args: 84 f: The function to decorate. 85 Returns: 86 The decorated function. 87 """ 88 get_timeout = lambda *a, **kw: kw['timeout'] 89 get_retries = lambda *a, **kw: kw['retries'] 90 return _TimeoutRetryWrapper(f, get_timeout, get_retries) 91 92 93def WithTimeoutAndConditionalRetries(retry_if_func): 94 """Returns a decorator that handles timeouts and, in some cases, retries. 95 96 'timeout' and 'retries' kwargs must be passed to the function. 97 98 Args: 99 retry_if_func: A unary callable that takes an exception and returns 100 whether failures should be retried. 101 Returns: 102 The actual decorator. 103 """ 104 105 def decorator(f): 106 get_timeout = lambda *a, **kw: kw['timeout'] 107 get_retries = lambda *a, **kw: kw['retries'] 108 return _TimeoutRetryWrapper( 109 f, get_timeout, get_retries, retry_if_func=retry_if_func) 110 111 return decorator 112 113 114def WithExplicitTimeoutAndRetries(timeout, retries): 115 """Returns a decorator that handles timeouts and retries. 116 117 The provided |timeout| and |retries| values are always used. 118 119 Args: 120 timeout: The number of seconds to wait for the decorated function to 121 return. Always used. 122 retries: The number of times the decorated function should be retried on 123 failure. Always used. 124 Returns: 125 The actual decorator. 126 """ 127 128 def decorator(f): 129 get_timeout = lambda *a, **kw: timeout 130 get_retries = lambda *a, **kw: retries 131 return _TimeoutRetryWrapper(f, get_timeout, get_retries) 132 133 return decorator 134 135 136def WithTimeoutAndRetriesDefaults(default_timeout, default_retries): 137 """Returns a decorator that handles timeouts and retries. 138 139 The provided |default_timeout| and |default_retries| values are used only 140 if timeout and retries values are not provided. 141 142 Args: 143 default_timeout: The number of seconds to wait for the decorated function 144 to return. Only used if a 'timeout' kwarg is not passed 145 to the decorated function. 146 default_retries: The number of times the decorated function should be 147 retried on failure. Only used if a 'retries' kwarg is not 148 passed to the decorated function. 149 Returns: 150 The actual decorator. 151 """ 152 153 def decorator(f): 154 get_timeout = lambda *a, **kw: kw.get('timeout', default_timeout) 155 get_retries = lambda *a, **kw: kw.get('retries', default_retries) 156 return _TimeoutRetryWrapper(f, get_timeout, get_retries, pass_values=True) 157 158 return decorator 159 160 161def WithTimeoutAndRetriesFromInstance(default_timeout_name=DEFAULT_TIMEOUT_ATTR, 162 default_retries_name=DEFAULT_RETRIES_ATTR, 163 min_default_timeout=None): 164 """Returns a decorator that handles timeouts and retries. 165 166 The provided |default_timeout_name| and |default_retries_name| are used to 167 get the default timeout value and the default retries value from the object 168 instance if timeout and retries values are not provided. 169 170 Note that this should only be used to decorate methods, not functions. 171 172 Args: 173 default_timeout_name: The name of the default timeout attribute of the 174 instance. 175 default_retries_name: The name of the default retries attribute of the 176 instance. 177 min_timeout: Miniumum timeout to be used when using instance timeout. 178 Returns: 179 The actual decorator. 180 """ 181 182 def decorator(f): 183 def get_timeout(inst, *_args, **kwargs): 184 ret = getattr(inst, default_timeout_name) 185 if min_default_timeout is not None: 186 ret = max(min_default_timeout, ret) 187 return kwargs.get('timeout', ret) 188 189 def get_retries(inst, *_args, **kwargs): 190 return kwargs.get('retries', getattr(inst, default_retries_name)) 191 192 return _TimeoutRetryWrapper(f, get_timeout, get_retries, pass_values=True) 193 194 return decorator 195