feat: add thread-safe token bucket rate limiter utility

samcli/lib/utils/rate_limiter.py

@@ -0,0 +1,102 @@
+"""
+Thread-safe rate limiter using the token bucket algorithm.
+
+Useful for throttling API calls to AWS services or other external endpoints
+to stay within service quotas.
+"""
+
+import threading
+import time
+
+
+class RateLimiter:
+    """
+    A token bucket rate limiter.
+
+    Parameters
+    ----------
+    rate : float
+        Number of tokens added per second.
+    burst : int
+        Maximum number of tokens the bucket can hold.
+
+    Examples
+    --------
+    >>> limiter = RateLimiter(rate=10, burst=10)
+    >>> limiter.acquire()  # blocks until a token is available
+    >>> limiter.try_acquire()  # returns True/False without blocking
+    """
+
+    def __init__(self, rate: float, burst: int):
+        if rate <= 0:
+            raise ValueError("rate must be positive")
+        if burst <= 0:
+            raise ValueError("burst must be a positive integer")
+
+        self._rate = float(rate)
+        self._burst = burst
+        self._tokens = float(burst)
+        self._last_refill = time.monotonic()
+        self._lock = threading.Lock()
+
+    def _refill(self):
+        now = time.monotonic()
+        elapsed = now - self._last_refill
+        self._tokens = min(self._burst, self._tokens + elapsed * self._rate)
+        self._last_refill = now
+
+    def try_acquire(self, tokens: int = 1) -> bool:
+        """
+        Try to consume tokens without blocking.
+
+        Parameters
+        ----------
+        tokens : int
+            Number of tokens to consume.
+
+        Returns
+        -------
+        bool
+            True if tokens were acquired, False otherwise.
+        """
+        with self._lock:
+            self._refill()
+            if self._tokens >= tokens:
+                self._tokens -= tokens
+                return True
+            return False
+
+    def acquire(self, tokens: int = 1, timeout: float = None) -> bool:
+        """
+        Block until tokens are available or timeout is reached.
+
+        Parameters
+        ----------
+        tokens : int
+            Number of tokens to consume.
+        timeout : float, optional
+            Maximum seconds to wait. None means wait indefinitely.
+
+        Returns
+        -------
+        bool
+            True if tokens were acquired, False if timed out.
+        """
+        deadline = None if timeout is None else time.monotonic() + timeout
+
+        while True:
+            with self._lock:
+                self._refill()
+                if self._tokens >= tokens:
+                    self._tokens -= tokens
+                    return True
+                deficit = tokens - self._tokens
+                wait_time = deficit / self._rate
+
+            if deadline is not None:
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    return False
+                wait_time = min(wait_time, remaining)
+
+            time.sleep(wait_time)

tests/unit/lib/utils/test_rate_limiter.py

@@ -0,0 +1,64 @@
+"""Tests for the rate limiter utility."""
+
+import threading
+import time
+from unittest import TestCase
+
+from samcli.lib.utils.rate_limiter import RateLimiter
+
+
+class TestRateLimiter(TestCase):
+    def test_try_acquire_succeeds_when_tokens_available(self):
+        limiter = RateLimiter(rate=10, burst=5)
+        self.assertTrue(limiter.try_acquire())
+
+    def test_try_acquire_fails_when_exhausted(self):
+        limiter = RateLimiter(rate=10, burst=2)
+        self.assertTrue(limiter.try_acquire())
+        self.assertTrue(limiter.try_acquire())
+        self.assertFalse(limiter.try_acquire())
+
+    def test_tokens_refill_over_time(self):
+        limiter = RateLimiter(rate=100, burst=1)
+        self.assertTrue(limiter.try_acquire())
+        self.assertFalse(limiter.try_acquire())
+        time.sleep(0.02)  # wait for refill
+        self.assertTrue(limiter.try_acquire())
+
+    def test_acquire_with_timeout_returns_false_on_expiry(self):
+        limiter = RateLimiter(rate=1, burst=1)
+        limiter.try_acquire()  # drain
+        self.assertFalse(limiter.acquire(timeout=0.01))
+
+    def test_acquire_blocks_until_available(self):
+        limiter = RateLimiter(rate=100, burst=1)
+        limiter.try_acquire()  # drain
+        start = time.monotonic()
+        self.assertTrue(limiter.acquire(timeout=1.0))
+        elapsed = time.monotonic() - start
+        self.assertLess(elapsed, 0.5)
+
+    def test_invalid_rate_raises(self):
+        with self.assertRaises(ValueError):
+            RateLimiter(rate=0, burst=1)
+
+    def test_invalid_burst_raises(self):
+        with self.assertRaises(ValueError):
+            RateLimiter(rate=1, burst=0)
+
+    def test_thread_safety(self):
+        limiter = RateLimiter(rate=1000, burst=100)
+        results = []
+
+        def worker():
+            acquired = sum(1 for _ in range(20) if limiter.try_acquire())
+            results.append(acquired)
+
+        threads = [threading.Thread(target=worker) for _ in range(10)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        # Total acquired should not exceed burst (100)
+        self.assertLessEqual(sum(results), 100)