From 3ab2c2fbcf877f4ac1aa9957f0d3ffbfc3f470c1 Mon Sep 17 00:00:00 2001
From: Ankush Menat <ankush@frappe.io>
Date: Thu, 26 Dec 2024 16:27:46 +0530
Subject: [PATCH] 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.
---
 frappe/rate_limiter.py            | 30 +++++++++++++++++++++---------
 frappe/tests/test_rate_limiter.py |  2 --
 2 files changed, 21 insertions(+), 11 deletions(-)

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)