Coverage for src/case_router/limits.py: 100%

1"""Per-tenant request limits. 

2 

3Each tenant gets its own token bucket sized from its requests-per-minute allowance, so one 

4busy org cannot use up capacity that belongs to the others. 

5""" 

6 

7from __future__ import annotations 

8 

9import threading 

10import time 

11from collections.abc import Callable 

12from dataclasses import dataclass 

13 

14from .models import Tenant 

15 

16Clock = Callable[[], float] 

17 

18 

19@dataclass(frozen=True, slots=True) 

20class Acquisition: 

21 allowed: bool 

22 retry_after: float = 0.0 

23 

24 

25class TokenBucket: 

26 """Thread-safe token bucket. 

27 

28 Holds up to ``capacity`` tokens and refills continuously at ``refill_per_second``. 

29 Every request spends one token; with none left, the caller learns how long to wait. 

30 """ 

31 

32 def __init__(self, capacity: int, refill_per_second: float, clock: Clock = time.monotonic): 

33 if capacity < 1: 

34 raise ValueError("capacity must be at least 1") 

35 if refill_per_second <= 0: 

36 raise ValueError("refill_per_second must be positive") 

37 self.capacity = capacity 

38 self.refill_per_second = refill_per_second 

39 self._clock = clock 

40 self._tokens = float(capacity) 

41 self._updated = clock() 

42 self._lock = threading.Lock() 

43 

44 def _refill(self) -> None: 

45 now = self._clock() 

46 if now <= self._updated: 

47 return # no time has passed, or the clock stepped back: add nothing 

48 elapsed = now - self._updated 

49 self._tokens = min(float(self.capacity), self._tokens + elapsed * self.refill_per_second) 

50 self._updated = now 

51 

52 def try_acquire(self) -> Acquisition: 

53 with self._lock: 

54 self._refill() 

55 if self._tokens >= 1: 

56 self._tokens -= 1 

57 return Acquisition(allowed=True) 

58 return Acquisition( 

59 allowed=False, retry_after=(1 - self._tokens) / self.refill_per_second 

60 ) 

61 

62 @property 

63 def available(self) -> float: 

64 with self._lock: 

65 self._refill() 

66 return self._tokens 

67 

68 

69class TenantLimiter: 

70 """Keeps one bucket per tenant and rebuilds it if the tenant's allowance changes.""" 

71 

72 def __init__(self, clock: Clock = time.monotonic) -> None: 

73 self._clock = clock 

74 self._buckets: dict[str, TokenBucket] = {} 

75 self._lock = threading.Lock() 

76 

77 def bucket_for(self, tenant: Tenant) -> TokenBucket: 

78 with self._lock: 

79 bucket = self._buckets.get(tenant.id) 

80 if bucket is None or bucket.capacity != tenant.requests_per_minute: 

81 bucket = TokenBucket( 

82 capacity=tenant.requests_per_minute, 

83 refill_per_second=tenant.requests_per_minute / 60, 

84 clock=self._clock, 

85 ) 

86 self._buckets[tenant.id] = bucket 

87 return bucket 

88 

89 def acquire(self, tenant: Tenant) -> Acquisition: 

90 return self.bucket_for(tenant).try_acquire()