Coverage for src / ss_utils_circuit_breaker / breaker.py: 20%

1"""Async circuit breaker with optional Redis (coredis) support.""" 

2 

3from datetime import UTC, datetime 

4from enum import Enum 

5from functools import wraps 

6from typing import Any 

7 

8 

9class CircuitState(str, Enum): 

10 """Circuit breaker states.""" 

11 

12 CLOSED = "closed" 

13 OPEN = "open" 

14 HALF_OPEN = "half_open" 

15 

16 

17class CircuitBreakerError(Exception): 

18 """Raised when circuit breaker is open and preventing calls.""" 

19 

20 def __init__(self, breaker_name: str) -> None: 

21 super().__init__(f"Circuit breaker '{breaker_name}' is open") 

22 self.breaker_name = breaker_name 

23 

24 

25class CircuitBreaker: 

26 """Async circuit breaker with in-memory state storage. 

27 

28 Args: 

29 name: Unique name for this circuit breaker 

30 failure_threshold: Number of failures before opening circuit 

31 recovery_timeout: Seconds to wait before transitioning to half-open 

32 expected_exception: Exception type(s) that count as failures 

33 

34 Example: 

35 >>> breaker = CircuitBreaker("my_service", failure_threshold=3, recovery_timeout=30) 

36 >>> 

37 >>> @breaker 

38 >>> async def call_external_api(): 

39 >>> # Your async API call here 

40 >>> pass 

41 """ 

42 

43 def __init__( 

44 self, 

45 name: str, 

46 failure_threshold: int = 5, 

47 recovery_timeout: int = 60, 

48 expected_exception: type[Exception] | tuple[type[Exception], ...] = Exception, 

49 ) -> None: 

50 self.name = name 

51 self.failure_threshold = failure_threshold 

52 self.recovery_timeout = recovery_timeout 

53 self.expected_exception = expected_exception 

54 

55 # In-memory state 

56 self._state = CircuitState.CLOSED 

57 self._failure_count = 0 

58 self._opened_at: datetime | None = None 

59 

60 @property 

61 def state(self) -> CircuitState: 

62 """Get current state, checking for automatic half-open transition.""" 

63 if self._state == CircuitState.OPEN and self._opened_at: 

64 elapsed = (datetime.now(UTC) - self._opened_at).total_seconds() 

65 if elapsed >= self.recovery_timeout: 

66 self._state = CircuitState.HALF_OPEN 

67 return self._state 

68 

69 @property 

70 def failure_count(self) -> int: 

71 """Get current failure count.""" 

72 return self._failure_count 

73 

74 def open(self) -> None: 

75 """Open the circuit breaker.""" 

76 self._state = CircuitState.OPEN 

77 self._opened_at = datetime.now(UTC) 

78 

79 def close(self) -> None: 

80 """Close the circuit breaker and reset failure count.""" 

81 self._state = CircuitState.CLOSED 

82 self._failure_count = 0 

83 self._opened_at = None 

84 

85 def half_open(self) -> None: 

86 """Transition to half-open state.""" 

87 self._state = CircuitState.HALF_OPEN 

88 

89 def reset(self) -> None: 

90 """Reset the circuit breaker to initial closed state.""" 

91 self.close() 

92 

93 def _on_success(self) -> None: 

94 """Handle successful call.""" 

95 if self._state == CircuitState.HALF_OPEN: 

96 self._state = CircuitState.CLOSED 

97 self._failure_count = 0 

98 self._opened_at = None 

99 

100 def _on_failure(self) -> None: 

101 """Handle failed call.""" 

102 self._failure_count += 1 

103 if self._state == CircuitState.HALF_OPEN: 

104 self.open() 

105 elif self._failure_count >= self.failure_threshold: 

106 self.open() 

107 

108 def __call__(self, func: Any) -> Any: 

109 """Decorator for async functions.""" 

110 

111 @wraps(func) 

112 async def wrapper(*args: Any, **kwargs: Any) -> Any: 

113 # Check if circuit is open 

114 if self.state == CircuitState.OPEN: 

115 raise CircuitBreakerError(self.name) 

116 

117 try: 

118 result = await func(*args, **kwargs) 

119 self._on_success() 

120 return result 

121 except self.expected_exception: 

122 self._on_failure() 

123 raise 

124 

125 return wrapper 

126 

127 

128class RedisCircuitBreaker(CircuitBreaker): 

129 """Async circuit breaker with Redis state storage via coredis. 

130 

131 Requires the `redis` extra: `pip install ss-utils-circuit-breaker[redis]` 

132 

133 Args: 

134 name: Unique name for this circuit breaker 

135 redis_url: Redis connection URL (format: "host:port") 

136 namespace: Redis key namespace prefix 

137 failure_threshold: Number of failures before opening circuit 

138 recovery_timeout: Seconds to wait before transitioning to half-open 

139 expected_exception: Exception type(s) that count as failures 

140 

141 Example: 

142 >>> breaker = RedisCircuitBreaker( 

143 >>> "my_service", 

144 >>> redis_url="localhost:6379", 

145 >>> failure_threshold=3, 

146 >>> ) 

147 >>> 

148 >>> @breaker 

149 >>> async def call_external_api(): 

150 >>> # Your async API call here 

151 >>> pass 

152 """ 

153 

154 def __init__( 

155 self, 

156 name: str, 

157 redis_url: str | None = None, 

158 namespace: str = "circuit_breaker", 

159 **kwargs: Any, 

160 ) -> None: 

161 super().__init__(name=name, **kwargs) 

162 self._redis_url = redis_url 

163 self._namespace = namespace 

164 self._redis: Any = None 

165 self._redis_available = False 

166 

167 def _key(self, suffix: str) -> str: 

168 """Generate Redis key with namespace.""" 

169 return f"{self._namespace}:{self.name}:{suffix}" 

170 

171 async def _ensure_redis(self) -> bool: 

172 """Ensure Redis connection is available (lazy initialization).""" 

173 if self._redis is not None: 

174 return self._redis_available 

175 

176 if not self._redis_url: 

177 self._redis_available = False 

178 return False 

179 

180 try: 

181 import coredis 

182 

183 # Parse redis_url (format: "host:port") 

184 parts = self._redis_url.split(":") 

185 expected_parts = 2 

186 if len(parts) != expected_parts: 

187 self._redis_available = False 

188 return False 

189 

190 host, port_str = parts 

191 port = int(port_str) 

192 

193 self._redis = coredis.Redis(host=host, port=port) 

194 # Test connection 

195 await self._redis.ping() 

196 self._redis_available = True 

197 except Exception: 

198 self._redis_available = False 

199 

200 return self._redis_available 

201 

202 async def _get_state_from_redis(self) -> CircuitState | None: 

203 """Get state from Redis.""" 

204 if not await self._ensure_redis(): 

205 return None 

206 

207 try: 

208 state_bytes = await self._redis.get(self._key("state")) 

209 if state_bytes: 

210 return CircuitState(state_bytes.decode("utf-8")) 

211 except Exception: 

212 pass 

213 return None 

214 

215 async def _set_state_in_redis(self, state: CircuitState) -> None: 

216 """Set state in Redis.""" 

217 if not await self._ensure_redis(): 

218 return 

219 

220 try: 

221 await self._redis.set(self._key("state"), state.value) 

222 except Exception: 

223 pass 

224 

225 async def _get_failure_count_from_redis(self) -> int | None: 

226 """Get failure count from Redis.""" 

227 if not await self._ensure_redis(): 

228 return None 

229 

230 try: 

231 count_bytes = await self._redis.get(self._key("failure_count")) 

232 if count_bytes: 

233 return int(count_bytes.decode("utf-8")) 

234 except Exception: 

235 pass 

236 return None 

237 

238 async def _set_failure_count_in_redis(self, count: int) -> None: 

239 """Set failure count in Redis.""" 

240 if not await self._ensure_redis(): 

241 return 

242 

243 try: 

244 await self._redis.set(self._key("failure_count"), str(count)) 

245 except Exception: 

246 pass 

247 

248 async def _get_opened_at_from_redis(self) -> datetime | None: 

249 """Get opened_at timestamp from Redis.""" 

250 if not await self._ensure_redis(): 

251 return None 

252 

253 try: 

254 ts_bytes = await self._redis.get(self._key("opened_at")) 

255 if ts_bytes: 

256 return datetime.fromisoformat(ts_bytes.decode("utf-8")) 

257 except Exception: 

258 pass 

259 return None 

260 

261 async def _set_opened_at_in_redis(self, dt: datetime | None) -> None: 

262 """Set opened_at timestamp in Redis.""" 

263 if not await self._ensure_redis(): 

264 return 

265 

266 try: 

267 if dt: 

268 await self._redis.set(self._key("opened_at"), dt.isoformat()) 

269 else: 

270 await self._redis.delete(self._key("opened_at")) 

271 except Exception: 

272 pass 

273 

274 async def _sync_from_redis(self) -> None: 

275 """Sync local state from Redis.""" 

276 state = await self._get_state_from_redis() 

277 if state: 

278 self._state = state 

279 

280 count = await self._get_failure_count_from_redis() 

281 if count is not None: 

282 self._failure_count = count 

283 

284 opened_at = await self._get_opened_at_from_redis() 

285 if opened_at: 

286 self._opened_at = opened_at 

287 

288 async def _sync_to_redis(self) -> None: 

289 """Sync local state to Redis.""" 

290 await self._set_state_in_redis(self._state) 

291 await self._set_failure_count_in_redis(self._failure_count) 

292 await self._set_opened_at_in_redis(self._opened_at) 

293 

294 async def get_state_async(self) -> CircuitState: 

295 """Get current state asynchronously (checks Redis if available).""" 

296 await self._sync_from_redis() 

297 

298 # Check for automatic half-open transition 

299 if self._state == CircuitState.OPEN and self._opened_at: 

300 elapsed = (datetime.now(UTC) - self._opened_at).total_seconds() 

301 if elapsed >= self.recovery_timeout: 

302 self._state = CircuitState.HALF_OPEN 

303 await self._set_state_in_redis(self._state) 

304 

305 return self._state 

306 

307 async def _on_success_async(self) -> None: 

308 """Handle successful call (async version).""" 

309 self._on_success() 

310 await self._sync_to_redis() 

311 

312 async def _on_failure_async(self) -> None: 

313 """Handle failed call (async version).""" 

314 self._on_failure() 

315 await self._sync_to_redis() 

316 

317 def __call__(self, func: Any) -> Any: 

318 """Decorator for async functions with Redis support.""" 

319 

320 @wraps(func) 

321 async def wrapper(*args: Any, **kwargs: Any) -> Any: 

322 # Get state (syncs from Redis if available) 

323 current_state = await self.get_state_async() 

324 

325 if current_state == CircuitState.OPEN: 

326 raise CircuitBreakerError(self.name) 

327 

328 try: 

329 result = await func(*args, **kwargs) 

330 await self._on_success_async() 

331 return result 

332 except self.expected_exception: 

333 await self._on_failure_async() 

334 raise 

335 

336 return wrapper 

337