🐛 Fixes Retry-After header in global_rate_limit_route (ITISFoundation#3379)

Author: pcrespov

services/web/server/src/simcore_service_webserver/utils_rate_limiting.py

@@ -1,12 +1,19 @@
 import json
+from dataclasses import dataclass
 from datetime import datetime, timedelta
 from functools import wraps
+from math import ceil
+from typing import NamedTuple
 
-import attr
 from aiohttp.web_exceptions import HTTPTooManyRequests
 
 
-def global_rate_limit_route(number_of_requests: int, interval_seconds: int):
+class RateLimitSetup(NamedTuple):
+    number_of_requests: int
+    interval_seconds: float
+
+
+def global_rate_limit_route(number_of_requests: int, interval_seconds: float):
     """
     Limits the requests per given interval to this endpoint
     from all incoming sources.
@@ -19,14 +26,14 @@ def global_rate_limit_route(number_of_requests: int, interval_seconds: int):
     """
 
     # compute the amount of requests per
-    def decorating_function(decorated_function):
-        @attr.s(auto_attribs=True)
-        class Context:
+    def decorator(decorated_function):
+        @dataclass
+        class _Context:
             max_allowed: int  # maximum allowed requests per interval
             remaining: int  # remaining requests
-            rate_limit_reset: int  # utc timestamp
+            rate_limit_reset: float  # utc timestamp
 
-        context = Context(
+        context = _Context(
             max_allowed=number_of_requests,
             remaining=number_of_requests,
             rate_limit_reset=0,
@@ -35,24 +42,24 @@ class Context:
         @wraps(decorated_function)
         async def wrapper(*args, **kwargs):
             utc_now = datetime.utcnow()
-            current_utc_timestamp = datetime.timestamp(utc_now)
+            utc_now_timestamp = datetime.timestamp(utc_now)
 
             # reset counter & first time initialization
-            if current_utc_timestamp >= context.rate_limit_reset:
+            if utc_now_timestamp >= context.rate_limit_reset:
                 context.rate_limit_reset = datetime.timestamp(
                     utc_now + timedelta(seconds=interval_seconds)
                 )
                 context.remaining = context.max_allowed
 
-            if (
-                current_utc_timestamp <= context.rate_limit_reset
-                and context.remaining <= 0
-            ):
-                # show error and return from here
+            if utc_now_timestamp <= context.rate_limit_reset and context.remaining <= 0:
+                # SEE https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/429
+                retry_after_sec = int(
+                    ceil(context.rate_limit_reset - utc_now_timestamp)
+                )
                 raise HTTPTooManyRequests(
                     headers={
                         "Content-Type": "application/json",
-                        "Retry-After": str(int(context.rate_limit_reset)),
+                        "Retry-After": f"{retry_after_sec}",
                     },
                     text=json.dumps(
                         {
@@ -68,7 +75,7 @@ async def wrapper(*args, **kwargs):
             context.remaining -= 1
             return await decorated_function(*args, **kwargs)
 
-        wrapper.rate_limit_setup = (number_of_requests, interval_seconds)
+        wrapper.rate_limit_setup = RateLimitSetup(number_of_requests, interval_seconds)
         return wrapper
 
-    return decorating_function
+    return decorator

services/web/server/tests/unit/isolated/test_utils_rate_limiting.py

@@ -4,12 +4,16 @@
 
 import asyncio
 import time
+from typing import Callable
 
 import pytest
 from aiohttp import web
-from aiohttp.web_exceptions import HTTPTooManyRequests
+from aiohttp.test_utils import TestClient
+from aiohttp.web_exceptions import HTTPOk, HTTPTooManyRequests
+from pydantic import ValidationError, conint, parse_obj_as
 from simcore_service_webserver.utils_rate_limiting import global_rate_limit_route
 
+TOTAL_TEST_TIME = 1  # secs
 MAX_NUM_REQUESTS = 3
 MEASURE_INTERVAL = 0.5
 MAX_REQUEST_RATE = MAX_NUM_REQUESTS / MEASURE_INTERVAL
@@ -22,35 +26,43 @@ async def get_ok_handler(_request: web.Request):
     return web.json_response({"value": 1})
 
 
-@pytest.mark.parametrize(
-    "requests_per_second",
-    [0.5 * MAX_REQUEST_RATE, MAX_REQUEST_RATE, 2 * MAX_REQUEST_RATE],
-)
-async def test_global_rate_limit_route(requests_per_second, aiohttp_client):
-    #
+@pytest.fixture
+def client(event_loop, aiohttp_client: Callable) -> TestClient:
     app = web.Application()
     app.router.add_get("/", get_ok_handler)
 
-    client = await aiohttp_client(app)
-    # ---
+    return event_loop.run_until_complete(aiohttp_client(app))
 
+
+def test_rate_limit_route_decorator():
     # decorated function keeps setup
     assert get_ok_handler.rate_limit_setup == (MAX_NUM_REQUESTS, MEASURE_INTERVAL)
 
+
+@pytest.mark.parametrize(
+    "requests_per_second",
+    [0.5 * MAX_REQUEST_RATE, MAX_REQUEST_RATE, 2 * MAX_REQUEST_RATE],
+)
+async def test_global_rate_limit_route(requests_per_second: float, client: TestClient):
+    # WARNING: this test has some timings and might fail when using breakpoints
+
     # Creates desired stream of requests for 1 second
-    TOTAL_TEST_TIME = 1  # secs
     num_requests = int(requests_per_second * TOTAL_TEST_TIME)
     time_between_requests = 1.0 / requests_per_second
 
-    futures = []
+    tasks = []
     t0 = time.time()
-    while len(futures) < num_requests:
+    while len(tasks) < num_requests:
         t1 = time.time()
-        futures.append(asyncio.create_task(client.get("/")))
-        time.sleep(time_between_requests - (time.time() - t1))
+        tasks.append(asyncio.create_task(client.get("/")))
+        elapsed_on_creation = time.time() - t1  # ANE is really precise here ;-)
+
+        # NOTE: I am not sure why using asyncio.sleep here would make some tests fail the check "after"
+        # await asyncio.sleep(time_between_requests - create_gap)
+        time.sleep(time_between_requests - elapsed_on_creation)
 
     elapsed = time.time() - t0
-    count = len(futures)
+    count = len(tasks)
     print(
         count,
         "requests in",
@@ -63,20 +75,42 @@ async def test_global_rate_limit_route(requests_per_second, aiohttp_client):
     assert count == num_requests
     assert elapsed == pytest.approx(TOTAL_TEST_TIME, abs=0.1)
 
-    for i, fut in enumerate(futures):
-        while not fut.done():
-            await asyncio.sleep(0.1)
-        assert not fut.cancelled()
-        assert not fut.exception()
-        print("%2d" % i, fut.result().status)
-
-    expected_status = 200
+    msg = []
+    for i, task in enumerate(tasks):
+        while not task.done():
+            await asyncio.sleep(0.01)
+        assert not task.cancelled()
+        assert not task.exception()
+        msg.append(
+            (
+                "request # %2d" % i,
+                f"status={task.result().status}",
+                f"retry-after={task.result().headers.get('Retry-After')}",
+            )
+        )
+        print(*msg[-1])
+
+    expected_status = HTTPOk.status_code
 
     # first requests are OK
-    assert all(f.result().status == expected_status for f in futures[:MAX_NUM_REQUESTS])
+    assert all(
+        t.result().status == expected_status for t in tasks[:MAX_NUM_REQUESTS]
+    ), f" Failed with { msg[:MAX_NUM_REQUESTS]}"
 
     if requests_per_second >= MAX_REQUEST_RATE:
         expected_status = HTTPTooManyRequests.status_code
 
     # after ...
-    assert all(f.result().status == expected_status for f in futures[MAX_NUM_REQUESTS:])
+    assert all(
+        t.result().status == expected_status for t in tasks[MAX_NUM_REQUESTS:]
+    ), f" Failed with { msg[MAX_NUM_REQUESTS:]}"
+
+    # checks Retry-After header
+    failed = []
+    for t in tasks:
+        if retry_after := t.result().headers.get("Retry-After"):
+            try:
+                parse_obj_as(conint(ge=1), retry_after)
+            except ValidationError as err:
+                failed.append((retry_after, f"{err}"))
+    assert not failed