Coverage for qreward / utils / scheduler / limiter.py: 100.00%

1"""Sliding window rate limiter implementation.""" 

2 

3import asyncio (empty)

4import bisect (empty)

5import threading (empty)

6import time (empty)

7from typing import Optional (empty)

8 

9 

10class LimiterPool: (empty)

11 """Sliding window rate limiter with both sync and async support.""" 

12 

13 # Global lock for protecting global pool 

14 global_lock = threading.Lock() (empty)

15 

16 # Global pool dictionary, storing different pool instances by key 

17 global_limiter_pool: dict = {} (empty)

18 

19 @classmethod (empty)

20 def get_pool( (empty)

21 cls, key: str, rate: int, window: float 

22 ) -> Optional["LimiterPool"]: 

23 """Get or create a limiter pool instance for the given key. 

24 

25 Uses singleton pattern to ensure one pool per key. 

26 

27 Args: 

28 key: Unique identifier for the pool 

29 rate: Max requests allowed in window 

30 window: Window duration in seconds 

31 

32 Returns: 

33 LimiterPool instance or None if rate/window is invalid 

34 """ 

35 if rate <= 0 or window <= 0: (empty)

36 return None (empty)

37 with cls.global_lock: (empty)

38 if key not in cls.global_limiter_pool: (empty)

39 cls.global_limiter_pool[key] = cls(rate=rate, window=window) (empty)

40 return cls.global_limiter_pool[key] (empty)

41 

42 def __init__(self, rate: int, window: float, clock=time.monotonic): (empty)

43 """Initialize sliding window rate limiter. 

44 

45 Args: 

46 rate: Max requests allowed in window 

47 window: Window duration in seconds 

48 clock: Time function, defaults to time.monotonic 

49 """ 

50 if rate <= 0 or window <= 0: (empty)

51 raise ValueError("rate / window must be positive") (empty)

52 self.rate = rate (empty)

53 self.window = float(window) (empty)

54 self._clock = clock (empty)

55 # Timestamp list, monotonically increasing 

56 self._times: list[float] = [] (empty)

57 # Condition variable for efficient sync waiting (replaces busy-wait) 

58 self._condition = threading.Condition(threading.RLock()) (empty)

59 # Alias for backward compatibility — the underlying lock 

60 self._lock = self._condition (empty)

61 # Async lock 

62 self._aio_lock = asyncio.Lock() (empty)

63 

64 def allow(self, timeout: Optional[float] = None) -> bool: (empty)

65 """Synchronously acquire a token. 

66 

67 Uses threading.Condition to efficiently wait for token availability 

68 instead of busy-waiting with time.sleep. 

69 

70 Args: 

71 timeout: Optional timeout in seconds 

72 

73 Returns: 

74 True if token acquired, False if timeout 

75 """ 

76 deadline = None if timeout is None else self._clock() + timeout (empty)

77 with self._condition: (empty)

78 while True: (empty)

79 ok = self._check_and_add() (empty)

80 if ok: (empty)

81 self._condition.notify_all() (empty)

82 return True (empty)

83 sleep_t = self._sleep_time() (empty)

84 if deadline and self._clock() + sleep_t > deadline: (empty)

85 return False (empty)

86 self._condition.wait(timeout=sleep_t) (empty)

87 

88 async def async_allow(self, timeout: Optional[float] = None) -> bool: (empty)

89 """Asynchronously acquire a token. 

90 

91 Args: 

92 timeout: Optional timeout in seconds 

93 

94 Returns: 

95 True if token acquired, False if timeout 

96 """ 

97 deadline = None if timeout is None else self._clock() + timeout (empty)

98 while True: (empty)

99 async with self._aio_lock: (empty)

100 ok = self._check_and_add() (empty)

101 if ok: (empty)

102 return True (empty)

103 sleep_t = self._sleep_time() (empty)

104 if deadline and self._clock() + sleep_t > deadline: (empty)

105 return False (empty)

106 await asyncio.sleep(sleep_t) (empty)

107 

108 def _check_and_add(self) -> bool: (empty)

109 """Check if window is full and add timestamp if not. 

110 

111 Must be called while holding lock. 

112 

113 Returns: 

114 True if request allowed, False if window full 

115 """ 

116 now = self._clock() (empty)

117 cutoff = now - self.window (empty)

118 # Clean expired timestamps 

119 idx = bisect.bisect_left(self._times, cutoff) (empty)

120 self._times = self._times[idx:] (empty)

121 # Check if window is full 

122 if len(self._times) < self.rate: (empty)

123 bisect.insort(self._times, now) (empty)

124 return True (empty)

125 return False (empty)

126 

127 def _sleep_time(self) -> float: (empty)

128 """Calculate time to wait until earliest record expires. 

129 

130 Must be called while holding lock. 

131 

132 Returns: 

133 Sleep time in seconds 

134 """ 

135 if not self._times: (empty)

136 return 0.01 (empty)

137 earliest = self._times[0] (empty)

138 return max(0.0, earliest + self.window - self._clock()) (empty)