• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020 Google LLC
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#      http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Retry decorator. Copied from ClusterFuzz source."""
15
16import functools
17import inspect
18import logging
19import sys
20import time
21
22# pylint: disable=too-many-arguments,broad-except
23
24
25def sleep(seconds):
26  """Invoke time.sleep. This is to avoid the flakiness of time.sleep. See:
27    crbug.com/770375"""
28  time.sleep(seconds)
29
30
31def get_delay(num_try, delay, backoff):
32  """Compute backoff delay."""
33  return delay * (backoff**(num_try - 1))
34
35
36def wrap(retries,
37         delay,
38         backoff=2,
39         exception_type=Exception,
40         retry_on_false=False):
41  """Retry decorator for a function."""
42
43  assert delay > 0
44  assert backoff >= 1
45  assert retries >= 0
46
47  def decorator(func):
48    """Decorator for the given function."""
49    tries = retries + 1
50    is_generator = inspect.isgeneratorfunction(func)
51    function_with_type = func.__qualname__
52    if is_generator:
53      function_with_type += ' (generator)'
54
55    def handle_retry(num_try, exception=None):
56      """Handle retry."""
57      if (exception is None or
58          isinstance(exception, exception_type)) and num_try < tries:
59        logging.log('Retrying on %s failed with %s. Retrying again.',
60                    function_with_type,
61                    sys.exc_info()[1])
62        sleep(get_delay(num_try, delay, backoff))
63        return True
64
65      logging.error('Retrying on %s failed with %s. Raise.', function_with_type,
66                    sys.exc_info()[1])
67      return False
68
69    @functools.wraps(func)
70    def _wrapper(*args, **kwargs):
71      """Regular function wrapper."""
72      for num_try in range(1, tries + 1):
73        try:
74          result = func(*args, **kwargs)
75          if retry_on_false and not result:
76            if not handle_retry(num_try):
77              return result
78
79            continue
80          return result
81        except Exception as error:
82          if not handle_retry(num_try, exception=error):
83            raise
84
85    @functools.wraps(func)
86    def _generator_wrapper(*args, **kwargs):
87      """Generator function wrapper."""
88      # This argument is not applicable for generator functions.
89      assert not retry_on_false
90      already_yielded_element_count = 0
91      for num_try in range(1, tries + 1):
92        try:
93          for index, result in enumerate(func(*args, **kwargs)):
94            if index >= already_yielded_element_count:
95              yield result
96              already_yielded_element_count += 1
97          break
98        except Exception as error:
99          if not handle_retry(num_try, exception=error):
100            raise
101
102    if is_generator:
103      return _generator_wrapper
104    return _wrapper
105
106  return decorator
107