perf: speedup rate limiter by ~1.2x (#28920)

* perf: reuse current time

now_datetime is site-tz-aware, we don't need it here.

* perf: dont need redis transactions

* perf: use `time.time()` instead of datetime

Using `datetime.timestamp()` is a round-about way to use `time.time()`
with extra cost of dealing with datetime and timezones.

* perf: define slots for rate_limiter

* fix!: Remove used rate limit header

This just shares how much was consumed in current request, people can
just time requests to get an approximation for this, not sure why is this
useful.
This commit is contained in:
Ankush Menat 2024-12-26 16:27:46 +05:30 committed by GitHub
parent 304dae9136
commit 3ab2c2fbcf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 21 additions and 11 deletions

View file

@ -1,7 +1,7 @@
# Copyright (c) 2020, Frappe Technologies Pvt. Ltd. and Contributors
# License: MIT. See LICENSE
import datetime
import time
from collections.abc import Callable
from functools import wraps
@ -30,14 +30,28 @@ def respond():
class RateLimiter:
__slots__ = (
"counter",
"duration",
"end",
"key",
"limit",
"rejected",
"remaining",
"reset",
"spent",
"start",
"window",
"window_number",
)
def __init__(self, limit, window):
self.limit = int(limit * 1000000)
self.window = window
self.start = datetime.datetime.now(datetime.timezone.utc)
timestamp = int(frappe.utils.now_datetime().timestamp())
self.start = time.time()
self.window_number, self.spent = divmod(timestamp, self.window)
self.window_number, self.spent = divmod(int(self.start), self.window)
self.key = frappe.cache.make_key(f"rate-limit-counter-{self.window_number}")
self.counter = cint(frappe.cache.get(self.key))
self.remaining = max(self.limit - self.counter, 0)
@ -57,7 +71,7 @@ class RateLimiter:
def update(self):
self.record_request_end()
pipeline = frappe.cache.pipeline()
pipeline = frappe.cache.pipeline(transaction=False)
pipeline.incrby(self.key, self.duration)
pipeline.expire(self.key, self.window)
pipeline.execute()
@ -71,16 +85,14 @@ class RateLimiter:
}
if self.rejected:
headers["Retry-After"] = self.reset
else:
headers["X-RateLimit-Used"] = self.duration
return headers
def record_request_end(self):
if self.end is not None:
return
self.end = datetime.datetime.now(datetime.timezone.utc)
self.duration = int((self.end - self.start).total_seconds() * 1000000)
self.end = time.time()
self.duration = int((self.end - self.start) * 1000000)
def respond(self):
if self.rejected:

View file

@ -45,7 +45,6 @@ class TestRateLimiter(IntegrationTestCase):
headers = frappe.local.rate_limiter.headers()
self.assertIn("Retry-After", headers)
self.assertNotIn("X-RateLimit-Used", headers)
self.assertIn("X-RateLimit-Reset", headers)
self.assertIn("X-RateLimit-Limit", headers)
self.assertIn("X-RateLimit-Remaining", headers)
@ -75,7 +74,6 @@ class TestRateLimiter(IntegrationTestCase):
self.assertNotIn("Retry-After", headers)
self.assertIn("X-RateLimit-Reset", headers)
self.assertTrue(int(headers["X-RateLimit-Reset"] < 86400))
self.assertEqual(int(headers["X-RateLimit-Used"]), frappe.local.rate_limiter.duration)
self.assertEqual(int(headers["X-RateLimit-Limit"]), 10000)
self.assertEqual(int(headers["X-RateLimit-Remaining"]), 10000)