Coverage for async-durable-execution/src/async_durable_execution/composite/with_retry.py: 100%

1from __future__ import annotations 

2 

3import asyncio 

4from dataclasses import dataclass 

5from typing import TYPE_CHECKING, Awaitable, Callable, TypeVar 

6 

7from ..config import Duration 

8from ..config import RetryStrategy 

9from ..context import bind_current_context 

10from ..primitive.child import ( 

11 DurableContext, 

12 get_durable_context, 

13 run_in_child_context, 

14) 

15from ..primitive.wait import wait 

16 

17if TYPE_CHECKING: 

18 from .parallel import SummaryGenerator 

19 from ..serdes import SerDes 

20 

21T = TypeVar("T") 

22 

23 

24@dataclass(frozen=True) 

25class WithRetryContext(DurableContext): 

26 """Context available while a with_retry body is executing.""" 

27 

28 attempt: int = 1 

29 

30 

31def with_retry( 

32 func: Callable[[], Awaitable[T]], 

33 *, 

34 name: str | None = None, 

35 retry_strategy: Callable[[Exception, int], Duration | None] | None = None, 

36 serdes: SerDes | None = None, 

37 summary_generator: SummaryGenerator | None = None, 

38 is_virtual: bool = False, 

39) -> asyncio.Task[T]: 

40 """Retry a block of durable logic with configurable backoff. 

41 

42 Args: 

43 func: Async callable to retry. Use get_current_context().attempt inside 

44 the callable to access the current attempt number. 

45 name: Optional durable operation name. 

46 retry_strategy: Optional strategy that returns a retry delay or None to stop. 

47 serdes: Optional serializer for the child context result. 

48 summary_generator: Optional summary generator for large child results. 

49 is_virtual: Whether the child context should skip lifecycle checkpoints. 

50 """ 

51 

52 async def run_loop() -> T: 

53 retry = retry_strategy or RetryStrategy() 

54 attempt = 0 

55 while True: 

56 attempt += 1 

57 try: 

58 context = get_durable_context("with_retry") 

59 retry_context = WithRetryContext( 

60 execution_state=context.execution_state, 

61 operation_identifier=context.operation_identifier, 

62 step_id_prefix=context.step_id_prefix, 

63 replaying=context.is_replaying(), 

64 attempt=attempt, 

65 ) 

66 if "step_counter" in context.__dict__: 

67 retry_context.__dict__["step_counter"] = context.__dict__[ 

68 "step_counter" 

69 ] 

70 with bind_current_context(retry_context): 

71 return await func() 

72 except Exception as err: 

73 delay = retry(err, attempt) 

74 if delay is None: 

75 raise 

76 

77 wait_name = ( 

78 f"{name}-backoff-{attempt}" if name else f"backoff-{attempt}" 

79 ) 

80 await wait(duration=delay, name=wait_name) 

81 

82 return run_in_child_context( 

83 run_loop, 

84 name=name or "with-retry", 

85 serdes=serdes, 

86 summary_generator=summary_generator, 

87 is_virtual=is_virtual, 

88 )