Fix: use proper bucket time for the rate limit

This commit is contained in:
Adrià Casajús 2024-02-20 11:13:06 +01:00
parent d0a6b8ed79
commit 363b851f61
No known key found for this signature in database
GPG Key ID: F0033226A5AFC9B9
1 changed files with 5 additions and 3 deletions

View File

@ -6,7 +6,7 @@ import redis.exceptions
import werkzeug.exceptions import werkzeug.exceptions
from limits.storage import RedisStorage from limits.storage import RedisStorage
from app.log import log from app.log import LOG
lock_redis: Optional[RedisStorage] = None lock_redis: Optional[RedisStorage] = None
@ -22,17 +22,19 @@ def check_bucket_limit(
bucket_seconds: int = 3600, bucket_seconds: int = 3600,
): ):
# Calculate current bucket time # Calculate current bucket time
bucket_id = int(datetime.utcnow().timestamp()) % bucket_seconds int_time = int(datetime.utcnow().timestamp())
bucket_id = int_time - (int_time % bucket_seconds)
bucket_lock_name = f"bl:{lock_name}:{bucket_id}" bucket_lock_name = f"bl:{lock_name}:{bucket_id}"
if not lock_redis: if not lock_redis:
return return
try: try:
value = lock_redis.incr(bucket_lock_name, bucket_seconds) value = lock_redis.incr(bucket_lock_name, bucket_seconds)
if value > max_hits: if value > max_hits:
LOG.i(f"Rate limit hit for {bucket_lock_name} -> {value}/{max_hits}")
newrelic.agent.record_custom_event( newrelic.agent.record_custom_event(
"BucketRateLimit", "BucketRateLimit",
{"lock_name": lock_name, "bucket_seconds": bucket_seconds}, {"lock_name": lock_name, "bucket_seconds": bucket_seconds},
) )
raise werkzeug.exceptions.TooManyRequests() raise werkzeug.exceptions.TooManyRequests()
except (redis.exceptions.RedisError, AttributeError): except (redis.exceptions.RedisError, AttributeError):
log.e("Cannot connect to redis") LOG.e("Cannot connect to redis")