Make multiple APIs share the flow control (#33)

README.md

@@ -54,7 +54,7 @@ app.add_middleware(
 )
 ```
 
-> :warning: **The pattern's order is important, rules are set on the first match**: Be careful here !
+:warning: **The pattern's order is important, rules are set on the first match**: Be careful here !
 
 Next, provide a custom authenticate function, or use one of the [existing auth methods](#built-in-auth-functions).
 
@@ -71,11 +71,8 @@ async def AUTH_FUNCTION(scope: Scope) -> Tuple[str, str]:
     # or use the function in the following document directly.
     return USER_UNIQUE_ID, GROUP_NAME
 
-rate_limit = RateLimitMiddleware(
-    ASGI_APP,
-    AUTH_FUNCTION,
-    ...
-)
+
+rate_limit = RateLimitMiddleware(ASGI_APP, AUTH_FUNCTION, ...)
 ```
 
 The `Rule` type takes a time unit (e.g. `"second"`) and/or a `"group"`, as a param. If the `"group"` param is not specified then the `"authenticate"` method needs to return the "default group".
@@ -125,6 +122,19 @@ Example for a "admin" group with higher limits.
     ...
 ```
 
+Sometimes you may want to specify that some APIs share the same flow control pool. In other words, flow control is performed on the entire set of APIs instead of a single specific API. Only the `zone` parameter needs to be used. **Note**: You can give different rules the same `zone` value, and all rules with the same `zone` value share the same flow control pool.
+
+```python
+    ...
+    config={
+        r"/user/\d+": [
+            Rule(minute=200, zone="user-api"),
+            Rule(second=100, zone="user-api", group="admin"),
+        ],
+    }
+    ...
+```
+
 ### Block time
 
 When the user's request frequency triggers the upper limit, all requests in the following period of time will be returned with a `429` status code.

ratelimit/core.py

@@ -74,7 +74,10 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
         if not [name for name in RULENAMES if getattr(rule, name) is not None]:
             return await self.app(scope, receive, send)
 
-        retry_after = await self.backend.retry_after(url_path, user, rule)
+        if rule.zone is None:
+            retry_after = await self.backend.retry_after(url_path, user, rule)
+        else:
+            retry_after = await self.backend.retry_after(rule.zone, user, rule)
         if retry_after == 0:
             return await self.app(scope, receive, send)
 

ratelimit/rule.py

@@ -14,6 +14,8 @@ class Rule:
 
     block_time: Optional[int] = None
 
+    zone: Optional[str] = None
+
     def ruleset(self, path: str, user: str) -> Dict[str, Tuple[int, int]]:
         """
         builds a dictionnary of keys, values where keys are the redis keys, and values

tests/test_core.py

@@ -127,3 +127,27 @@ async def test_custom_blocked():
         response = await client.get("/", headers={"user": "user", "group": "default"})
         assert response.status_code == 429
         assert response.content == b"custom 429 page"
+
+
+@pytest.mark.asyncio
+async def test_rule_zone():
+    rate_limit = RateLimitMiddleware(
+        hello_world,
+        auth_func,
+        RedisBackend(),
+        {
+            r"/message": [Rule(second=1, zone="commom")],
+            r"/\d+": [Rule(second=1, zone="commom")],
+        },
+    )
+    async with httpx.AsyncClient(
+        app=rate_limit, base_url="http://testserver"
+    ) as client:  # type: httpx.AsyncClient
+
+        response = await client.get("/10", headers={"user": "user", "group": "default"})
+        assert response.status_code == 200
+
+        response = await client.get(
+            "/message", headers={"user": "user", "group": "default"}
+        )
+        assert response.status_code == 429