summaryrefslogtreecommitdiff
path: root/systrace/catapult/devil/devil/android/decorators.py
blob: 11d2494b5a6016f2990734e276b0171bcd208aa8 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
# Copyright 2014 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""
Function/method decorators that provide timeout and retry logic.
"""

import functools
import itertools
import sys

import six

from devil.android import device_errors
from devil.utils import cmd_helper
from devil.utils import reraiser_thread
from devil.utils import timeout_retry

DEFAULT_TIMEOUT_ATTR = '_default_timeout'
DEFAULT_RETRIES_ATTR = '_default_retries'


def _TimeoutRetryWrapper(f,
                         timeout_func,
                         retries_func,
                         retry_if_func=timeout_retry.AlwaysRetry,
                         pass_values=False):
  """ Wraps a funcion with timeout and retry handling logic.

  Args:
    f: The function to wrap.
    timeout_func: A callable that returns the timeout value.
    retries_func: A callable that returns the retries value.
    pass_values: If True, passes the values returned by |timeout_func| and
                 |retries_func| to the wrapped function as 'timeout' and
                 'retries' kwargs, respectively.
  Returns:
    The wrapped function.
  """

  @functools.wraps(f)
  def timeout_retry_wrapper(*args, **kwargs):
    timeout = timeout_func(*args, **kwargs)
    retries = retries_func(*args, **kwargs)
    if pass_values:
      kwargs['timeout'] = timeout
      kwargs['retries'] = retries

    @functools.wraps(f)
    def impl():
      return f(*args, **kwargs)

    try:
      if timeout_retry.CurrentTimeoutThreadGroup():
        # Don't wrap if there's already an outer timeout thread.
        return impl()
      else:
        desc = '%s(%s)' % (f.__name__, ', '.join(
            itertools.chain(
                (str(a) for a in args),
                ('%s=%s' % (k, str(v)) for k, v in six.iteritems(kwargs)))))
        return timeout_retry.Run(
            impl, timeout, retries, desc=desc, retry_if_func=retry_if_func)
    except reraiser_thread.TimeoutError as e:
      six.reraise(
          device_errors.CommandTimeoutError,
          device_errors.CommandTimeoutError(str(e)),
          sys.exc_info()[2])
    except cmd_helper.TimeoutError as e:
      six.reraise(
          device_errors.CommandTimeoutError,
          device_errors.CommandTimeoutError(str(e), output=e.output),
          sys.exc_info()[2])

  return timeout_retry_wrapper


def WithTimeoutAndRetries(f):
  """A decorator that handles timeouts and retries.

  'timeout' and 'retries' kwargs must be passed to the function.

  Args:
    f: The function to decorate.
  Returns:
    The decorated function.
  """
  get_timeout = lambda *a, **kw: kw['timeout']
  get_retries = lambda *a, **kw: kw['retries']
  return _TimeoutRetryWrapper(f, get_timeout, get_retries)


def WithTimeoutAndConditionalRetries(retry_if_func):
  """Returns a decorator that handles timeouts and, in some cases, retries.

  'timeout' and 'retries' kwargs must be passed to the function.

  Args:
    retry_if_func: A unary callable that takes an exception and returns
      whether failures should be retried.
  Returns:
    The actual decorator.
  """

  def decorator(f):
    get_timeout = lambda *a, **kw: kw['timeout']
    get_retries = lambda *a, **kw: kw['retries']
    return _TimeoutRetryWrapper(
        f, get_timeout, get_retries, retry_if_func=retry_if_func)

  return decorator


def WithExplicitTimeoutAndRetries(timeout, retries):
  """Returns a decorator that handles timeouts and retries.

  The provided |timeout| and |retries| values are always used.

  Args:
    timeout: The number of seconds to wait for the decorated function to
             return. Always used.
    retries: The number of times the decorated function should be retried on
             failure. Always used.
  Returns:
    The actual decorator.
  """

  def decorator(f):
    get_timeout = lambda *a, **kw: timeout
    get_retries = lambda *a, **kw: retries
    return _TimeoutRetryWrapper(f, get_timeout, get_retries)

  return decorator


def WithTimeoutAndRetriesDefaults(default_timeout, default_retries):
  """Returns a decorator that handles timeouts and retries.

  The provided |default_timeout| and |default_retries| values are used only
  if timeout and retries values are not provided.

  Args:
    default_timeout: The number of seconds to wait for the decorated function
                     to return. Only used if a 'timeout' kwarg is not passed
                     to the decorated function.
    default_retries: The number of times the decorated function should be
                     retried on failure. Only used if a 'retries' kwarg is not
                     passed to the decorated function.
  Returns:
    The actual decorator.
  """

  def decorator(f):
    get_timeout = lambda *a, **kw: kw.get('timeout', default_timeout)
    get_retries = lambda *a, **kw: kw.get('retries', default_retries)
    return _TimeoutRetryWrapper(f, get_timeout, get_retries, pass_values=True)

  return decorator


def WithTimeoutAndRetriesFromInstance(default_timeout_name=DEFAULT_TIMEOUT_ATTR,
                                      default_retries_name=DEFAULT_RETRIES_ATTR,
                                      min_default_timeout=None):
  """Returns a decorator that handles timeouts and retries.

  The provided |default_timeout_name| and |default_retries_name| are used to
  get the default timeout value and the default retries value from the object
  instance if timeout and retries values are not provided.

  Note that this should only be used to decorate methods, not functions.

  Args:
    default_timeout_name: The name of the default timeout attribute of the
                          instance.
    default_retries_name: The name of the default retries attribute of the
                          instance.
    min_timeout: Miniumum timeout to be used when using instance timeout.
  Returns:
    The actual decorator.
  """

  def decorator(f):
    def get_timeout(inst, *_args, **kwargs):
      ret = getattr(inst, default_timeout_name)
      if min_default_timeout is not None:
        ret = max(min_default_timeout, ret)
      return kwargs.get('timeout', ret)

    def get_retries(inst, *_args, **kwargs):
      return kwargs.get('retries', getattr(inst, default_retries_name))

    return _TimeoutRetryWrapper(f, get_timeout, get_retries, pass_values=True)

  return decorator