Coverage for async-durable-execution/src/async_durable_execution/config.py: 98%

1"""Configuration types.""" 

2 

3from __future__ import annotations 

4 

5import math 

6import random 

7import re 

8from dataclasses import dataclass, field 

9from datetime import timedelta 

10from enum import Enum 

11from typing import TypeAlias 

12 

13from .exceptions import ValidationError 

14 

15Duration: TypeAlias = int | timedelta 

16 

17 

18def duration_to_seconds(duration: Duration, field_name: str = "duration") -> int: 

19 """Convert a seconds integer or timedelta to whole seconds.""" 

20 if isinstance(duration, bool) or not isinstance(duration, int | timedelta): 

21 msg = f"{field_name} must be an int number of seconds or a timedelta" 

22 raise ValidationError(msg) 

23 

24 seconds = ( 

25 int(duration.total_seconds()) if isinstance(duration, timedelta) else duration 

26 ) 

27 if seconds < 0: 

28 msg = f"{field_name} must be non-negative" 

29 raise ValidationError(msg) 

30 return seconds 

31 

32 

33class JitterStrategy(str, Enum): 

34 """ 

35 Jitter strategies are used to introduce noise when attempting to retry 

36 an invoke. We introduce noise to prevent a thundering-herd effect where 

37 a group of accesses (e.g. invokes) happen at once. 

38 

39 Jitter is meant to be used to spread operations across time. 

40 

41 Based on AWS Architecture Blog: https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/ 

42 

43 members: 

44 :NONE: No jitter; use the exact calculated delay 

45 :FULL: Full jitter; random delay between 0 and calculated delay 

46 :HALF: Equal jitter; random delay between 0.5x and 1.0x of the calculated delay 

47 """ 

48 

49 NONE = "NONE" 

50 FULL = "FULL" 

51 HALF = "HALF" 

52 

53 def apply_jitter(self, delay: float) -> float: 

54 """Apply jitter to a delay value and return the final delay. 

55 

56 Args: 

57 delay: The base delay value to apply jitter to 

58 

59 Returns: 

60 The final delay after applying jitter strategy 

61 """ 

62 match self: 

63 case JitterStrategy.NONE: 

64 return delay 

65 case JitterStrategy.HALF: 

66 # Equal jitter: delay/2 + random(0, delay/2) 

67 return delay / 2 + random.random() * (delay / 2) # noqa: S311 

68 case _: # default is FULL 

69 # Full jitter: random(0, delay) 

70 return random.random() * delay # noqa: S311 

71 

72 def finalize_delay(self, base_delay: float) -> int: 

73 """Apply jitter, round up, and clamp to a minimum of 1 second.""" 

74 return max(1, math.ceil(self.apply_jitter(base_delay))) 

75 

76 

77@dataclass 

78class _DelayStrategy: 

79 """Common delay configuration for retry and polling strategies.""" 

80 

81 max_attempts: int = 6 

82 initial_delay: Duration = 5 

83 max_delay: Duration = 60 

84 backoff_rate: int | float = 2 

85 jitter_strategy: JitterStrategy = field(default=JitterStrategy.FULL) 

86 increment: Duration | None = None 

87 

88 def __post_init__(self): 

89 self.initial_delay = duration_to_seconds(self.initial_delay, "initial_delay") 

90 self.max_delay = duration_to_seconds(self.max_delay, "max_delay") 

91 if self.increment is not None: 

92 self.increment = duration_to_seconds(self.increment, "increment") 

93 

94 @property 

95 def initial_delay_seconds(self) -> int: 

96 """Get initial delay in seconds.""" 

97 return duration_to_seconds(self.initial_delay, "initial_delay") 

98 

99 @property 

100 def max_delay_seconds(self) -> int: 

101 """Get max delay in seconds.""" 

102 return duration_to_seconds(self.max_delay, "max_delay") 

103 

104 @property 

105 def increment_seconds(self) -> int | None: 

106 """Get linear delay increment in seconds.""" 

107 if self.increment is None: 

108 return None 

109 return duration_to_seconds(self.increment, "increment") 

110 

111 def calculate_delay(self, attempts_made: int) -> int: 

112 """Calculate a whole-second delay for exponential or linear strategies.""" 

113 increment_seconds = self.increment_seconds 

114 if increment_seconds is None: 

115 base_delay: float = self.initial_delay_seconds * ( 

116 self.backoff_rate ** (attempts_made - 1) 

117 ) 

118 else: 

119 base_delay = self.initial_delay_seconds + increment_seconds * ( 

120 attempts_made - 1 

121 ) 

122 

123 return self.jitter_strategy.finalize_delay( 

124 min(base_delay, self.max_delay_seconds) 

125 ) 

126 

127 

128@dataclass 

129class RetryStrategy(_DelayStrategy): 

130 """Exponential-backoff retry strategy for durable operations.""" 

131 

132 retryable_errors: list[str | re.Pattern] | None = None 

133 retryable_error_types: list[type[Exception]] | None = None 

134 

135 def __call__(self, error: Exception, attempts_made: int) -> Duration | None: 

136 """Return retry delay, or None if the error should not be retried.""" 

137 default_retryable_error_pattern = re.compile(r".*") 

138 should_use_default_errors: bool = ( 

139 self.retryable_errors is None and self.retryable_error_types is None 

140 ) 

141 

142 retryable_errors: list[str | re.Pattern] = ( 

143 self.retryable_errors 

144 if self.retryable_errors is not None 

145 else ( 

146 [default_retryable_error_pattern] if should_use_default_errors else [] 

147 ) 

148 ) 

149 retryable_error_types: list[type[Exception]] = self.retryable_error_types or [] 

150 

151 if attempts_made >= self.max_attempts: 

152 return None 

153 

154 is_retryable_error_message: bool = any( 

155 pattern.search(str(error)) 

156 if isinstance(pattern, re.Pattern) 

157 else pattern in str(error) 

158 for pattern in retryable_errors 

159 ) 

160 is_retryable_error_type: bool = any( 

161 isinstance(error, error_type) for error_type in retryable_error_types 

162 ) 

163 

164 if not is_retryable_error_message and not is_retryable_error_type: 

165 return None 

166 

167 return self.calculate_delay(attempts_made) 

168 

169 @classmethod 

170 def none(cls) -> RetryStrategy: 

171 """No retries.""" 

172 return cls(max_attempts=1, max_delay=300, backoff_rate=2.0) 

173 

174 @classmethod 

175 def default(cls) -> RetryStrategy: 

176 """Default retries, used automatically when no retry strategy is provided.""" 

177 return cls() 

178 

179 @classmethod 

180 def transient(cls) -> RetryStrategy: 

181 """Quick retries for transient errors.""" 

182 return cls( 

183 max_attempts=3, 

184 max_delay=300, 

185 backoff_rate=2, 

186 jitter_strategy=JitterStrategy.HALF, 

187 ) 

188 

189 @classmethod 

190 def resource_availability(cls) -> RetryStrategy: 

191 """Longer retries for resource availability.""" 

192 return cls( 

193 max_attempts=5, 

194 initial_delay=5, 

195 max_delay=300, 

196 backoff_rate=2, 

197 ) 

198 

199 @classmethod 

200 def critical(cls) -> RetryStrategy: 

201 """Aggressive retries for critical operations.""" 

202 return cls( 

203 max_attempts=10, 

204 initial_delay=1, 

205 max_delay=60, 

206 backoff_rate=1.5, 

207 jitter_strategy=JitterStrategy.NONE, 

208 ) 

209 

210 @classmethod 

211 def linear(cls) -> RetryStrategy: 

212 """Linearly increasing delay between retries: 1s, 2s, 3s, 4s, 5s.""" 

213 return cls( 

214 max_attempts=6, 

215 initial_delay=1, 

216 increment=1, 

217 max_delay=300, 

218 backoff_rate=2, 

219 jitter_strategy=JitterStrategy.NONE, 

220 )