Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright 2023 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.
15"""
16Generator wrapper for retryable streaming RPCs.
17"""
19from __future__ import annotations
21import functools
22import sys
23import time
24from typing import (
25 TYPE_CHECKING,
26 Any,
27 Callable,
28 Generator,
29 Iterable,
30 List,
31 Optional,
32 Tuple,
33 TypeVar,
34)
36from google.api_core.retry import (
37 RetryFailureReason,
38 build_retry_error,
39 exponential_sleep_generator,
40)
41from google.api_core.retry.retry_base import _BaseRetry, _retry_error_helper
43if TYPE_CHECKING:
44 if sys.version_info >= (3, 10):
45 from typing import ParamSpec
46 else:
47 from typing_extensions import ParamSpec
49 _P = ParamSpec("_P") # target function call parameters
50 _Y = TypeVar("_Y") # yielded values
53def retry_target_stream(
54 target: Callable[_P, Iterable[_Y]],
55 predicate: Callable[[Exception], bool],
56 sleep_generator: Iterable[float],
57 timeout: Optional[float] = None,
58 on_error: Optional[Callable[[Exception], None]] = None,
59 exception_factory: Callable[
60 [List[Exception], RetryFailureReason, Optional[float]],
61 Tuple[Exception, Optional[Exception]],
62 ] = build_retry_error,
63 init_args: tuple = (),
64 init_kwargs: dict = {},
65 **kwargs,
66) -> Generator[_Y, Any, None]:
67 """Create a generator wrapper that retries the wrapped stream if it fails.
69 This is the lowest-level retry helper. Generally, you'll use the
70 higher-level retry helper :class:`Retry`.
72 Args:
73 target: The generator function to call and retry.
74 predicate: A callable used to determine if an
75 exception raised by the target should be considered retryable.
76 It should return True to retry or False otherwise.
77 sleep_generator: An infinite iterator that determines
78 how long to sleep between retries.
79 timeout: How long to keep retrying the target.
80 Note: timeout is only checked before initiating a retry, so the target may
81 run past the timeout value as long as it is healthy.
82 on_error: If given, the on_error callback will be called with each
83 retryable exception raised by the target. Any error raised by this
84 function will *not* be caught.
85 exception_factory: A function that is called when the retryable reaches
86 a terminal failure state, used to construct an exception to be raised.
87 It takes a list of all exceptions encountered, a retry.RetryFailureReason
88 enum indicating the failure cause, and the original timeout value
89 as arguments. It should return a tuple of the exception to be raised,
90 along with the cause exception if any. The default implementation will raise
91 a RetryError on timeout, or the last exception encountered otherwise.
92 init_args: Positional arguments to pass to the target function.
93 init_kwargs: Keyword arguments to pass to the target function.
95 Returns:
96 Generator: A retryable generator that wraps the target generator function.
98 Raises:
99 ValueError: If the sleep generator stops yielding values.
100 Exception: a custom exception specified by the exception_factory if provided.
101 If no exception_factory is provided:
102 google.api_core.RetryError: If the timeout is exceeded while retrying.
103 Exception: If the target raises an error that isn't retryable.
104 """
106 timeout = kwargs.get("deadline", timeout)
107 deadline: Optional[float] = (
108 time.monotonic() + timeout if timeout is not None else None
109 )
110 error_list: list[Exception] = []
111 sleep_iter = iter(sleep_generator)
113 # continue trying until an attempt completes, or a terminal exception is raised in _retry_error_helper
114 # TODO: support max_attempts argument: https://github.com/googleapis/python-api-core/issues/535
115 while True:
116 # Start a new retry loop
117 try:
118 # Note: in the future, we can add a ResumptionStrategy object
119 # to generate new args between calls. For now, use the same args
120 # for each attempt.
121 subgenerator = target(*init_args, **init_kwargs)
122 return (yield from subgenerator)
123 # handle exceptions raised by the subgenerator
124 # pylint: disable=broad-except
125 # This function explicitly must deal with broad exceptions.
126 except Exception as exc:
127 # defer to shared logic for handling errors
128 next_sleep = _retry_error_helper(
129 exc,
130 deadline,
131 sleep_iter,
132 error_list,
133 predicate,
134 on_error,
135 exception_factory,
136 timeout,
137 )
138 # if exception not raised, sleep before next attempt
139 time.sleep(next_sleep)
142class StreamingRetry(_BaseRetry):
143 """Exponential retry decorator for streaming synchronous RPCs.
145 This class returns a Generator when called, which wraps the target
146 stream in retry logic. If any exception is raised by the target, the
147 entire stream will be retried within the wrapper.
149 Although the default behavior is to retry transient API errors, a
150 different predicate can be provided to retry other exceptions.
152 Important Note: when a stream encounters a retryable error, it will
153 silently construct a fresh iterator instance in the background
154 and continue yielding (likely duplicate) values as if no error occurred.
155 This is the most general way to retry a stream, but it often is not the
156 desired behavior. Example: iter([1, 2, 1/0]) -> [1, 2, 1, 2, ...]
158 There are two ways to build more advanced retry logic for streams:
160 1. Wrap the target
161 Use a ``target`` that maintains state between retries, and creates a
162 different generator on each retry call. For example, you can wrap a
163 network call in a function that modifies the request based on what has
164 already been returned:
166 .. code-block:: python
168 def attempt_with_modified_request(target, request, seen_items=[]):
169 # remove seen items from request on each attempt
170 new_request = modify_request(request, seen_items)
171 new_generator = target(new_request)
172 for item in new_generator:
173 yield item
174 seen_items.append(item)
176 retry_wrapped_fn = StreamingRetry()(attempt_with_modified_request)
177 retryable_generator = retry_wrapped_fn(target, request)
179 2. Wrap the retry generator
180 Alternatively, you can wrap the retryable generator itself before
181 passing it to the end-user to add a filter on the stream. For
182 example, you can keep track of the items that were successfully yielded
183 in previous retry attempts, and only yield new items when the
184 new attempt surpasses the previous ones:
186 .. code-block:: python
188 def retryable_with_filter(target):
189 stream_idx = 0
190 # reset stream_idx when the stream is retried
191 def on_error(e):
192 nonlocal stream_idx
193 stream_idx = 0
194 # build retryable
195 retryable_gen = StreamingRetry(...)(target)
196 # keep track of what has been yielded out of filter
197 seen_items = []
198 for item in retryable_gen():
199 if stream_idx >= len(seen_items):
200 seen_items.append(item)
201 yield item
202 elif item != seen_items[stream_idx]:
203 raise ValueError("Stream differs from last attempt")
204 stream_idx += 1
206 filter_retry_wrapped = retryable_with_filter(target)
208 Args:
209 predicate (Callable[Exception]): A callable that should return ``True``
210 if the given exception is retryable.
211 initial (float): The minimum amount of time to delay in seconds. This
212 must be greater than 0.
213 maximum (float): The maximum amount of time to delay in seconds.
214 multiplier (float): The multiplier applied to the delay.
215 timeout (float): How long to keep retrying, in seconds.
216 Note: timeout is only checked before initiating a retry, so the target may
217 run past the timeout value as long as it is healthy.
218 on_error (Callable[Exception]): A function to call while processing
219 a retryable exception. Any error raised by this function will
220 *not* be caught.
221 deadline (float): DEPRECATED: use `timeout` instead. For backward
222 compatibility, if specified it will override the ``timeout`` parameter.
223 """
225 def __call__(
226 self,
227 func: Callable[_P, Iterable[_Y]],
228 on_error: Callable[[Exception], Any] | None = None,
229 ) -> Callable[_P, Generator[_Y, Any, None]]:
230 """Wrap a callable with retry behavior.
232 Args:
233 func (Callable): The callable to add retry behavior to.
234 on_error (Optional[Callable[Exception]]): If given, the
235 on_error callback will be called with each retryable exception
236 raised by the wrapped function. Any error raised by this
237 function will *not* be caught. If on_error was specified in the
238 constructor, this value will be ignored.
240 Returns:
241 Callable: A callable that will invoke ``func`` with retry
242 behavior.
243 """
244 if self._on_error is not None:
245 on_error = self._on_error
247 @functools.wraps(func)
248 def retry_wrapped_func(
249 *args: _P.args, **kwargs: _P.kwargs
250 ) -> Generator[_Y, Any, None]:
251 """A wrapper that calls target function with retry."""
252 sleep_generator = exponential_sleep_generator(
253 self._initial, self._maximum, multiplier=self._multiplier
254 )
255 return retry_target_stream(
256 func,
257 predicate=self._predicate,
258 sleep_generator=sleep_generator,
259 timeout=self._timeout,
260 on_error=on_error,
261 init_args=args,
262 init_kwargs=kwargs,
263 )
265 return retry_wrapped_func