feat(ratelimits): Add metric buckets rate limit (#11506)

closes #11336
getsentry · Apr 10, 2024 · ab40116 · ab40116
1 parent 9d861fe
commit ab40116
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 3 deletions.
diff --git a/packages/utils/src/ratelimit.ts b/packages/utils/src/ratelimit.ts
@@ -67,7 +67,7 @@ export function updateRateLimits(
      * rate limit headers are of the form
      *     <header>,<header>,..
      * where each <header> is of the form
-     *     <retry_after>: <categories>: <scope>: <reason_code>
+     *     <retry_after>: <categories>: <scope>: <reason_code>: <namespaces>
      * where
      *     <retry_after> is a delay in seconds
      *     <categories> is the event type(s) (error, transaction, etc) being rate limited and is of the form
@@ -78,14 +78,21 @@ export function updateRateLimits(
      *         Only present if rate limit applies to the metric_bucket data category.
      */
     for (const limit of rateLimitHeader.trim().split(',')) {
-      const [retryAfter, categories] = limit.split(':', 2);
+      const [retryAfter, categories, , , namespaces] = limit.split(':', 5);
       const headerDelay = parseInt(retryAfter, 10);
       const delay = (!isNaN(headerDelay) ? headerDelay : 60) * 1000; // 60sec default
       if (!categories) {
         updatedRateLimits.all = now + delay;
       } else {
         for (const category of categories.split(';')) {
-          updatedRateLimits[category] = now + delay;
+          if (category === 'metric_bucket') {
+            // namespaces will be present when category === 'metric_bucket'
+            if (!namespaces || namespaces.split(';').includes('custom')) {
+              updatedRateLimits[category] = now + delay;
+            }
+          } else {
+            updatedRateLimits[category] = now + delay;
+          }
         }
       }
     }

diff --git a/packages/utils/test/ratelimit.test.ts b/packages/utils/test/ratelimit.test.ts
@@ -197,3 +197,43 @@ describe('updateRateLimits()', () => {
     expect(updatedRateLimits.all).toEqual(60_000);
   });
 });
+
+describe('data category "metric_bucket"', () => {
+  test('should add limit for `metric_bucket` category when namespaces contain "custom"', () => {
+    const rateLimits: RateLimits = {};
+    const headers = {
+      'retry-after': null,
+      'x-sentry-rate-limits': '42:metric_bucket:::custom',
+    };
+    const updatedRateLimits = updateRateLimits(rateLimits, { headers }, 0);
+    expect(updatedRateLimits.metric_bucket).toEqual(42 * 1000);
+  });
+
+  test('should not add limit for `metric_bucket` category when namespaces do not contain "custom"', () => {
+    const rateLimits: RateLimits = {};
+    const headers = {
+      'retry-after': null,
+      'x-sentry-rate-limits': '42:metric_bucket:::namespace1;namespace2',
+    };
+    const updatedRateLimits = updateRateLimits(rateLimits, { headers }, 0);
+    expect(updatedRateLimits.metric_bucket).toBeUndefined();
+  });
+
+  test('should add limit for `metric_bucket` category when namespaces are empty', () => {
+    const rateLimits: RateLimits = {};
+
+    const headers1 = {
+      'retry-after': null,
+      'x-sentry-rate-limits': '42:metric_bucket', // without semicolon at the end
+    };
+    const updatedRateLimits1 = updateRateLimits(rateLimits, { headers: headers1 }, 0);
+    expect(updatedRateLimits1.metric_bucket).toEqual(42 * 1000);
+
+    const headers2 = {
+      'retry-after': null,
+      'x-sentry-rate-limits': '42:metric_bucket:organization:quota_exceeded:', // with semicolon at the end
+    };
+    const updatedRateLimits2 = updateRateLimits(rateLimits, { headers: headers2 }, 0);
+    expect(updatedRateLimits2.metric_bucket).toEqual(42 * 1000);
+  });
+});