def __init__(self, api_key: str, http_client: HTTPClient = None) -> None: try: api_key = os.environ[api_key] except KeyError: pass self._key = api_key if http_client is None: self._client = HTTPClient() else: self._client = http_client self._rate_limiter = MultiRateLimiter( FixedWindowRateLimiter(600, 3000), FixedWindowRateLimiter(10, 50)) self._cached_data = {}
def _construct_limiters(self, limits: List[List[int]]): # Creates the necessary FixedWindowRateLimiters from the rates in the headers assert len(self._limiters) == 0 # Create the rate limiters for permits, window in limits: self._limiters.append( FixedWindowRateLimiter(window_seconds=window, window_permits=permits))
def _get_rate_limiter(self, platform: Platform, endpoint: str): try: limiter = self._rate_limiters[(platform, endpoint)] except KeyError: # TODO: Move to settings and don't force dev limits. # TODO This is a circular import and instead this functionality should somehow be moved into settings. from cassiopeia.configuration import settings limits = settings.rate_limits limiter = MultiRateLimiter(*([ FixedWindowRateLimiter(window, permits) for permits, window in limits ] + [ FixedWindowRateLimiter( *reversed(_METHOD_RATE_LIMITS[endpoint])) ])) self._rate_limiters[(platform, endpoint)] = limiter return limiter
def test_window_decorator_simple(): limiter = FixedWindowRateLimiter(SECONDS, PERMITS) @limiter.limit def call(): return True assert call()
def test_window_permit_count(): limiter = FixedWindowRateLimiter(SECONDS, MANY_PERMITS) assert limiter.permits_issued == 0 @limiter.limit def call(): pass for i in range(LARGE_VALUE_COUNT // 2): with limiter: pass assert limiter.permits_issued == i + 1 limiter.reset_permits_issued() assert limiter.permits_issued == 0 for i in range(LARGE_VALUE_COUNT // 2): call() assert limiter.permits_issued == i + 1
def test_window_acquire_across_windows(): from time import monotonic, sleep limiter = FixedWindowRateLimiter(SECONDS, 1) # This should take up two fixed windows for the first task, and the second shouldn't execute until the third (after 2 x SECONDS). with limiter: first = monotonic() sleep(SECONDS * 1.25) with limiter: second = monotonic() assert (SECONDS - EPSILON) * 3 >= second - first >= (SECONDS - EPSILON) * 2
def test_window_acquire_timing(): from time import monotonic limiter = FixedWindowRateLimiter(SECONDS, PERMITS) times = [] for _ in range(VALUE_COUNT): with limiter: times.append(monotonic()) start_indexes = [i for i in range(VALUE_COUNT) if i % PERMITS == 0] last = -SECONDS for index in start_indexes: assert times[index] - last >= SECONDS - EPSILON last = times[index]
def test_window_decorator_across_windows(): from time import monotonic, sleep limiter = FixedWindowRateLimiter(SECONDS, 1) @limiter.limit def call(): t = monotonic() sleep(SECONDS * 1.25) return t # This should take up two fixed windows for the first task, and the second shouldn't execute until the third (after 2 x SECONDS). first = call() second = call() assert (SECONDS - EPSILON) * 3 >= second - first >= (SECONDS - EPSILON) * 2
def test_window_decorator_timing(): from time import monotonic limiter = FixedWindowRateLimiter(SECONDS, PERMITS) @limiter.limit def call(): return monotonic() times = [] for _ in range(VALUE_COUNT): times.append(call()) start_indexes = [i for i in range(VALUE_COUNT) if i % PERMITS == 0] last = -SECONDS for index in start_indexes: assert times[index] - last >= SECONDS - EPSILON last = times[index]
def test_window_acquire_simple(): limiter = FixedWindowRateLimiter(SECONDS, PERMITS) x = False with limiter: x = True assert x