Coverage for src/beautyspot/limiter.py: 92%

1# src/beautyspot/limiter.py

2

3 import time

4 import asyncio

5 import threading

6 from typing import Protocol, runtime_checkable

7

8

9@runtime_checkable

10 class LimiterProtocol(Protocol):

11 def consume(self, cost: int) -> None: ...

12

13 async def consume_async(self, cost: int) -> None: ...

14

15

16 class Gcra(LimiterProtocol):

17 """

18 A smooth rate limiter based on the GCRA (Generic Cell Rate Algorithm).

19

20 Features:

21 - No burst after long idle (Strict Pacing).

22 - No start-up delay for the very first request.

23 - Fails fast if a task cost exceeds the TPM limit.

24 - Thread-safe and Async-compatible.

25 - Uses monotonic clock.

26 """

27

28 def __init__(self, tokens_per_minute: int):

29 if tokens_per_minute <= 0:

30 raise ValueError("tokens_per_minute must be positive")

31

32 # Rate: tokens per second

33 self.rate = float(tokens_per_minute) / 60.0

34

35 # Maximum allowed cost per task.

36 # A single task consuming more than the TPM limit is physically impossible

37 # to process within the rate window, so it should be forbidden.

38 self.max_cost = int(tokens_per_minute)

39

40 # Theoretical Arrival Time (TAT)

41 self.tat = time.monotonic()

42 self.lock = threading.Lock()

43

44 def _consume_reservation(self, cost: int) -> float:

45 """

46 Calculates wait time and updates TAT.

47 Returns seconds to wait.

48 """

49 if cost <= 0: 49 ↛ 50line 49 didn't jump to line 50 because the condition on line 49 was never true

50 return 0.0

51

52 # Guard: Prevent requests that exceed the rate limit capacity entirely

53 if cost > self.max_cost:

54 raise ValueError(

55 f"Requested cost ({cost}) exceeds the maximum limit of {self.max_cost} tokens per minute. "

56 "This task cannot be processed within the defined rate limit."

57 )

58

59 increment = cost / self.rate

60

61 with self.lock:

62 now = time.monotonic()

63 if now > self.tat:

64 self.tat = now

65

66 wait_time = self.tat - now

67 if wait_time < 0: 67 ↛ 68line 67 didn't jump to line 68 because the condition on line 67 was never true

68 wait_time = 0.0

69

70 self.tat += increment

71

72 return wait_time

73

74 def consume(self, cost: int):

75 """

76 Acquire tokens from the bucket, blocking if necessary.

77

78 If the bucket does not have enough tokens, this method sleeps (blocks the thread)

79 until the tokens become available based on the refill rate.

80

81 Args:

82 cost (int): Number of tokens to consume.

83

84 Raises:

85 ValueError: If `cost` exceeds the bucket's total capacity (`tpm`).

86 (i.e., the request is too expensive to ever be processed)

87 """

88 wait_time = self._consume_reservation(cost)

89 if wait_time > 0:

90 time.sleep(wait_time)

91

92 async def consume_async(self, cost: int):

93 """

94 Acquire tokens asynchronously.

95

96 If the bucket does not have enough tokens, this method awaits (non-blocking sleep)

97 until the tokens become available.

98

99 Args:

100 cost (int): Number of tokens to consume.

101

102 Raises:

103 ValueError: If `cost` exceeds the bucket's total capacity.

104 """

105 wait_time = self._consume_reservation(cost)

106 if wait_time > 0:

107 await asyncio.sleep(wait_time)