mirror of https://gitlab.com/litecord/litecord.git
ratelimits.handler: five better retry_after and global flag
- run: add X-RateLimit-Global and Retry-After headers
This commit is contained in:
parent
33f893c0ff
commit
a96b9c5e7f
|
|
@ -10,8 +10,11 @@ async def _check_bucket(bucket):
|
|||
request.bucket = bucket
|
||||
|
||||
if retry_after:
|
||||
raise Ratelimited('You are being ratelimited.', {
|
||||
'retry_after': retry_after
|
||||
request.retry_after = retry_after
|
||||
|
||||
raise Ratelimited('You are being rate limited.', {
|
||||
'retry_after': int(retry_after * 1000),
|
||||
'global': request.bucket_global,
|
||||
})
|
||||
|
||||
|
||||
|
|
@ -22,6 +25,7 @@ async def _handle_global(ratelimit):
|
|||
except Unauthorized:
|
||||
user_id = request.remote_addr
|
||||
|
||||
request.bucket_global = True
|
||||
bucket = ratelimit.get_bucket(user_id)
|
||||
await _check_bucket(bucket)
|
||||
|
||||
|
|
@ -59,6 +63,12 @@ async def ratelimit_handler():
|
|||
# methods have different ratelimits
|
||||
rule_path = rule.endpoint
|
||||
|
||||
# some request ratelimit context.
|
||||
# TODO: maybe put those in a namedtuple or contextvar of sorts?
|
||||
request.bucket = None
|
||||
request.retry_after = None
|
||||
request.bucket_global = False
|
||||
|
||||
try:
|
||||
ratelimit = app.ratelimiter.get_ratelimit(rule_path)
|
||||
await _handle_specific(ratelimit)
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ class RatelimitManager:
|
|||
"""Manager for the bucket managers"""
|
||||
def __init__(self):
|
||||
self._ratelimiters = {}
|
||||
self.global_bucket = Ratelimit(50, 1)
|
||||
self.global_bucket = Ratelimit(1, 1)
|
||||
self._fill_rtl()
|
||||
|
||||
def _fill_rtl(self):
|
||||
|
|
|
|||
11
run.py
11
run.py
|
|
@ -122,9 +122,20 @@ async def app_set_ratelimit_headers(resp):
|
|||
"""Set the specific ratelimit headers."""
|
||||
try:
|
||||
bucket = request.bucket
|
||||
|
||||
if bucket is None:
|
||||
raise AttributeError()
|
||||
|
||||
resp.headers['X-RateLimit-Limit'] = str(bucket.requests)
|
||||
resp.headers['X-RateLimit-Remaining'] = str(bucket._tokens)
|
||||
resp.headers['X-RateLimit-Reset'] = str(bucket._window + bucket.second)
|
||||
|
||||
resp.headers['X-RateLimit-Global'] = str(request.bucket_global).lower()
|
||||
|
||||
# only add Retry-After if we actually hit a ratelimit
|
||||
retry_after = request.retry_after
|
||||
if request.retry_after:
|
||||
resp.headers['Retry-After'] = str(retry_after)
|
||||
except AttributeError:
|
||||
pass
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue