diff --git a/frappe/rate_limiter.py b/frappe/rate_limiter.py index 8cdc4adb1f..a09cfe96ed 100644 --- a/frappe/rate_limiter.py +++ b/frappe/rate_limiter.py @@ -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: diff --git a/frappe/tests/test_rate_limiter.py b/frappe/tests/test_rate_limiter.py index 97d82714ea..aa1393ab08 100644 --- a/frappe/tests/test_rate_limiter.py +++ b/frappe/tests/test_rate_limiter.py @@ -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)