Coverage for bzfs_main/util/retry.py: 100%

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

7# http://www.apache.org/licenses/LICENSE-2.0

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"""Customizable generic retry framework; defaults to jittered exponential backoff with cap unless specified otherwise.

17Purpose:

18--------

19- Provide a reusable retry helper for transient failures using customizable policy and callbacks.

20- Centralize backoff, jitter, logging and metrics behavior while keeping call sites compact.

21- Prevent accidental retries: the loop retries only when the developer explicitly raises a ``RetryableError``, which reduces

22 the risk of retrying non-idempotent operations.

23- Provide a thread-safe, fast implementation; avoid shared RNG contention.

24- Provide both sync and async API, both with the same semantics, except the async API awaits ``fn``, awaitable

25 ``before_attempt`` / ``after_attempt`` / ``on_retryable_error`` / ``on_exhaustion`` results, and non-blocking sleep.

26- Avoid unnecessary complexity and add zero dependencies beyond the Python standard library. Everything you need is in this

27 single Python file.

29Usage:

30------

31- Wrap work in a callable ``fn(retry: Retry)`` and therein raise ``RetryableError`` for failures that should be retried.

32- Construct a policy via ``RetryPolicy(...)`` that specifies how ``RetryableError`` shall be retried.

33- Invoke ``call_with_retries(fn=fn, policy=policy, log=logger)`` or call_with_retries_async with a standard logging.Logger.

34- On success, the result of calling ``fn`` is returned.

35- By default on exhaustion, call_with_retries* either re-raises the last underlying ``RetryableError.__cause__``, or raises

36 ``RetryError`` (wrapping the last ``RetryableError``), like so:

37 - if ``RetryPolicy.reraise`` is True and the last ``RetryableError.__cause__`` is not None, re-raise the last

38 ``RetryableError.__cause__`` with its original traceback.

39 - Otherwise, raise ``RetryError`` (wrapping the last ``RetryableError``, preserving its ``__cause__`` chain).

40 - The default is ``RetryPolicy.reraise=True``.

42Advanced Configuration:

43-----------------------

44- Tune ``RetryPolicy`` parameters to control maximum retries, sleep bounds, elapsed-time budget, logging, etc.

45- Use ``RetryPolicy.config: RetryConfig`` to control logging settings.

46- Set ``log=None`` to disable logging, or customize ``info_loglevel`` / ``warning_loglevel`` for structured logs.

47- Supply a ``giveup(AttemptOutcome)`` callback to stop retrying based on domain-specific logic (for example, decisions based

48 on time budget/quota or the previous N most recent AttemptOutcome objects (via AttemptOutcome.retry.previous_outcomes).

49- Use the ``any_giveup()`` / ``all_giveup()`` helper to consult more than one callback handler in ``giveup(AttemptOutcome)``.

50- Supply an ``on_exhaustion(AttemptOutcome)`` callback to customize behavior when giving up; it may raise an error or return

51 a fallback value.

53Observability:

54--------------

55- Supply an ``after_attempt(AttemptOutcome)`` callback to collect per-attempt metrics such as success flag,

56 exhausted/terminated state, attempt number, total elapsed duration (in nanoseconds), sleep duration (in nanoseconds), etc.

57- ``AttemptOutcome.result`` is either the successful result or the most recent ``RetryableError``, enabling integration with

58 metrics and tracing systems without coupling the retry loop to any specific backend.

59- Supply an ``after_attempt(AttemptOutcome)`` callback to customize logging 100%, if necessary.

60- Use the ``multi_after_attempt()`` helper to invoke more than one callback handler in ``after_attempt(AttemptOutcome)``.

62Expert Configuration:

63---------------------

64- Supply a ``backoff(BackoffContext)`` callback to plug in a custom backoff algorithm (e.g., decorrelated-jitter or

65 retry-after HTTP 429). The default is full-jitter exponential backoff with cap (aka industry standard).

66- Supply a ``before_attempt(Retry)`` callback for optional rate limiting or other forms of internal backpressure.

67- Supply a ``on_retryable_error(AttemptOutcome)`` callback, e.g. to count failures (RetryableError) caught by the retry loop.

68- Set ``RetryPolicy.max_previous_outcomes > 0`` to pass the N most recent AttemptOutcome objects to callbacks (default is 0).

69- If ``RetryPolicy.max_previous_outcomes > 0``, you can use ``RetryableError(..., attachment=...)`` to carry domain-specific

70 state from a failed attempt to the next attempt via ``retry.previous_outcomes``. This pattern helps if attempt N+1 is a

71 function of attempt N or all prior attempts (e.g., switching endpoints or resuming from an offset).

72- Use ``[Async]RetryTemplate`` as a bag-of-knobs configuration template for functions that shall be retried in similar ways.

73- Or package up all knobs plus a ``fn(retry: Retry)`` function into a self-contained auto-retrying higher level function by

74 constructing a ``[Async]RetryTemplate`` object (which is a ``Callable`` function itself).

75- To keep calling code retry-transparent, set ``RetryPolicy.reraise=True`` (the default) *and* raise retryable failures as

76 ``raise RetryableError(...) from exc``. Client code now won't notice whether call_with_retries* is used or not.

77- To make exhaustion observable to calling code, set ``RetryPolicy.reraise=False``: by default call_with_retries* now always

78 raises ``RetryError`` (wrapping the last ``RetryableError``) on exhaustion, so callers now catch ``RetryError`` and can

79 inspect the last underlying exception via ``err.outcome``, ``err.__cause__``, and even ``err.__cause__.__cause__`` when

80 present.

81- Set ``RetryPolicy.timing`` to customize reading the current monotonic time, sleeping and optional async termination.

82- The callback API is powerful enough to easily plug in advanced retry algorithms such as:

83 - Google SRE Client-Side Adaptive Throttling - https://sre.google/sre-book/handling-overload/

84 - gRPC retry throttling - https://grpc.io/docs/guides/retry/

85 - AWS SDK adaptive retry mode - https://docs.aws.amazon.com/sdkref/latest/guide/feature-retry-behavior.html

86 - Circuit breakers - https://martinfowler.com/bliki/CircuitBreaker.html (e.g. via `pybreaker` third-party library)

87 - Rate limiting with Fixed Window, Moving Window, and Sliding Windows (e.g. via `limits` third-party library)

88 - Various example backoff strategies such as decorrelated-jitter or retry-after HTTP 429 "Too Many Requests" responses,

89 etc can be found in test_retry_examples.py.

91Example Usage:

92--------------

93 import logging

94 from bzfs_main.util.retry import Retry, RetryPolicy, RetryableError, call_with_retries

96 def unreliable_operation(retry: Retry) -> str:

97 try:

98 if retry.count < 3:

99 raise ValueError("temporary failure connecting to foo.example.com")

100 return "ok"

101 except ValueError as exc:

102 # Preserve the underlying cause for correct error propagation and logging

103 raise RetryableError(display_msg="connect") from exc

104

105 retry_policy = RetryPolicy(

106 max_retries=10,

107 min_sleep_secs=0,

108 initial_max_sleep_secs=0.125,

109 max_sleep_secs=10,

110 max_elapsed_secs=60,

111 )

112 log = logging.getLogger(__name__)

113 result: str = call_with_retries(fn=unreliable_operation, policy=retry_policy, log=log)

114 print(result)

115

116 # Sample log output:

117 # INFO:Retrying connect [1/10] in 8.79ms ...

118 # INFO:Retrying connect [2/10] in 90.1ms ...

119 # INFO:Retrying connect [3/10] in 372ms ...

120 # ok

121

122Background:

123-----------

124For background on exponential backoff and jitter, see for example

125https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter

126"""

127

128from __future__ import (

129 annotations,

130)

131import argparse

132import dataclasses

133import functools

134import inspect

135import logging

136import random

137import threading

138import time

139from collections.abc import (

140 Awaitable,

141 Iterable,

142 Mapping,

143 Sequence,

144)

145from dataclasses import (

146 dataclass,

147)

148from typing import (

149 TYPE_CHECKING,

150 Any,

151 Callable,

152 Final,

153 Generic,

154 NamedTuple,

155 NoReturn,

156 TypeVar,

157 Union,

158 cast,

159 final,

160)

161

162if TYPE_CHECKING:

163 import asyncio

164

165from bzfs_main.util.utils import (

166 human_readable_duration,

167)

168

169# constants:

170INFINITY_MAX_RETRIES: Final[int] = 2**90 - 1 # a number that's essentially infinity for all practical retry purposes

171

172

173#############################################################################

174def full_jitter_backoff_strategy(context: BackoffContext) -> tuple[int, int]:

175 """Default implementation of ``backoff`` callback for call_with_retries(); computes delay time before next retry attempt,

176 after failure.

177

178 Full-jitter picks a random sleep_nanos duration from the range [min_sleep_nanos, curr_max_sleep_nanos] and applies

179 exponential backoff with cap to the next attempt; thread-safe. Typically, min_sleep_nanos is 0 and exponential_base is 2.

180 Example curr_max_sleep_nanos sequence: 125ms --> 250ms --> 500ms --> 1s --> 2s --> 4s --> 8s --> 10s --> 10s...

181 Full-jitter provides optimal balance between reducing server load and minimizing retry latency.

182 """

183 policy: RetryPolicy = context.retry.policy

184 curr_max_sleep_nanos: int = context.curr_max_sleep_nanos

185 if policy.min_sleep_nanos == curr_max_sleep_nanos:

186 sleep_nanos = curr_max_sleep_nanos # perf

187 else:

188 sleep_nanos = context.rng.randint(policy.min_sleep_nanos, curr_max_sleep_nanos) # nanos to delay until next attempt

189 curr_max_sleep_nanos = round(curr_max_sleep_nanos * policy.exponential_base) # exponential backoff

190 curr_max_sleep_nanos = min(curr_max_sleep_nanos, policy.max_sleep_nanos) # ... with cap for next attempt

191 return sleep_nanos, curr_max_sleep_nanos

192

193

194def no_giveup(outcome: AttemptOutcome) -> object | None:

195 """Default implementation of ``giveup`` callback for call_with_retries(); never gives up; returning anything other than

196 ``None`` indicates to give up retrying; thread-safe."""

197 return None # don't give up retrying

198

199

200def giveup_if_backoff_exceeds_deadline(outcome: AttemptOutcome) -> object | None:

201 """Implementation of ``giveup`` callback for call_with_retries() that gives up if the computed backoff sleep exceeds the

202 ``RetryPolicy.max_elapsed_secs`` deadline; thread-safe."""

203 return (

204 None

205 if outcome.elapsed_nanos + outcome.sleep_nanos < outcome.retry.policy.max_elapsed_nanos

206 else "backoff exceeds deadline"

207 )

208

209

210def before_attempt_noop(retry: Retry) -> int:

211 """Default implementation of ``before_attempt`` callback for call_with_retries(); does nothing; thread-safe."""

212 return 0

213

214

215def after_attempt_log_failure(outcome: AttemptOutcome) -> None:

216 """Default implementation of ``after_attempt`` callback for call_with_retries(); performs simple logging of retry attempt

217 failures; thread-safe."""

218 retry: Retry = outcome.retry

219 policy: RetryPolicy = retry.policy

220 config: RetryConfig = policy.config

221 if outcome.is_success or retry.log is None or not config.enable_logging:

222 return

223 log: logging.Logger = retry.log

224 assert isinstance(outcome.result, RetryableError)

225 retryable_error: RetryableError = outcome.result

226 if not outcome.is_exhausted:

227 if log.isEnabledFor(config.info_loglevel): # Retrying X in Y ms ...

228 m1: str = config.format_msg(config.display_msg, retryable_error)

229 m2: str = config.format_pair(retry.count + 1, policy.max_retries)

230 m3: str = config.format_duration(outcome.sleep_nanos)

231 log.log(config.info_loglevel, "%s", f"{m1}{m2} in {m3}{config.dots}", extra=config.extra)

232 else:

233 if policy.max_retries > 0 and log.isEnabledFor(config.warning_loglevel) and not outcome.is_terminated:

234 reason: str = "" if outcome.giveup_reason is None else f"{outcome.giveup_reason}; "

235 format_duration: Callable[[int], str] = config.format_duration # lambda: nanos

236 log.log(

237 config.warning_loglevel,

238 "%s",

239 f"{config.format_msg(config.display_msg, retryable_error)}"

240 f"exhausted; giving up because {reason}the last "

241 f"{config.format_pair(retry.count, policy.max_retries)} retries across "

242 f"{config.format_pair(format_duration(outcome.elapsed_nanos), format_duration(policy.max_elapsed_nanos))} "

243 "failed",

244 exc_info=retryable_error if config.exc_info else None,

245 stack_info=config.stack_info,

246 extra=config.extra,

247 )

248

249

250def noop(outcome: AttemptOutcome) -> None:

251 """Default implementation of ``on_retryable_error`` callback for call_with_retries(); does nothing; thread-safe."""

252

253

254def on_exhaustion_raise(outcome: AttemptOutcome) -> NoReturn:

255 """Default implementation of ``on_exhaustion`` callback for call_with_retries(); always raises; thread-safe."""

256 assert outcome.is_exhausted

257 assert isinstance(outcome.result, RetryableError)

258 retryable_error: RetryableError = outcome.result

259 policy: RetryPolicy = outcome.retry.policy

260 cause: BaseException | None = retryable_error.__cause__

261 if policy.reraise and cause is not None:

262 raise cause.with_traceback(cause.__traceback__)

263 else:

264 raise RetryError(outcome=outcome) from retryable_error

265

266

267def _update_previous_outcomes(

268 previous_outcomes: tuple[AttemptOutcome, ...], outcome: AttemptOutcome, policy: RetryPolicy, retry: Retry

269) -> tuple[AttemptOutcome, ...]:

270 """Computes value of previous_outcomes for next retry iteration."""

271 n: int = policy.max_previous_outcomes

272 if n > 0: # outcome will be passed to next attempt via Retry.previous_outcomes

273 if previous_outcomes: # detach to reduce memory footprint

274 outcome = outcome.copy(retry=retry.copy(previous_outcomes=()))

275 previous_outcomes = previous_outcomes[len(previous_outcomes) - n + 1 :] + (outcome,) # immutable deque

276 return previous_outcomes

277

278

279#############################################################################

280_T = TypeVar("_T")

281

282

283def call_with_retries(

284 fn: Callable[[Retry], _T], # typically a lambda; wraps work and raises RetryableError for failures that shall be retried

285 policy: RetryPolicy, # specifies how ``RetryableError`` shall be retried

286 *,

287 backoff: BackoffStrategy = full_jitter_backoff_strategy, # computes delay time before next retry attempt, after failure

288 giveup: Callable[[AttemptOutcome], object | None] = no_giveup, # stop retrying based on domain-specific logic, e.g. time

289 before_attempt: Callable[[Retry], int] = before_attempt_noop, # e.g. wait due to rate limiting or internal backpressure

290 after_attempt: Callable[[AttemptOutcome], None] = after_attempt_log_failure, # e.g. record metrics and/or custom logging

291 on_retryable_error: Callable[[AttemptOutcome], None] = noop, # e.g. count failures (RetryableError) caught by retry loop

292 on_exhaustion: Callable[[AttemptOutcome], _T] = on_exhaustion_raise, # raise error or return fallback value

293 log: logging.Logger | None = None, # set this to ``None`` to disable logging

294) -> _T:

295 """Runs the function ``fn`` and returns its result; retries on failure as indicated by policy; thread-safe.

296

297 By default on exhaustion, call_with_retries() either re-raises the last underlying ``RetryableError.__cause__``, or raises

298 ``RetryError`` (wrapping the last ``RetryableError``), like so:

299 - if ``RetryPolicy.reraise`` is True and the last ``RetryableError.__cause__`` is not None, re-raise the last

300 ``RetryableError.__cause__`` with its original traceback.

301 - Otherwise, raise ``RetryError`` (wrapping the last ``RetryableError``, preserving its ``__cause__`` chain).

302 - The default is ``RetryPolicy.reraise=True``.

303

304 On the exhaustion path, ``on_exhaustion`` will be called exactly once (after the final after_attempt). The default

305 implementation raises as described above; custom ``on_exhaustion`` impls may return a fallback value instead of an error.

306 """

307 rng: random.Random | None = None

308 retry_count: int = 0

309 idle_nanos: int = 0

310 curr_max_sleep_nanos: int = policy.initial_max_sleep_nanos

311 previous_outcomes: tuple[AttemptOutcome, ...] = () # for safety pass *immutable* deque to callbacks

312 timing: RetryTiming = policy.timing

313 sleep: Callable[[int, Retry], None] = timing.sleep

314 monotonic_ns: Callable[[], int] = timing.monotonic_ns

315 call_start_nanos: Final[int] = monotonic_ns()

316 while True:

317 before_attempt_nanos: int = monotonic_ns() if retry_count != 0 else call_start_nanos

318 prev: tuple[AttemptOutcome, ...] = previous_outcomes

319 retry: Retry = Retry(

320 retry_count, call_start_nanos, before_attempt_nanos, before_attempt_nanos, idle_nanos, policy, log, prev

321 )

322 try:

323 if before_attempt is not before_attempt_noop:

324 before_attempt_sleep_nanos: int = before_attempt(retry)

325 assert before_attempt_sleep_nanos >= 0, before_attempt_sleep_nanos

326 if before_attempt_sleep_nanos > 0:

327 sleep(before_attempt_sleep_nanos, retry) # e.g. wait due to rate limiting or internal backpressure

328 attempt_start_nanos: int = monotonic_ns()

329 idle_nanos += attempt_start_nanos - before_attempt_nanos

330 retry = Retry(

331 retry_count, call_start_nanos, before_attempt_nanos, attempt_start_nanos, idle_nanos, policy, log, prev

332 )

333 timing.on_before_attempt(retry)

334 result: _T = fn(retry) # Call the target function and supply retry attempt number and other metadata

335 except RetryableError as retryable_error:

336 elapsed_nanos: int = monotonic_ns() - call_start_nanos

337 is_terminated: Callable[[Retry], bool] = timing.is_terminated

338 giveup_reason: object | None = None

339 sleep_nanos: int = 0

340 outcome = AttemptOutcome(retry, False, False, False, None, elapsed_nanos, sleep_nanos, retryable_error)

341 on_retryable_error(outcome) # e.g. count failures (RetryableError) caught by retry loop

342 if retry_count < policy.max_retries and elapsed_nanos < policy.max_elapsed_nanos:

343 if policy.max_sleep_nanos == 0 and backoff is full_jitter_backoff_strategy:

344 pass # perf: e.g. spin-before-block

345 elif retry_count == 0 and retryable_error.retry_immediately_once:

346 pass # retry once immediately without backoff

347 else: # jitter: default backoff strategy picks random sleep_nanos in [min_sleep_nanos, curr_max_sleep_nanos]

348 rng = _thread_local_rng() if rng is None else rng

349 sleep_nanos, curr_max_sleep_nanos = backoff( # compute delay before next retry attempt, after failure

350 BackoffContext(retry, curr_max_sleep_nanos, rng, elapsed_nanos, retryable_error)

351 )

352 assert sleep_nanos >= 0 and curr_max_sleep_nanos >= 0, sleep_nanos

353

354 if sleep_nanos > 0:

355 outcome = AttemptOutcome(retry, False, False, False, None, elapsed_nanos, sleep_nanos, retryable_error)

356 if (not is_terminated(retry)) and (giveup_reason := giveup(outcome)) is None:

357 after_attempt(outcome)

358 sleep(sleep_nanos, retry)

359 idle_nanos += sleep_nanos

360 if not is_terminated(retry):

361 previous_outcomes = _update_previous_outcomes(previous_outcomes, outcome, policy, retry)

362 del outcome # help gc

363 retry_count += 1

364 continue # continue retry loop with next attempt

365 else:

366 sleep_nanos = 0

367 outcome = AttemptOutcome(

368 retry, False, True, is_terminated(retry), giveup_reason, elapsed_nanos, sleep_nanos, retryable_error

369 )

370 after_attempt(outcome)

371 return on_exhaustion(outcome) # raise error or return fallback value

372 else: # success

373 if after_attempt is not after_attempt_log_failure:

374 elapsed_nanos = monotonic_ns() - call_start_nanos

375 outcome = AttemptOutcome(retry, True, False, False, None, elapsed_nanos, 0, result)

376 after_attempt(outcome)

377 return result

378

379

380async def call_with_retries_async(

381 fn: Callable[[Retry], Awaitable[_T]], # wraps work and raises RetryableError for failures that shall be retried

382 policy: RetryPolicy, # specifies how ``RetryableError`` shall be retried

383 *,

384 backoff: BackoffStrategy = full_jitter_backoff_strategy, # computes delay time before next retry attempt, after failure

385 giveup: Callable[[AttemptOutcome], object | None] = no_giveup, # stop retrying based on domain-specific logic, e.g. time

386 before_attempt: Callable[[Retry], int | Awaitable[int]] = before_attempt_noop, # e.g rate limiting/internal backpressure

387 after_attempt: Callable[[AttemptOutcome], None | Awaitable[None]] = after_attempt_log_failure, # e.g. metrics/logging

388 on_retryable_error: Callable[[AttemptOutcome], None | Awaitable[None]] = noop, # e.g. count RetryableError failures

389 on_exhaustion: Callable[[AttemptOutcome], _T | Awaitable[_T]] = on_exhaustion_raise, # raise error or return fallback

390 log: logging.Logger | None = None, # set this to ``None`` to disable logging

391) -> _T:

392 """Async version of call_with_retries() with the same semantics except it awaits ``fn``, awaitable ``before_attempt`` /

393 ``after_attempt`` / ``on_retryable_error`` / ``on_exhaustion`` results and non-blocking sleep. Note that ``backoff``,

394 ``giveup``, and ``RetryTiming.is_terminated`` are not async as they are intentionally fast **synchronous** decisions;

395 they must not block."""

396 rng: random.Random | None = None

397 retry_count: int = 0

398 idle_nanos: int = 0

399 curr_max_sleep_nanos: int = policy.initial_max_sleep_nanos

400 previous_outcomes: tuple[AttemptOutcome, ...] = () # for safety pass *immutable* deque to callbacks

401 timing: RetryTiming = policy.timing

402 sleep: Callable[[int, Retry], Awaitable[None]] = timing.sleep_async

403 monotonic_ns: Callable[[], int] = timing.monotonic_ns

404 call_start_nanos: Final[int] = monotonic_ns()

405 while True:

406 before_attempt_nanos: int = monotonic_ns() if retry_count != 0 else call_start_nanos

407 prev: tuple[AttemptOutcome, ...] = previous_outcomes

408 retry: Retry = Retry(

409 retry_count, call_start_nanos, before_attempt_nanos, before_attempt_nanos, idle_nanos, policy, log, prev

410 )

411 try:

412 if before_attempt is not before_attempt_noop:

413 before_attempt_sleep_nanos: int = await _await_result(before_attempt(retry))

414 assert before_attempt_sleep_nanos >= 0, before_attempt_sleep_nanos

415 if before_attempt_sleep_nanos > 0:

416 await sleep(before_attempt_sleep_nanos, retry) # e.g. wait due to rate limiting or internal backpressure

417 attempt_start_nanos: int = monotonic_ns()

418 idle_nanos += attempt_start_nanos - before_attempt_nanos

419 retry = Retry(

420 retry_count, call_start_nanos, before_attempt_nanos, attempt_start_nanos, idle_nanos, policy, log, prev

421 )

422 timing.on_before_attempt(retry)

423 result: _T = await fn(retry) # Call the target function and supply retry attempt number and other metadata

424 except RetryableError as retryable_error:

425 elapsed_nanos: int = monotonic_ns() - call_start_nanos

426 is_terminated: Callable[[Retry], bool] = timing.is_terminated

427 giveup_reason: object | None = None

428 sleep_nanos: int = 0

429 outcome = AttemptOutcome(retry, False, False, False, None, elapsed_nanos, sleep_nanos, retryable_error)

430 await _await_result(on_retryable_error(outcome)) # e.g. count RetryableError failures

431 if retry_count < policy.max_retries and elapsed_nanos < policy.max_elapsed_nanos:

432 if policy.max_sleep_nanos == 0 and backoff is full_jitter_backoff_strategy:

433 pass # perf: e.g. spin-before-block

434 elif retry_count == 0 and retryable_error.retry_immediately_once:

435 pass # retry once immediately without backoff

436 else: # jitter: default backoff strategy picks random sleep_nanos in [min_sleep_nanos, curr_max_sleep_nanos]

437 rng = _thread_local_rng() if rng is None else rng

438 sleep_nanos, curr_max_sleep_nanos = backoff( # compute delay before next retry attempt, after failure

439 BackoffContext(retry, curr_max_sleep_nanos, rng, elapsed_nanos, retryable_error)

440 )

441 assert sleep_nanos >= 0 and curr_max_sleep_nanos >= 0, sleep_nanos

442

443 if sleep_nanos > 0:

444 outcome = AttemptOutcome(retry, False, False, False, None, elapsed_nanos, sleep_nanos, retryable_error)

445 if (not is_terminated(retry)) and (giveup_reason := giveup(outcome)) is None:

446 await _await_result(after_attempt(outcome))

447 await sleep(sleep_nanos, retry)

448 idle_nanos += sleep_nanos

449 if not is_terminated(retry):

450 previous_outcomes = _update_previous_outcomes(previous_outcomes, outcome, policy, retry)

451 del outcome # help gc

452 retry_count += 1

453 continue # continue retry loop with next attempt

454 else:

455 sleep_nanos = 0

456 outcome = AttemptOutcome(

457 retry, False, True, is_terminated(retry), giveup_reason, elapsed_nanos, sleep_nanos, retryable_error

458 )

459 await _await_result(after_attempt(outcome))

460 return await _await_result(on_exhaustion(outcome)) # raise error or return fallback value

461 else: # success

462 if after_attempt is not after_attempt_log_failure:

463 elapsed_nanos = monotonic_ns() - call_start_nanos

464 outcome = AttemptOutcome(retry, True, False, False, None, elapsed_nanos, 0, result)

465 await _await_result(after_attempt(outcome))

466 return result

467

468

469async def _await_result(result: _T | Awaitable[_T]) -> _T:

470 """Returns a callback result, awaiting it first when it is awaitable."""

471 if inspect.isawaitable(result):

472 return await result

473 else:

474 return result

475

476

477def multi_after_attempt(handlers: Iterable[Callable[[AttemptOutcome], None]]) -> Callable[[AttemptOutcome], None]:

478 """Composes independent ``after_attempt`` handlers into one ``call_with_retries(after_attempt=...)`` callback that

479 invokes each handler in order; thread-safe."""

480 handlers = tuple(handlers)

481 if len(handlers) == 1:

482 return handlers[0] # perf

483

484 def _after_attempt(outcome: AttemptOutcome) -> None:

485 for handler in handlers:

486 handler(outcome)

487

488 return _after_attempt

489

490

491def multi_after_attempt_async(

492 handlers: Iterable[Callable[[AttemptOutcome], None | Awaitable[None]]],

493) -> Callable[[AttemptOutcome], Awaitable[None]]:

494 """Async version of ``multi_after_attempt()``."""

495 handlers = tuple(handlers)

496

497 async def _after_attempt_async(outcome: AttemptOutcome) -> None:

498 for handler in handlers:

499 await _await_result(handler(outcome))

500

501 return _after_attempt_async

502

503

504def any_giveup(handlers: Iterable[Callable[[AttemptOutcome], object | None]]) -> Callable[[AttemptOutcome], object | None]:

505 """Composes independent ``giveup`` handlers into one ``call_with_retries(giveup=...)`` callback that gives up retrying if

506 *any* handler gives up; that is if any handler returns a non-``None`` reason; thread-safe.

507

508 Handlers are evaluated in order and short-circuit: On giving up returns the first handler's reason for giving up.

509 """

510 handlers = tuple(handlers)

511 if len(handlers) == 1:

512 return handlers[0] # perf

513

514 def _giveup(outcome: AttemptOutcome) -> object | None:

515 for handler in handlers:

516 giveup_reason: object | None = handler(outcome)

517 if giveup_reason is not None:

518 return giveup_reason

519 return None # don't give up retrying

520

521 return _giveup

522

523

524def all_giveup(handlers: Iterable[Callable[[AttemptOutcome], object | None]]) -> Callable[[AttemptOutcome], object | None]:

525 """Composes independent ``giveup`` handlers into one ``call_with_retries(giveup=...)`` callback that gives up retrying if

526 *all* handlers give up; that is if all handlers return a non-``None`` reason; thread-safe.

527

528 Handlers are evaluated in order and short-circuit: stops at first ``None``; else returns the last non-``None`` reason.

529 """

530 handlers = tuple(handlers)

531 if len(handlers) == 1:

532 return handlers[0] # perf

533

534 def _giveup(outcome: AttemptOutcome) -> object | None:

535 giveup_reason: object | None = None

536 for handler in handlers:

537 giveup_reason = handler(outcome)

538 if giveup_reason is None:

539 return None # don't give up retrying

540 return giveup_reason

541

542 return _giveup

543

544

545#############################################################################

546class RetryableError(Exception):

547 """Indicates that the task that caused the underlying exception can be retried and might eventually succeed;

548 ``call_with_retries()`` will pass this exception to callbacks via ``AttemptOutcome.result``; can be subclassed."""

549

550 def __init__(

551 self,

552 *exc_args: object, # optional args passed into super().__init__()

553 display_msg: object = None, # for logging

554 retry_immediately_once: bool = False, # retry once immediately without backoff?

555 category: object = None, # optional classification e.g. "CONCURRENCY, "SERVER_ISSUE", "THROTTLING", "TRANSIENT", ...

556 attachment: object = None, # optional domain specific info passed to next attempt via Retry.previous_outcomes if

557 # RetryPolicy.max_previous_outcomes > 0. This helps when retrying is not just 'try again later', but

558 # 'try again differently based on what just happened'.

559 # Examples: switching network endpoints, adjusting per-attempt timeouts, capping retries by error-class, resuming

560 # with a token/offset, maintaining failure history for this invocation of call_with_retries().

561 # Example: 'cap retries to 3 for ECONNREFUSED but 12 for ETIMEDOUT' via attachment=collections.Counter

562 ) -> None:

563 super().__init__(*exc_args)

564 self.display_msg: object = display_msg

565 self.retry_immediately_once: bool = retry_immediately_once

566 self.category: object = category

567 self.attachment: object = attachment

568

569 def display_msg_str(self) -> str:

570 """Returns the display_msg as a str; for logging."""

571 return "" if self.display_msg is None else str(self.display_msg)

572

573

574#############################################################################

575@final

576class RetryError(Exception):

577 """Indicates that retries have been exhausted; the last RetryableError is in RetryError.__cause__."""

578

579 outcome: Final[AttemptOutcome]

580 """Metadata that describes why and how call_with_retries() gave up."""

581

582 def __init__(self, outcome: AttemptOutcome) -> None:

583 super().__init__(outcome)

584 self.outcome = outcome

585

586

587#############################################################################

588@final

589class Retry(NamedTuple):

590 """Attempt metadata provided to callback functions; includes the current retry attempt number; immutable."""

591

592 count: int # type: ignore[assignment]

593 """Attempt number; count=0 is the first attempt, count=1 is the second attempt aka first retry."""

594

595 call_start_time_nanos: int

596 """Value of time.monotonic_ns() at start of call_with_retries() invocation."""

597

598 before_attempt_start_time_nanos: int

599 """Value of time.monotonic_ns() at start of before_attempt() invocation."""

600

601 attempt_start_time_nanos: int

602 """Value of time.monotonic_ns() at start of fn() invocation."""

603

604 idle_nanos: int

605 """Sum of all before_attempt_sleep_nanos() plus AttemptOutcome.sleep_nanos across this call_with_retries() invocation, at

606 the start of fn() invocation."""

607

608 policy: RetryPolicy

609 """Policy that was passed into call_with_retries()."""

610

611 log: logging.Logger | None

612 """Logger that was passed into call_with_retries()."""

613

614 previous_outcomes: Sequence[AttemptOutcome]

615 """History/state of the N=max_previous_outcomes most recent outcomes for the current call_with_retries() invocation."""

616

617 def copy(self, **override_kwargs: Any) -> Retry:

618 """Creates a new object copying an existing one with the specified fields overridden for customization."""

619 return self._replace(**override_kwargs)

620

621 def before_attempt_sleep_nanos(self) -> int:

622 """Returns duration between the start of before_attempt() and the start of fn() attempt."""

623 return self.attempt_start_time_nanos - self.before_attempt_start_time_nanos

624

625 def __repr__(self) -> str:

626 return (

627 f"{type(self).__name__}(count={self.count!r}, call_start_time_nanos={self.call_start_time_nanos!r}, "

628 f"before_attempt_start_time_nanos={self.before_attempt_start_time_nanos!r}, "

629 f"attempt_start_time_nanos={self.attempt_start_time_nanos!r}, idle_nanos={self.idle_nanos!r})"

630 )

631

632 def __eq__(self, other: object) -> bool:

633 return self is other

634

635 def __hash__(self) -> int:

636 return object.__hash__(self)

637

638

639#############################################################################

640@final

641class AttemptOutcome(NamedTuple):

642 """Captures per-attempt state for ``after_attempt`` callbacks; immutable."""

643

644 retry: Retry

645 """Attempt metadata passed into fn(retry)."""

646

647 is_success: bool

648 """False if fn(retry) raised a RetryableError; True otherwise."""

649

650 is_exhausted: bool

651 """True if the loop is giving up retrying (possibly even due to is_terminated); False otherwise."""

652

653 is_terminated: bool

654 """True if the termination predicate has become true; False otherwise."""

655

656 giveup_reason: object | None

657 """Reason returned by giveup(); None means giveup() was not called or decided to not give up."""

658

659 elapsed_nanos: int

660 """Total duration between the start of call_with_retries() invocation and the end of this fn() attempt."""

661

662 sleep_nanos: int

663 """Duration of current sleep period."""

664

665 result: RetryableError | object

666 """Result of fn(retry); a RetryableError on retryable failure, or some other object on success."""

667

668 def attempt_elapsed_nanos(self) -> int:

669 """Returns duration between the start of this fn() attempt and the end of this fn() attempt."""

670 return self.elapsed_nanos + self.retry.call_start_time_nanos - self.retry.attempt_start_time_nanos

671

672 def copy(self, **override_kwargs: Any) -> AttemptOutcome:

673 """Creates a new outcome copying an existing one with the specified fields overridden for customization."""

674 return self._replace(**override_kwargs)

675

676 def __repr__(self) -> str:

677 return (

678 f"{type(self).__name__}("

679 f"retry={self.retry!r}, "

680 f"is_success={self.is_success!r}, "

681 f"is_exhausted={self.is_exhausted!r}, "

682 f"is_terminated={self.is_terminated!r}, "

683 f"giveup_reason={self.giveup_reason!r}, "

684 f"elapsed_nanos={self.elapsed_nanos!r}, "

685 f"sleep_nanos={self.sleep_nanos!r})"

686 )

687

688 def __eq__(self, other: object) -> bool:

689 return self is other

690

691 def __hash__(self) -> int:

692 return object.__hash__(self)

693

694

695#############################################################################

696@final

697class BackoffContext(NamedTuple):

698 """Captures per-backoff state for ``backoff`` callbacks."""

699

700 retry: Retry

701 """Attempt metadata passed into fn(retry)."""

702

703 curr_max_sleep_nanos: int

704 """Current maximum duration (in nanoseconds) to sleep before the next retry attempt;

705 Typically: ``RetryPolicy.initial_max_sleep_nanos <= curr_max_sleep_nanos <= RetryPolicy.max_sleep_nanos``."""

706

707 rng: random.Random

708 """Thread-local random number generator instance."""

709

710 elapsed_nanos: int

711 """Total duration between the start of call_with_retries() invocation and the end of this fn() attempt."""

712

713 retryable_error: RetryableError

714 """Result of failed fn(retry) attempt."""

715

716 def copy(self, **override_kwargs: Any) -> BackoffContext:

717 """Creates a new object copying an existing one with the specified fields overridden for customization."""

718 return self._replace(**override_kwargs)

719

720 def __repr__(self) -> str:

721 return (

722 f"{type(self).__name__}("

723 f"retry={self.retry!r}, "

724 f"curr_max_sleep_nanos={self.curr_max_sleep_nanos!r}, "

725 f"elapsed_nanos={self.elapsed_nanos!r})"

726 )

727

728 def __eq__(self, other: object) -> bool:

729 return self is other

730

731 def __hash__(self) -> int:

732 return object.__hash__(self)

733

734

735BackoffStrategy = Callable[[BackoffContext], tuple[int, int]] # typealias; returns sleep_nanos:int, curr_max_sleep_nanos:int

736"""Strategy that implements a backoff algorithm that reduces server load while minimizing retry latency; default is full

737jitter; various other example backoff strategies such as decorrelated-jitter or retry-after HTTP 429 "Too Many Requests"

738responses, etc can be found in test_retry_examples.py."""

739

740

741#############################################################################

742def _default_timing_is_terminated(retry: Retry) -> bool:

743 return False

744

745

746def _default_timing_sleep(sleep_nanos: int, retry: Retry) -> None:

747 time.sleep(sleep_nanos / 1_000_000_000)

748

749

750async def _default_timing_sleep_asyncio(sleep_nanos: int, retry: Retry) -> None:

751 import asyncio

752

753 await asyncio.sleep(sleep_nanos / 1_000_000_000)

754

755

756def _default_timing_on_before_attempt(retry: Retry) -> None:

757 if retry.policy.timing.is_terminated(retry):

758 raise RetryableError(display_msg="terminated before attempt") from RetryTerminationError()

759

760

761@final

762class RetryTerminationError(InterruptedError):

763 """Termination signal raised when retry loop exits before starting the next attempt."""

764

765

766@dataclass(frozen=True)

767@final

768class RetryTiming:

769 """Customizable callbacks for reading the current monotonic time, sleeping and optional async termination; immutable."""

770

771 monotonic_ns: Callable[[], int] = time.monotonic_ns

772 """Returns the system's current monotonic time in nanoseconds."""

773

774 is_terminated: Callable[[Retry], bool] = _default_timing_is_terminated

775 """Returns whether a predicate has become true; if so causes the retry loop to exit early between attempts; can be used

776 to indicate system shutdown or similar cancellation conditions; default is to always return ``False``; this function

777 should complete quickly without any blocking or sleeping."""

778

779 sleep: Callable[[int, Retry], None] = _default_timing_sleep

780 """Sleeps N nanoseconds between attempts; override to inject custom sleeping or for early wake-ups; thread-safe."""

781

782 sleep_async: Callable[[int, Retry], Awaitable[None]] = _default_timing_sleep_asyncio

783 """Sleeps N nanoseconds between attempts; override to inject custom sleeping or for early wake-ups; thread-safe."""

784

785 on_before_attempt: Callable[[Retry], None] = _default_timing_on_before_attempt

786 """Typically (but not necessarily) raises an error if ``is_terminated()`` is True; otherwise fn() will still run; this

787 function should complete quickly without any blocking or sleeping.

788

789 To disable this behavior: RetryTiming.make_from(...).copy(on_before_attempt=lambda retry: None).

790 """

791

792 def copy(self, **override_kwargs: Any) -> RetryTiming:

793 """Creates a new object copying an existing one with the specified fields overridden for customization; thread-

794 safe."""

795 return dataclasses.replace(self, **override_kwargs)

796

797 @staticmethod

798 def make_from(termination_event: threading.Event | None) -> RetryTiming:

799 """Convenience factory that creates a RetryTiming that performs async termination when termination_event is set."""

800 if termination_event is None:

801 return RetryTiming()

802

803 def _is_terminated(retry: Retry) -> bool:

804 return termination_event.is_set()

805

806 def _sleep(sleep_nanos: int, retry: Retry) -> None:

807 termination_event.wait(sleep_nanos / 1_000_000_000) # allow early wakeup on async termination

808

809 return RetryTiming(is_terminated=_is_terminated, sleep=_sleep)

810

811 @staticmethod

812 def make_from_asyncio(termination_event: asyncio.Event | None) -> RetryTiming:

813 """Convenience factory that creates a RetryTiming that performs async termination when termination_event is set;

814 Write an analog version of this function if you wish to replace asyncio with Trio or AnyIO or similar."""

815 if termination_event is None:

816 return RetryTiming()

817

818 def _is_terminated(retry: Retry) -> bool:

819 return termination_event.is_set()

820

821 async def _sleep_async(sleep_nanos: int, retry: Retry) -> None:

822 import asyncio

823

824 if sleep_nanos <= 0:

825 await asyncio.sleep(0) # perf: cooperative yield with less overhead than asyncio.wait_for(..., 0)

826 return

827 try:

828 await asyncio.wait_for(termination_event.wait(), timeout=sleep_nanos / 1_000_000_000)

829 except asyncio.TimeoutError:

830 pass # expected

831

832 return RetryTiming(is_terminated=_is_terminated, sleep_async=_sleep_async)

833

834

835#############################################################################

836def _format_msg(display_msg: str, retryable_error: RetryableError) -> str:

837 """Default implementation of ``format_msg`` callback for RetryConfig; creates simple log message; thread-safe."""

838 msg = display_msg + " " if display_msg else ""

839 errmsg: str = retryable_error.display_msg_str()

840 msg = msg + errmsg + " " if errmsg else msg

841 msg = msg if msg else "Retrying "

842 return msg

843

844

845def _format_pair(first: object, second: object) -> str:

846 """Default implementation of ``format_pair`` callback for RetryConfig; creates simple log message part; thread-safe."""

847 second = "∞" if INFINITY_MAX_RETRIES == second else second # noqa: SIM300

848 return f"[{first}/{second}]"

849

850

851@dataclass(frozen=True)

852@final

853class RetryConfig:

854 """Configures logging for call_with_retries(); all defaults work out of the box; immutable."""

855

856 display_msg: str = "Retrying" # message prefix for retry log messages

857 dots: str = " ..." # suffix appended to retry log messages

858 format_msg: Callable[[str, RetryableError], str] = _format_msg # lambda: display_msg, retryable_error

859 format_pair: Callable[[object, object], str] = _format_pair # lambda: first, second

860 format_duration: Callable[[int], str] = human_readable_duration # lambda: nanos

861 info_loglevel: int = logging.INFO # loglevel used when not giving up

862 warning_loglevel: int = logging.WARNING # loglevel used when giving up

863 enable_logging: bool = True # set to False to disable logging

864 exc_info: bool = False # passed into Logger.log()

865 stack_info: bool = False # passed into Logger.log()

866 extra: Mapping[str, object] | None = dataclasses.field(default=None, repr=False, compare=False) # passed to Logger.log()

867 context: object = dataclasses.field(default=None, repr=False, compare=False) # optional domain specific info

868

869 def copy(self, **override_kwargs: Any) -> RetryConfig:

870 """Creates a new config copying an existing one with the specified fields overridden for customization."""

871 return dataclasses.replace(self, **override_kwargs)

872

873

874#############################################################################

875@dataclass(frozen=True)

876@final

877class RetryPolicy:

878 """Configuration of maximum retries, sleep bounds, elapsed-time budget, logging, etc for call_with_retries(); immutable.

879

880 By default uses full jitter which works as follows: The maximum duration to sleep between attempts initially starts with

881 ``initial_max_sleep_secs`` and doubles on each retry, up to the final maximum of ``max_sleep_secs``.

882 Example: 125ms --> 250ms --> 500ms --> 1s --> 2s --> 4s --> 8s --> 10s --> 10s...

883 On each retry a random sleep duration in the range ``[min_sleep_secs, current max]`` is picked.

884 In a nutshell: ``0 <= min_sleep_secs <= initial_max_sleep_secs <= max_sleep_secs``. Typically, min_sleep_secs=0.

885 """

886

887 max_retries: int = INFINITY_MAX_RETRIES

888 """The maximum number of times ``fn`` will be invoked additionally after the first attempt invocation; must be >= 0."""

889

890 min_sleep_secs: float = 0

891 """The minimum duration to sleep between any two attempts."""

892

893 initial_max_sleep_secs: float = 0.125

894 """The initial maximum duration to sleep between any two attempts."""

895

896 max_sleep_secs: float = 10

897 """The final max duration to sleep between any two attempts; 0 <= min_sleep_secs <= initial_max_sleep_secs <=

898 max_sleep_secs."""

899

900 max_elapsed_secs: float = 47

901 """``fn`` will not be retried (or not retried anymore) once this much time has elapsed since the initial start of

902 call_with_retries(); set this to 365 * 86400 seconds or similar to effectively disable the time limit."""

903

904 exponential_base: float = 2

905 """Growth factor (aka multiplier) for backoff algorithm to calculate sleep duration; must be >= 1."""

906

907 max_elapsed_nanos: int = dataclasses.field(init=False, repr=False) # derived value

908 min_sleep_nanos: int = dataclasses.field(init=False, repr=False) # derived value

909 initial_max_sleep_nanos: int = dataclasses.field(init=False, repr=False) # derived value

910 max_sleep_nanos: int = dataclasses.field(init=False, repr=False) # derived value

911

912 reraise: bool = True

913 """On exhaustion, the default (``True``) is to re-raise the underlying exception when present."""

914

915 max_previous_outcomes: int = 0

916 """Pass the N=max_previous_outcomes most recent AttemptOutcome objects to callbacks via Retry.previous_outcomes."""

917

918 config: RetryConfig = dataclasses.field(default=RetryConfig(), repr=False, compare=False)

919 """Configures logging behavior."""

920

921 timing: RetryTiming = dataclasses.field(default=RetryTiming(), repr=False)

922 """Customizable callbacks for reading the current monotonic time, sleeping and optional async termination."""

923

924 context: object = dataclasses.field(default=None, repr=False, compare=False)

925 """Optional domain specific info."""

926

927 @classmethod

928 def from_namespace(cls, args: argparse.Namespace) -> RetryPolicy:

929 """Factory that reads the policy from argparse.ArgumentParser via args."""

930 return cls(

931 max_retries=getattr(args, "max_retries", INFINITY_MAX_RETRIES),

932 min_sleep_secs=getattr(args, "retry_min_sleep_secs", 0),

933 initial_max_sleep_secs=getattr(args, "retry_initial_max_sleep_secs", 0.125),

934 max_sleep_secs=getattr(args, "retry_max_sleep_secs", 10),

935 max_elapsed_secs=getattr(args, "retry_max_elapsed_secs", 47),

936 exponential_base=getattr(args, "retry_exponential_base", 2),

937 reraise=getattr(args, "retry_reraise", True),

938 max_previous_outcomes=getattr(args, "retry_max_previous_outcomes", 0),

939 config=getattr(args, "retry_config", RetryConfig()),

940 timing=getattr(args, "retry_timing", RetryTiming()),

941 context=getattr(args, "retry_context", None),

942 )

943

944 @classmethod

945 def no_retries(cls) -> RetryPolicy:

946 """Returns a policy that never retries."""

947 return cls(

948 max_retries=0,

949 min_sleep_secs=0,

950 initial_max_sleep_secs=0,

951 max_sleep_secs=0,

952 max_elapsed_secs=0,

953 )

954

955 def __post_init__(self) -> None: # validate and compute derived values

956 self._validate_min("max_retries", self.max_retries, 0)

957 self._validate_min("exponential_base", self.exponential_base, 1)

958 self._validate_min("min_sleep_secs", self.min_sleep_secs, 0)

959 self._validate_min("initial_max_sleep_secs", self.initial_max_sleep_secs, 0)

960 self._validate_min("max_sleep_secs", self.max_sleep_secs, 0)

961 self._validate_min("max_elapsed_secs", self.max_elapsed_secs, 0)

962 object.__setattr__(self, "max_elapsed_nanos", int(self.max_elapsed_secs * 1_000_000_000)) # derived value

963 min_sleep_nanos: int = int(self.min_sleep_secs * 1_000_000_000)

964 initial_max_sleep_nanos: int = int(self.initial_max_sleep_secs * 1_000_000_000)

965 max_sleep_nanos: int = int(self.max_sleep_secs * 1_000_000_000)

966 max_sleep_nanos = max(min_sleep_nanos, max_sleep_nanos)

967 initial_max_sleep_nanos = min(max_sleep_nanos, max(min_sleep_nanos, initial_max_sleep_nanos))

968 object.__setattr__(self, "min_sleep_nanos", min_sleep_nanos) # derived value

969 object.__setattr__(self, "initial_max_sleep_nanos", initial_max_sleep_nanos) # derived value

970 object.__setattr__(self, "max_sleep_nanos", max_sleep_nanos) # derived value

971 self._validate_min("max_previous_outcomes", self.max_previous_outcomes, 0)

972 assert 0 <= self.min_sleep_nanos <= self.initial_max_sleep_nanos <= self.max_sleep_nanos

973 if not isinstance(self.reraise, bool):

974 raise TypeError(f"{type(self).__name__}.reraise must be bool")

975

976 def _validate_min(self, attr_name: str, value: float, minimum: float) -> None:

977 if value < minimum:

978 raise ValueError(f"Invalid {type(self).__name__}.{attr_name}: must be >= {minimum} but got {value}")

979

980 def copy(self, **override_kwargs: Any) -> RetryPolicy:

981 """Creates a new policy copying an existing one with the specified fields overridden for customization; thread-safe.

982

983 Example usage: policy = retry_policy.copy(max_sleep_secs=2, max_elapsed_secs=10)

984 """

985 return dataclasses.replace(self, **override_kwargs)

986

987

988#############################################################################

989def _fn_not_implemented(_retry: Retry) -> NoReturn:

990 """Default implementation of ``fn`` callback for RetryTemplate; always raises."""

991 raise NotImplementedError("Provide fn when calling RetryTemplate")

992

993

994NO_LOGGER: Final[logging.Logger] = logging.Logger("NULL") # noqa: LOG001 do not register dummy logger with Logger.manager

995NO_LOGGER.addHandler(logging.NullHandler()) # prevents lastResort fallback

996NO_LOGGER.disabled = True

997NO_LOGGER.propagate = False

998_R = TypeVar("_R")

999

1000

1001@dataclass(frozen=True)

1002@final

1003class RetryTemplate(Generic[_T]):

1004 """Convenience class that aggregates all knobs for call_with_retries(); and is itself callable too; immutable."""

1005

1006 fn: Callable[[Retry], _T] = _fn_not_implemented # set this to make the RetryTemplate object itself callable

1007 policy: RetryPolicy = RetryPolicy() # specifies how ``RetryableError`` shall be retried

1008 backoff: BackoffStrategy = full_jitter_backoff_strategy # computes delay time before next retry attempt, after failure

1009 giveup: Callable[[AttemptOutcome], object | None] = no_giveup # stop retrying based on domain-specific logic, e.g. time

1010 before_attempt: Callable[[Retry], int] = before_attempt_noop # e.g. wait due to rate limiting or internal backpressure

1011 after_attempt: Callable[[AttemptOutcome], None] = after_attempt_log_failure # e.g. record metrics and/or custom logging

1012 on_retryable_error: Callable[[AttemptOutcome], None] = noop # e.g. count failures (RetryableError) caught by retry loop

1013 on_exhaustion: Callable[[AttemptOutcome], _T] = on_exhaustion_raise # raise error or return fallback value

1014 log: logging.Logger | None = None # set this to ``None`` to disable logging

1015

1016 def copy(self, **override_kwargs: Any) -> RetryTemplate[_T]:

1017 """Creates a new object copying an existing one with the specified fields overridden for customization; thread-safe.

1018

1019 Example usage: retry_template.copy(policy=policy.copy(max_sleep_secs=2, max_elapsed_secs=10), log=None)

1020 """

1021 return dataclasses.replace(self, **override_kwargs)

1022

1023 def __call__(self) -> _T:

1024 """Invokes ``self.fn`` via the call_with_retries() retry loop using the stored parameters; thread-safe.

1025

1026 Example Usage: result: str = retry_template.copy(fn=...)()

1027 """

1028 return call_with_retries(

1029 fn=self.fn,

1030 policy=self.policy,

1031 backoff=self.backoff,

1032 giveup=self.giveup,

1033 before_attempt=self.before_attempt,

1034 after_attempt=self.after_attempt,

1035 on_retryable_error=self.on_retryable_error,

1036 on_exhaustion=self.on_exhaustion,

1037 log=self.log,

1038 )

1039

1040 def call_with_retries(

1041 self,

1042 fn: Callable[[Retry], _R],

1043 policy: RetryPolicy | None = None,

1044 *,

1045 backoff: BackoffStrategy | None = None,

1046 giveup: Callable[[AttemptOutcome], object | None] | None = None,

1047 before_attempt: Callable[[Retry], int] | None = None,

1048 after_attempt: Callable[[AttemptOutcome], None] | None = None,

1049 on_retryable_error: Callable[[AttemptOutcome], None] | None = None,

1050 on_exhaustion: Callable[[AttemptOutcome], _R] | None = None,

1051 log: logging.Logger | None = None, # pass NO_LOGGER to override template logger and disable logging for this call

1052 ) -> _R:

1053 """Invokes ``fn`` via the call_with_retries() retry loop using the stored or overridden params; thread-safe.

1054

1055 Example Usage: result: str = retry_template.call_with_retries(fn=...)

1056 """

1057 return call_with_retries(

1058 fn=fn,

1059 policy=self.policy if policy is None else policy,

1060 backoff=self.backoff if backoff is None else backoff,

1061 giveup=self.giveup if giveup is None else giveup,

1062 before_attempt=self.before_attempt if before_attempt is None else before_attempt,

1063 after_attempt=self.after_attempt if after_attempt is None else after_attempt,

1064 on_retryable_error=self.on_retryable_error if on_retryable_error is None else on_retryable_error,

1065 on_exhaustion=(

1066 cast(Callable[[AttemptOutcome], _R], self.on_exhaustion) if on_exhaustion is None else on_exhaustion

1067 ),

1068 log=None if log is NO_LOGGER else self.log if log is None else log,

1069 )

1070

1071 def wraps(self, fn: Callable[..., _R]) -> Callable[..., _R]:

1072 """Returns a wrapper function that forwards all arguments to ``fn`` and retries it using this template; thread-safe.

1073

1074 Example Usage:

1075 def fn(x: int) -> int:

1076 return x * 2

1077 func: Callable[[int], int] = retry_template.wraps(fn)

1078 y: int = func(5) # returns 10

1079 """

1080

1081 @functools.wraps(fn)

1082 def wrapped(*args: Any, **kwargs: Any) -> _R:

1083 return self.call_with_retries(fn=lambda _retry: fn(*args, **kwargs))

1084

1085 return wrapped

1086

1087

1088#############################################################################

1089@dataclass(frozen=True)

1090@final

1091class AsyncRetryTemplate(Generic[_T]):

1092 """Async version of ``RetryTemplate`` with the same semantics except it awaits ``fn``, awaitable ``before_attempt`` /

1093 ``after_attempt`` / ``on_retryable_error`` / ``on_exhaustion`` results and non-blocking sleep. Note that ``backoff``,

1094 ``giveup``, and ``RetryTiming.is_terminated`` are not async as they are intentionally fast **synchronous** decisions;

1095 they must not block."""

1096

1097 fn: Callable[[Retry], Awaitable[_T]] = _fn_not_implemented # set this to make the RetryTemplate object itself callable

1098 policy: RetryPolicy = RetryPolicy() # specifies how ``RetryableError`` shall be retried

1099 backoff: BackoffStrategy = full_jitter_backoff_strategy # computes delay time before next retry attempt, after failure

1100 giveup: Callable[[AttemptOutcome], object | None] = no_giveup # stop retrying based on domain-specific logic, e.g. time

1101 before_attempt: Callable[[Retry], int | Awaitable[int]] = before_attempt_noop # e.g. rate limiting/internal backpressure

1102 after_attempt: Callable[[AttemptOutcome], None | Awaitable[None]] = after_attempt_log_failure # e.g. metrics/logging

1103 on_retryable_error: Callable[[AttemptOutcome], None | Awaitable[None]] = noop # e.g. count RetryableError failures

1104 on_exhaustion: Callable[[AttemptOutcome], _T | Awaitable[_T]] = on_exhaustion_raise # raise or fallback

1105 log: logging.Logger | None = None # set this to ``None`` to disable logging

1106

1107 def copy(self, **override_kwargs: Any) -> AsyncRetryTemplate[_T]:

1108 """Creates a new object copying an existing one with the specified fields overridden for customization; thread-safe.

1109

1110 Example usage: retry_template.copy(policy=policy.copy(max_sleep_secs=2, max_elapsed_secs=10), log=None)

1111 """

1112 return dataclasses.replace(self, **override_kwargs)

1113

1114 async def __call__(self) -> _T:

1115 """Invokes ``self.fn`` via the call_with_retries_async() retry loop using the stored parameters; thread-safe.

1116

1117 Example Usage: result: str = await retry_template.copy(fn=...)()

1118 """

1119 return await call_with_retries_async(

1120 fn=self.fn,

1121 policy=self.policy,

1122 backoff=self.backoff,

1123 giveup=self.giveup,

1124 before_attempt=self.before_attempt,

1125 after_attempt=self.after_attempt,

1126 on_retryable_error=self.on_retryable_error,

1127 on_exhaustion=self.on_exhaustion,

1128 log=self.log,

1129 )

1130

1131 async def call_with_retries(

1132 self,

1133 fn: Callable[[Retry], Awaitable[_R]],

1134 policy: RetryPolicy | None = None,

1135 *,

1136 backoff: BackoffStrategy | None = None,

1137 giveup: Callable[[AttemptOutcome], object | None] | None = None,

1138 before_attempt: Callable[[Retry], int | Awaitable[int]] | None = None,

1139 after_attempt: Callable[[AttemptOutcome], None | Awaitable[None]] | None = None,

1140 on_retryable_error: Callable[[AttemptOutcome], None | Awaitable[None]] | None = None,

1141 on_exhaustion: Callable[[AttemptOutcome], _R | Awaitable[_R]] | None = None,

1142 log: logging.Logger | None = None, # pass NO_LOGGER to override template logger and disable logging for this call

1143 ) -> _R:

1144 """Invokes ``fn`` via the call_with_retries_async() retry loop using the stored or overridden params; thread-safe.

1145

1146 Example Usage: result: str = await retry_template.call_with_retries(fn=...)

1147 """

1148 return await call_with_retries_async(

1149 fn=fn,

1150 policy=self.policy if policy is None else policy,

1151 backoff=self.backoff if backoff is None else backoff,

1152 giveup=self.giveup if giveup is None else giveup,

1153 before_attempt=self.before_attempt if before_attempt is None else before_attempt,

1154 after_attempt=self.after_attempt if after_attempt is None else after_attempt,

1155 on_retryable_error=self.on_retryable_error if on_retryable_error is None else on_retryable_error,

1156 on_exhaustion=(

1157 cast(Callable[[AttemptOutcome], Union[_R, Awaitable[_R]]], self.on_exhaustion)

1158 if on_exhaustion is None

1159 else on_exhaustion

1160 ),

1161 log=None if log is NO_LOGGER else self.log if log is None else log,

1162 )

1163

1164 def wraps(self, fn: Callable[..., Awaitable[_R]]) -> Callable[..., Awaitable[_R]]:

1165 """Returns a wrapper function that forwards all arguments to ``fn`` and retries it using this template; thread-safe.

1166

1167 Example Usage:

1168 async def fn(x: int) -> int:

1169 return x * 2

1170 func: Callable[[int], Awaitable[int]] = retry_template.wraps(fn)

1171 y: int = await func(5) # returns 10

1172 """

1173

1174 @functools.wraps(fn)

1175 async def wrapped_async(*args: Any, **kwargs: Any) -> _R:

1176 return await self.call_with_retries(fn=lambda _retry: fn(*args, **kwargs))

1177

1178 return wrapped_async

1179

1180

1181#############################################################################

1182def raise_retryable_error_from(

1183 exc: BaseException,

1184 *,

1185 display_msg: object = None,

1186 retry_immediately_once: bool = False,

1187 category: object = None,

1188 attachment: object = None,

1189) -> NoReturn:

1190 """Convenience function that raises a generic RetryableError that wraps the given underlying exception."""

1191 raise RetryableError(

1192 display_msg=type(exc).__name__ if display_msg is None else display_msg,

1193 retry_immediately_once=retry_immediately_once,

1194 category=category,

1195 attachment=attachment,

1196 ) from exc

1197

1198

1199ExceptionPredicate = Union[bool, Callable[[BaseException], bool]] # Type alias

1200

1201

1202def call_with_exception_handlers(

1203 fn: Callable[[], _T], # typically a lambda

1204 *,

1205 continue_scanning_if_no_predicate_matches: bool = False,

1206 handlers: Mapping[type[BaseException], Sequence[tuple[ExceptionPredicate, Callable[[BaseException], _T]]]],

1207) -> _T:

1208 """Convenience function that calls ``fn`` and returns its result; on exception runs the first matching handler in a per-

1209 exception handler chain; composes independent handlers via predicates into one function, in Event-Predicate-Action style.

1210

1211 Lookup uses the exception type's Method Resolution Order (most-specific class in the exception class hierarchy wins). For

1212 the first class that exists as a key in ``handlers``, its chain is scanned in order. Each chain element is

1213 ``(predicate, handler)`` where ``predicate`` is either ``True`` (always matches), ``False`` (disabled), or

1214 ``predicate(exc) -> bool``. The first matching handler is called with the exception and its return value is returned. If

1215 no predicate matches then, by default, the original exception is re-raised and no less-specific handler chains are

1216 consulted. Set ``continue_scanning_if_no_predicate_matches=True`` to continue scanning exception base classes instead.

1217

1218 Typically (but not necessarily) the handler raises a ``RetryableError``, via ``raise_retryable_error_from`` or similar.

1219 Or it may raise another exception type (which will not be retried), or even return a fallback value instead of raising.

1220

1221 Example: turn transient ssh/zfs command failures into RetryableError for call_with_retries(), including feature flags:

1222

1223 def run_remote(retry: Retry) -> str:

1224 p = subprocess.run(["ssh", "foo.example.com", "zfs", "list", "-H"], text=True, capture_output=True, check=True)

1225 return p.stdout

1226

1227 def fn(retry: Retry) -> str:

1228 return call_with_exception_handlers(

1229 fn=lambda: run_remote(retry),

1230 handlers={

1231 TimeoutError: [(True, raise_retryable_error_from)],

1232 ConnectionResetError: [(True, lambda exc: raise_retryable_error_from(exc, display_msg="ssh reset"))],

1233 subprocess.CalledProcessError: [

1234 (lambda exc: exc.returncode == 255, lambda exc: raise_retryable_error_from(exc, display_msg="ssh error")),

1235 (lambda exc: "cannot receive" in (exc.stderr or ""), lambda exc: raise_retryable_error_from(exc, display_msg="zfs recv")),

1236 ],

1237 OSError: [

1238 (lambda exc: getattr(exc, "errno", None) in {errno.ETIMEDOUT, errno.EHOSTUNREACH},

1239 lambda exc: raise_retryable_error_from(exc, display_msg=f"network: {exc}")),

1240 (False, lambda exc: raise_retryable_error_from(exc, display_msg="disabled handler example")),

1241 ],

1242 },

1243 )

1244

1245 stdout: str = call_with_retries(fn=fn, policy=RetryPolicy(max_retries=3))

1246

1247 Example: return a fallback value (no retry loop required):

1248

1249 def read_optional_file(path: str) -> str:

1250 return call_with_exception_handlers(

1251 fn=lambda: open(path, encoding="utf-8").read(),

1252 handlers={FileNotFoundError: [(True, lambda _exc: "")]},

1253 )

1254 """

1255 try:

1256 return fn()

1257 except BaseException as exc:

1258 for cls in type(exc).__mro__:

1259 handler_chain = handlers.get(cls)

1260 if handler_chain is not None:

1261 for predicate, handler in handler_chain:

1262 if predicate is True or (predicate is not False and predicate(exc)):

1263 return handler(exc)

1264 if not continue_scanning_if_no_predicate_matches:

1265 raise

1266 raise

1267

1268

1269#############################################################################

1270@final

1271class _ThreadLocalRNG(threading.local):

1272 """Caches a per-thread random number generator."""

1273

1274 def __init__(self) -> None:

1275 self.rng: random.Random | None = None

1276

1277

1278_THREAD_LOCAL_RNG: Final[_ThreadLocalRNG] = _ThreadLocalRNG()

1279

1280

1281def _thread_local_rng() -> random.Random:

1282 """Returns a per-thread RNG for backoff jitter; for perf avoids locking and initializing a new random.Random() at high

1283 frequency."""

1284 threadlocal: _ThreadLocalRNG = _THREAD_LOCAL_RNG

1285 rng: random.Random | None = threadlocal.rng

1286 if rng is None:

1287 rng = random.Random() # noqa: S311 jitter isn't security sensitive, and random.SystemRandom.randint() is slow

1288 threadlocal.rng = rng

1289 return rng