Coverage for src/taipanstack/utils/rate_limit.py: 100%

1""" 

2Rate limiting utilities. 

3 

4Provides an in-memory token-bucket based rate limiting decorator 

5for both synchronous and asynchronous functions. The decorator 

6returns a ``Result`` type encapsulating the original return value 

7or a ``RateLimitError`` error. 

8""" 

9 

10import contextlib 

11import functools 

12import inspect 

13import math 

14import threading 

15import time 

16from collections.abc import Awaitable, Callable 

17from typing import ParamSpec, Protocol, TypeVar, cast, overload 

18 

19from taipanstack.core.result import Err, Ok, Result 

20 

21__all__ = ["RateLimitError", "RateLimiter", "rate_limit"] 

22 

23P = ParamSpec("P") 

24T = TypeVar("T") 

25 

26 

27class RateLimitError(Exception): 

28 """Exception raised when a rate limit is exceeded.""" 

29 

30 def __init__(self, message: str = "Rate limit exceeded") -> None: 

31 """Initialize the RateLimitError. 

32 

33 Args: 

34 message: The error message to display.Defaults to "Rate limit exceeded". 

35 

36 """ 

37 super().__init__(message) 

38 

39 

40class RateLimiter: 

41 """Token bucket rate limiter logic.""" 

42 

43 def _is_finite_number(self, val: float | int) -> bool: 

44 return isinstance(val, (int, float)) and math.isfinite(val) 

45 

46 def _validate_finite(self, max_calls: int, time_window: float) -> None: 

47 """Check if parameters are finite numbers.""" 

48 if not self._is_finite_number(max_calls) or not self._is_finite_number( 

49 time_window 

50 ): 

51 raise ValueError("max_calls and time_window must be finite numbers") 

52 

53 def _validate_positive(self, max_calls: int, time_window: float) -> None: 

54 """Check if parameters are positive.""" 

55 if max_calls <= 0: 

56 raise ValueError("max_calls and time_window must be > 0.0") 

57 if time_window <= 0: 

58 raise ValueError("max_calls and time_window must be > 0.0") 

59 

60 def _validate_init_params(self, max_calls: int, time_window: float) -> None: 

61 """Validate initialization parameters.""" 

62 self._validate_finite(max_calls, time_window) 

63 self._validate_positive(max_calls, time_window) 

64 

65 def __init__(self, max_calls: int, time_window: float) -> None: 

66 """Initialize the token bucket. 

67 

68 Args: 

69 max_calls: The maximum number of calls allowed in the time window. 

70 time_window: The time window in seconds. 

71 

72 """ 

73 self._validate_init_params(max_calls, time_window) 

74 self.capacity: float = float(max_calls) 

75 self.time_window: float = float(time_window) 

76 self.tokens: float = self.capacity 

77 self.last_update: float = time.monotonic() 

78 self._lock = threading.Lock() 

79 

80 def _is_valid_time_window(self) -> bool: 

81 """Check if time window is valid.""" 

82 if not isinstance(self.time_window, (int, float)): 

83 return False # type: ignore[unreachable] 

84 return math.isfinite(self.time_window) and self.time_window > 0.0 

85 

86 def _is_valid_capacity(self) -> bool: 

87 """Check if capacity is valid.""" 

88 if not isinstance(self.capacity, (int, float)): 

89 return False # type: ignore[unreachable] 

90 return math.isfinite(self.capacity) and self.capacity > 0.0 

91 

92 def _is_valid_bucket_state(self) -> bool: 

93 """Check if the bucket's time window and capacity are in a valid state.""" 

94 return self._is_valid_time_window() and self._is_valid_capacity() 

95 

96 def _calculate_new_tokens(self, elapsed: float) -> float | None: 

97 """Calculate new tokens based on elapsed time.""" 

98 try: 

99 new_tokens = elapsed * (self.capacity / self.time_window) 

100 return new_tokens if math.isfinite(new_tokens) else None 

101 except Exception: 

102 return None 

103 

104 def _validate_token_types(self, new_tokens: float) -> bool: 

105 """Validate the types of tokens.""" 

106 if not isinstance(self.tokens, (int, float)): 

107 self.tokens = self.capacity # type: ignore[unreachable] 

108 return False 

109 return isinstance(new_tokens, (int, float)) 

110 

111 def _update_tokens_value(self, new_tokens: float) -> bool: 

112 """Update tokens value safely.""" 

113 try: 

114 self.tokens += new_tokens 

115 if not math.isfinite(self.tokens): 

116 self.tokens = self.capacity 

117 return False 

118 self.tokens = min(self.tokens, self.capacity) 

119 return True 

120 except Exception: 

121 self.tokens = self.capacity 

122 return False 

123 

124 def _apply_new_tokens(self, new_tokens: float) -> bool: 

125 """Apply new tokens to the bucket.""" 

126 if not self._validate_token_types(new_tokens): 

127 return False 

128 return self._update_tokens_value(new_tokens) 

129 

130 def _calculate_elapsed(self, now: float) -> float | None: 

131 """Calculate the elapsed time since the last update.""" 

132 if not isinstance(self.last_update, (int, float)): 

133 return None # type: ignore[unreachable] 

134 try: 

135 raw_elapsed = now - self.last_update 

136 if not math.isfinite(raw_elapsed): 

137 return None 

138 return max(0.0, raw_elapsed) 

139 except Exception: 

140 return None 

141 

142 def _add_tokens(self, now: float) -> bool: 

143 """Calculate and add new tokens to the bucket based on elapsed time. 

144 

145 Args: 

146 now: Current monotonic time. 

147 

148 Returns: 

149 True if token update succeeds, False if state corruption is detected. 

150 

151 """ 

152 elapsed = self._calculate_elapsed(now) 

153 if elapsed is None: 

154 return False 

155 

156 # Prevent state corruption or infinite elapsed time 

157 if not self._is_valid_bucket_state(): 

158 return False 

159 

160 new_tokens = self._calculate_new_tokens(elapsed) 

161 if new_tokens is None: 

162 return False 

163 

164 if not self._apply_new_tokens(new_tokens): 

165 return False 

166 

167 self.last_update = now 

168 return True 

169 

170 def _check_tokens_validity(self) -> bool: 

171 """Check tokens validity.""" 

172 if not isinstance(self.tokens, (int, float)): 

173 return False # type: ignore[unreachable] 

174 if not math.isfinite(self.tokens): 

175 self.tokens = self.capacity 

176 return False 

177 return True 

178 

179 def _try_consume(self, tokens: float) -> bool: 

180 """Attempt to consume the tokens from the bucket if available.""" 

181 try: 

182 if not self._check_tokens_validity(): 

183 return False 

184 if self.tokens >= tokens: 

185 self.tokens -= tokens 

186 return True 

187 return False 

188 except Exception: 

189 self.tokens = self.capacity 

190 return False 

191 

192 def _get_current_time(self) -> float | None: 

193 """Get current monotonic time safely.""" 

194 try: 

195 return time.monotonic() 

196 except Exception: 

197 return None 

198 

199 def _validate_and_add_tokens(self, now: float | None) -> bool: 

200 """Validate time and add tokens. 

201 

202 Returns True if tokens were added successfully, False otherwise. 

203 """ 

204 if now is None: 

205 return False 

206 if not isinstance(now, (int, float)): 

207 return False # type: ignore[unreachable] 

208 if not math.isfinite(now): 

209 # If not finite, we can't add tokens, but we shouldn't fail 

210 # the consumption if there are already tokens. 

211 return True 

212 return self._add_tokens(now) 

213 

214 def _process_consumption(self, tokens: float) -> bool: 

215 """Process token consumption inside the lock.""" 

216 now = self._get_current_time() 

217 

218 # Prevent time corruption from poisoning the bucket state. 

219 # Only try to add tokens if time is finite. 

220 if not self._validate_and_add_tokens(now): 

221 return False 

222 

223 return self._try_consume(tokens) 

224 

225 def _is_valid_token_amount(self, tokens: float) -> bool: 

226 """Check if requested token amount is valid.""" 

227 if not isinstance(tokens, (int, float)): 

228 return False # type: ignore[unreachable] 

229 return math.isfinite(tokens) 

230 

231 def _acquire_lock_safe(self) -> bool: 

232 """Safely acquire lock.""" 

233 try: 

234 return self._lock.acquire(timeout=0.1) 

235 except Exception: 

236 return False 

237 

238 def consume(self, tokens: float = 1.0) -> bool: 

239 """Try to consume tokens. 

240 

241 Args: 

242 tokens: Number of tokens to consume. Defaults to 1.0. 

243 

244 Returns: 

245 True if tokens were consumed (allow), False otherwise (limit exceeded). 

246 

247 """ 

248 if not self._is_valid_token_amount(tokens): 

249 return False 

250 if tokens <= 0: 

251 return True 

252 

253 if not self._acquire_lock_safe(): 

254 return False 

255 try: 

256 return self._process_consumption(tokens) 

257 finally: 

258 with contextlib.suppress(RuntimeError): 

259 self._lock.release() 

260 

261 

262class RateLimitDecorator(Protocol): 

263 """Protocol for the rate limit decorator.""" 

264 

265 @overload 

266 def __call__( 

267 self, 

268 func: Callable[P, T], 

269 ) -> Callable[P, Result[T, RateLimitError]]: ... 

270 

271 @overload 

272 def __call__( 

273 self, 

274 func: Callable[P, Awaitable[T]], 

275 ) -> Callable[P, Awaitable[Result[T, RateLimitError]]]: ... 

276 

277 

278def _async_rate_limit_wrapper( 

279 func_coro: Callable[P, Awaitable[T]], 

280 limiter: RateLimiter, 

281) -> Callable[P, Awaitable[Result[T, RateLimitError]]]: 

282 @functools.wraps(func_coro) 

283 async def async_wrapper( 

284 *args: P.args, 

285 **kwargs: P.kwargs, 

286 ) -> Result[T, RateLimitError]: 

287 try: 

288 if not limiter.consume(): 

289 return Err(RateLimitError()) 

290 except Exception: 

291 return Err(RateLimitError()) 

292 return Ok(await func_coro(*args, **kwargs)) 

293 

294 return async_wrapper # type: ignore[misc] 

295 

296 

297def _sync_rate_limit_wrapper( 

298 func_sync: Callable[P, T], 

299 limiter: RateLimiter, 

300) -> Callable[P, Result[T, RateLimitError]]: 

301 @functools.wraps(func_sync) 

302 def wrapper(*args: P.args, **kwargs: P.kwargs) -> Result[T, RateLimitError]: 

303 try: 

304 if not limiter.consume(): 

305 return Err(RateLimitError()) 

306 except Exception: 

307 return Err(RateLimitError()) 

308 return Ok(func_sync(*args, **kwargs)) 

309 

310 return wrapper 

311 

312 

313def rate_limit( 

314 max_calls: int, 

315 time_window: float, 

316) -> RateLimitDecorator: 

317 """Decorate a function to apply rate limiting. 

318 

319 If the rate limit is exceeded, the wrapped function immediately returns 

320 an ``Err(RateLimitError)``. Uses an in-memory token bucket strategy. 

321 

322 Args: 

323 max_calls: Maximum function executions allowed in the defined window. 

324 time_window: Time window size in seconds. 

325 

326 Returns: 

327 Decorated function returning a ``Result[T, RateLimitError]``. 

328 

329 Example: 

330 >>> @rate_limit(max_calls=2, time_window=1.0) 

331 ... def fetch_data() -> str: 

332 ... return "data" 

333 >>> fetch_data() 

334 Ok('data') 

335 >>> fetch_data() 

336 Ok('data') 

337 >>> fetch_data() 

338 Err(RateLimitError('Rate limit exceeded')) 

339 

340 """ 

341 

342 def decorator( 

343 func: Callable[P, T] | Callable[P, Awaitable[T]], 

344 ) -> ( 

345 Callable[P, Result[T, RateLimitError]] 

346 | Callable[P, Awaitable[Result[T, RateLimitError]]] 

347 ): 

348 limiter = RateLimiter(max_calls, time_window) 

349 

350 if inspect.iscoroutinefunction(func): 

351 func_coro = cast(Callable[P, Awaitable[T]], func) 

352 return _async_rate_limit_wrapper(func_coro, limiter) 

353 

354 func_sync = cast(Callable[P, T], func) 

355 return _sync_rate_limit_wrapper(func_sync, limiter) 

356 

357 return cast(RateLimitDecorator, decorator)