From 627a5ae2534d26152d597ef2fdc547c8cd45c67e Mon Sep 17 00:00:00 2001 From: Harold Sun Date: Fri, 20 Mar 2026 10:30:42 -0700 Subject: [PATCH 1/2] feat: add thread-safe token bucket rate limiter utility --- samcli/lib/utils/rate_limiter.py | 102 +++++++++++++++++++++++++++++++ 1 file changed, 102 insertions(+) create mode 100644 samcli/lib/utils/rate_limiter.py diff --git a/samcli/lib/utils/rate_limiter.py b/samcli/lib/utils/rate_limiter.py new file mode 100644 index 00000000000..52d96cc861e --- /dev/null +++ b/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) From 2011c80e5a93043c113f18ef0f4e0436d185f656 Mon Sep 17 00:00:00 2001 From: Harold Sun Date: Fri, 20 Mar 2026 10:30:46 -0700 Subject: [PATCH 2/2] test: add unit tests for rate limiter --- tests/unit/lib/utils/test_rate_limiter.py | 64 +++++++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 tests/unit/lib/utils/test_rate_limiter.py diff --git a/tests/unit/lib/utils/test_rate_limiter.py b/tests/unit/lib/utils/test_rate_limiter.py new file mode 100644 index 00000000000..2d5a0d5714b --- /dev/null +++ b/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)