1# Copyright 2017 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
15"""Decorators for applying timeout arguments to functions.
16
17These decorators are used to wrap API methods to apply either a constant
18or exponential timeout argument.
19
20For example, imagine an API method that can take a while to return results,
21such as one that might block until a resource is ready:
22
23.. code-block:: python
24
25    def is_thing_ready(timeout=None):
26        response = requests.get('https://example.com/is_thing_ready')
27        response.raise_for_status()
28        return response.json()
29
30This module allows a function like this to be wrapped so that timeouts are
31automatically determined, for example:
32
33.. code-block:: python
34
35    timeout_ = timeout.ExponentialTimeout()
36    is_thing_ready_with_timeout = timeout_(is_thing_ready)
37
38    for n in range(10):
39        try:
40            is_thing_ready_with_timeout({'example': 'data'})
41        except:
42            pass
43
44In this example the first call to ``is_thing_ready`` will have a relatively
45small timeout (like 1 second). If the resource is available and the request
46completes quickly, the loop exits. But, if the resource isn't yet available
47and the request times out, it'll be retried - this time with a larger timeout.
48
49In the broader context these decorators are typically combined with
50:mod:`google.api_core.retry` to implement API methods with a signature that
51matches ``api_method(request, timeout=None, retry=None)``.
52"""
53
54from __future__ import unicode_literals
55
56import datetime
57import functools
58
59from google.api_core import datetime_helpers
60
61_DEFAULT_INITIAL_TIMEOUT = 5.0  # seconds
62_DEFAULT_MAXIMUM_TIMEOUT = 30.0  # seconds
63_DEFAULT_TIMEOUT_MULTIPLIER = 2.0
64# If specified, must be in seconds. If none, deadline is not used in the
65# timeout calculation.
66_DEFAULT_DEADLINE = None
67
68
69class ConstantTimeout(object):
70    """A decorator that adds a constant timeout argument.
71
72    This is effectively equivalent to
73    ``functools.partial(func, timeout=timeout)``.
74
75    Args:
76        timeout (Optional[float]): the timeout (in seconds) to applied to the
77            wrapped function. If `None`, the target function is expected to
78            never timeout.
79    """
80
81    def __init__(self, timeout=None):
82        self._timeout = timeout
83
84    def __call__(self, func):
85        """Apply the timeout decorator.
86
87        Args:
88            func (Callable): The function to apply the timeout argument to.
89                This function must accept a timeout keyword argument.
90
91        Returns:
92            Callable: The wrapped function.
93        """
94
95        @functools.wraps(func)
96        def func_with_timeout(*args, **kwargs):
97            """Wrapped function that adds timeout."""
98            kwargs["timeout"] = self._timeout
99            return func(*args, **kwargs)
100
101        return func_with_timeout
102
103    def __str__(self):
104        return "<ConstantTimeout timeout={:.1f}>".format(self._timeout)
105
106
107def _exponential_timeout_generator(initial, maximum, multiplier, deadline):
108    """A generator that yields exponential timeout values.
109
110    Args:
111        initial (float): The initial timeout.
112        maximum (float): The maximum timeout.
113        multiplier (float): The multiplier applied to the timeout.
114        deadline (float): The overall deadline across all invocations.
115
116    Yields:
117        float: A timeout value.
118    """
119    if deadline is not None:
120        deadline_datetime = datetime_helpers.utcnow() + datetime.timedelta(
121            seconds=deadline
122        )
123    else:
124        deadline_datetime = datetime.datetime.max
125
126    timeout = initial
127    while True:
128        now = datetime_helpers.utcnow()
129        yield min(
130            # The calculated timeout based on invocations.
131            timeout,
132            # The set maximum timeout.
133            maximum,
134            # The remaining time before the deadline is reached.
135            float((deadline_datetime - now).seconds),
136        )
137        timeout = timeout * multiplier
138
139
140class ExponentialTimeout(object):
141    """A decorator that adds an exponentially increasing timeout argument.
142
143    This is useful if a function is called multiple times. Each time the
144    function is called this decorator will calculate a new timeout parameter
145    based on the the number of times the function has been called.
146
147    For example
148
149    .. code-block:: python
150
151    Args:
152        initial (float): The initial timeout to pass.
153        maximum (float): The maximum timeout for any one call.
154        multiplier (float): The multiplier applied to the timeout for each
155            invocation.
156        deadline (Optional[float]): The overall deadline across all
157            invocations. This is used to prevent a very large calculated
158            timeout from pushing the overall execution time over the deadline.
159            This is especially useful in conjuction with
160            :mod:`google.api_core.retry`. If ``None``, the timeouts will not
161            be adjusted to accomodate an overall deadline.
162    """
163
164    def __init__(
165        self,
166        initial=_DEFAULT_INITIAL_TIMEOUT,
167        maximum=_DEFAULT_MAXIMUM_TIMEOUT,
168        multiplier=_DEFAULT_TIMEOUT_MULTIPLIER,
169        deadline=_DEFAULT_DEADLINE,
170    ):
171        self._initial = initial
172        self._maximum = maximum
173        self._multiplier = multiplier
174        self._deadline = deadline
175
176    def with_deadline(self, deadline):
177        """Return a copy of this teimout with the given deadline.
178
179        Args:
180            deadline (float): The overall deadline across all invocations.
181
182        Returns:
183            ExponentialTimeout: A new instance with the given deadline.
184        """
185        return ExponentialTimeout(
186            initial=self._initial,
187            maximum=self._maximum,
188            multiplier=self._multiplier,
189            deadline=deadline,
190        )
191
192    def __call__(self, func):
193        """Apply the timeout decorator.
194
195        Args:
196            func (Callable): The function to apply the timeout argument to.
197                This function must accept a timeout keyword argument.
198
199        Returns:
200            Callable: The wrapped function.
201        """
202        timeouts = _exponential_timeout_generator(
203            self._initial, self._maximum, self._multiplier, self._deadline
204        )
205
206        @functools.wraps(func)
207        def func_with_timeout(*args, **kwargs):
208            """Wrapped function that adds timeout."""
209            kwargs["timeout"] = next(timeouts)
210            return func(*args, **kwargs)
211
212        return func_with_timeout
213
214    def __str__(self):
215        return (
216            "<ExponentialTimeout initial={:.1f}, maximum={:.1f}, "
217            "multiplier={:.1f}, deadline={:.1f}>".format(
218                self._initial, self._maximum, self._multiplier, self._deadline
219            )
220        )
221