Coverage for async_durable_execution/_operation/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 .._core import ( 

8 Duration, 

9 DurableContext, 

10 OperationSubType, 

11 RetryStrategy, 

12 SerDes, 

13 bind_current_context, 

14 get_current_context, 

15 get_durable_context, 

16) 

17from ..extension import get_extension_context 

18 

19if TYPE_CHECKING: 

20 from .child import SummaryGenerator 

21 

22T = TypeVar("T") 

23 

24 

25def wait( 

26 duration: Duration, 

27 *, 

28 name: str | None = None, 

29) -> asyncio.Task[None]: 

30 """Run an SDK-owned wait operation through the stable operation SPI.""" 

31 return ( 

32 get_extension_context() 

33 ._reserve_sdk_operation(name) # noqa: SLF001 

34 ._run_wait( # noqa: SLF001 

35 duration, 

36 sub_type=OperationSubType.WAIT, 

37 ) 

38 ) 

39 

40 

41def run_in_child_context( 

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

43 *, 

44 name: str | None = None, 

45 serdes: SerDes | None = None, 

46 summary_generator: SummaryGenerator | None = None, 

47 is_virtual: bool = False, 

48) -> asyncio.Task[T]: 

49 """Run an SDK-owned retry scope through the stable operation SPI.""" 

50 return ( 

51 get_extension_context() 

52 ._reserve_sdk_operation(name) # noqa: SLF001 

53 ._run_in_child_context( # noqa: SLF001 

54 func, 

55 sub_type=OperationSubType.RUN_IN_CHILD_CONTEXT, 

56 serdes=serdes, 

57 summary_generator=summary_generator, 

58 is_virtual=is_virtual, 

59 ) 

60 ) 

61 

62 

63@dataclass(frozen=True) 

64class WithRetryContext(DurableContext): 

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

66 

67 attempt: int = 1 

68 

69 

70def get_with_retry_context() -> WithRetryContext: 

71 """Return the active `WithRetryContext`.""" 

72 current_context = get_current_context() 

73 if not isinstance(current_context, WithRetryContext): 

74 msg = ( 

75 "get_with_retry_context() can only be used while a with_retry body " 

76 "is executing." 

77 ) 

78 raise RuntimeError(msg) 

79 return current_context 

80 

81 

82def with_retry( 

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

84 *, 

85 name: str | None = None, 

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

87 serdes: SerDes | None = None, 

88 summary_generator: SummaryGenerator | None = None, 

89 is_virtual: bool = False, 

90) -> asyncio.Task[T]: 

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

92 

93 Args: 

94 func: Async callable to retry. Use get_with_retry_context().attempt inside 

95 the callable to access the current attempt number. 

96 name: Optional durable operation name. 

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

98 serdes: Optional serializer for the child context result. 

99 summary_generator: Optional summary generator for large child results. 

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

101 """ 

102 

103 async def run_loop() -> T: 

104 retry = retry_strategy or RetryStrategy() 

105 attempt = 0 

106 while True: 

107 attempt += 1 

108 try: 

109 context = get_durable_context() 

110 retry_context = WithRetryContext( 

111 execution_state=context.execution_state, 

112 operation_identifier=context.operation_identifier, 

113 step_id_prefix=context.step_id_prefix, 

114 replaying=context.is_replaying(), 

115 attempt=attempt, 

116 ) 

117 if "step_counter" in context.__dict__: 

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

119 "step_counter" 

120 ] 

121 with bind_current_context(retry_context): 

122 return await func() 

123 except Exception as err: 

124 delay = retry(err, attempt) 

125 if delay is None: 

126 raise 

127 

128 wait_name = ( 

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

130 ) 

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

132 

133 return run_in_child_context( 

134 run_loop, 

135 name=name or "with-retry", 

136 serdes=serdes, 

137 summary_generator=summary_generator, 

138 is_virtual=is_virtual, 

139 )