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