track ratelimit only on creation (#5950)

thirdweb-dev · Jan 15, 2025 · c290997 · c290997
1 parent 81b146a
commit c290997
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 37 deletions.
diff --git a/.changeset/green-tips-beam.md b/.changeset/green-tips-beam.md
@@ -0,0 +1,5 @@
+---
+"@thirdweb-dev/service-utils": minor
+---
+
+track usage call once per ratelimit window
diff --git a/packages/service-utils/src/core/rateLimit/index.ts b/packages/service-utils/src/core/rateLimit/index.ts
@@ -71,22 +71,25 @@ export async function rateLimit(args: {
     limitPerSecond * sampleRate * RATE_LIMIT_WINDOW_SECONDS;
 
   if (requestCount > limitPerWindow) {
-    // Report rate limit hits.
-    if (project?.id) {
-      await updateRateLimitedAt(project.id, serviceConfig);
+    /**
+     * Report rate limit hits.
+     * Only track rate limit when its hit for the first time.
+     * Not waiting for tracking to complete as user doesn't need to wait.
+     */
+    if (requestCount === limitPerWindow + 1 && project?.id) {
+      updateRateLimitedAt(project.id, serviceConfig).catch(() => {
+        // no-op
+      });
     }
 
-    // Reject requests when they've exceeded 2x the rate limit.
-    if (requestCount > 2 * limitPerWindow) {
-      return {
-        rateLimited: true,
-        requestCount,
-        rateLimit: limitPerWindow,
-        status: 429,
-        errorMessage: `You've exceeded your ${serviceScope} rate limit at ${limitPerSecond} reqs/sec. To get higher rate limits, contact us at https://thirdweb.com/contact-us.`,
-        errorCode: "RATE_LIMIT_EXCEEDED",
-      };
-    }
+    return {
+      rateLimited: true,
+      requestCount,
+      rateLimit: limitPerWindow,
+      status: 429,
+      errorMessage: `You've exceeded your ${serviceScope} rate limit at ${limitPerSecond} reqs/sec. To get higher rate limits, contact us at https://thirdweb.com/contact-us.`,
+      errorCode: "RATE_LIMIT_EXCEEDED",
+    };
   }
 
   return {

diff --git a/packages/service-utils/src/core/rateLimit/rateLimit.test.ts b/packages/service-utils/src/core/rateLimit/rateLimit.test.ts
@@ -10,7 +10,7 @@ const mockRedis = {
 
 // Mocking the updateRateLimitedAt function
 vi.mock("../../../src/core/api", () => ({
-  updateRateLimitedAt: vi.fn(),
+  updateRateLimitedAt: vi.fn().mockResolvedValue({}),
 }));
 
 describe("rateLimit", () => {
@@ -59,27 +59,8 @@ describe("rateLimit", () => {
     expect(mockRedis.expire).not.toHaveBeenCalled();
   });
 
-  it("should report rate limit if exceeded but not block", async () => {
-    mockRedis.incr.mockResolvedValue(51); // Current count is 51 requests in 10 seconds.
-
-    const result = await rateLimit({
-      project: validProjectResponse,
-      limitPerSecond: 5,
-      serviceConfig: validServiceConfig,
-      redis: mockRedis,
-    });
-
-    expect(result).toEqual({
-      rateLimited: false,
-      requestCount: 51,
-      rateLimit: 50,
-    });
-    expect(updateRateLimitedAt).toHaveBeenCalled();
-    expect(mockRedis.expire).not.toHaveBeenCalled();
-  });
-
   it("should rate limit if exceeded hard limit", async () => {
-    mockRedis.incr.mockResolvedValue(101);
+    mockRedis.incr.mockResolvedValue(51);
 
     const result = await rateLimit({
       project: validProjectResponse,
@@ -90,7 +71,7 @@ describe("rateLimit", () => {
 
     expect(result).toEqual({
       rateLimited: true,
-      requestCount: 101,
+      requestCount: 51,
       rateLimit: 50,
       status: 429,
       errorMessage: `You've exceeded your storage rate limit at 5 reqs/sec. To get higher rate limits, contact us at https://thirdweb.com/contact-us.`,
@@ -131,9 +112,13 @@ describe("rateLimit", () => {
     });
 
     expect(result).toEqual({
-      rateLimited: false,
+      rateLimited: true,
       requestCount: 10,
       rateLimit: 5,
+      status: 429,
+      errorMessage:
+        "You've exceeded your storage rate limit at 5 reqs/sec. To get higher rate limits, contact us at https://thirdweb.com/contact-us.",
+      errorCode: "RATE_LIMIT_EXCEEDED",
     });
   });