feat: different rate limiting handling (#765)

* feat: different rate limiting handling

* fix

* fix

* Fix

* fix

src/__tests__/extensions/sessionrecording.js

@@ -281,7 +281,7 @@ describe('SessionRecording', () => {
                     method: 'POST',
                     endpoint: '/s/',
                     _noTruncate: true,
-                    _batchKey: 'sessionRecording',
+                    _batchKey: 'recordings',
                     _metrics: expect.anything(),
                 }
             )
@@ -318,7 +318,7 @@ describe('SessionRecording', () => {
                     transport: 'XHR',
                     endpoint: '/s/',
                     _noTruncate: true,
-                    _batchKey: 'sessionRecording',
+                    _batchKey: 'recordings',
                     _metrics: expect.anything(),
                 }
             )
@@ -368,7 +368,7 @@ describe('SessionRecording', () => {
                     transport: 'XHR',
                     endpoint: '/s/',
                     _noTruncate: true,
-                    _batchKey: 'sessionRecording',
+                    _batchKey: 'recordings',
                     _metrics: expect.anything(),
                 }
             )

src/__tests__/rate-limiter.ts

@@ -5,7 +5,6 @@ describe('Rate Limiter', () => {
 
     beforeEach(() => {
         jest.useFakeTimers()
-
         rateLimiter = new RateLimiter()
     })
 
@@ -35,75 +34,80 @@ describe('Rate Limiter', () => {
         expect(rateLimiter.isRateLimited('the batch key')).toBe(true)
     })
 
-    it('sets the retryAfter on429Response', () => {
-        rateLimiter.on429Response({
-            status: 429,
-            getResponseHeader: (name: string) => (name === 'X-PostHog-Retry-After-Events' ? '150' : null),
-        } as unknown as XMLHttpRequest)
+    it('sets the events retryAfter on checkForLimiting', () => {
+        rateLimiter.checkForLimiting({
+            responseText: JSON.stringify({ quota_limited: ['events'] }),
+        })
 
-        expect(rateLimiter.limits).toStrictEqual({ events: new Date().getTime() + 150_000 })
+        const expectedRetry = new Date().getTime() + 60_000
+        expect(rateLimiter.limits).toStrictEqual({ events: expectedRetry })
     })
 
-    it('sets the retryAfter to a default if the header is not a number in on429Response', () => {
-        rateLimiter.on429Response({
-            status: 429,
-            getResponseHeader: (name: string) => (name === 'X-PostHog-Retry-After-Events' ? 'tomato' : null),
-        } as unknown as XMLHttpRequest)
+    it('sets the recordings retryAfter on checkForLimiting', () => {
+        rateLimiter.checkForLimiting({
+            responseText: JSON.stringify({ quota_limited: ['recordings'] }),
+        })
 
-        expect(rateLimiter.limits).toStrictEqual({ events: new Date().getTime() + 60_000 })
+        const expectedRetry = new Date().getTime() + 60_000
+        expect(rateLimiter.limits).toStrictEqual({ recordings: expectedRetry })
     })
 
-    it('keeps existing batch keys on429Response', () => {
-        rateLimiter.limits = { 'some-other-key': 4000 }
-        rateLimiter.on429Response({
-            status: 429,
-            getResponseHeader: (name: string) => (name === 'X-PostHog-Retry-After-Events' ? '150' : null),
-        } as unknown as XMLHttpRequest)
+    it('sets multiple retryAfter on checkForLimiting', () => {
+        rateLimiter.checkForLimiting({
+            responseText: JSON.stringify({ quota_limited: ['recordings', 'events', 'mystery'] }),
+        })
 
+        const expectedRetry = new Date().getTime() + 60_000
         expect(rateLimiter.limits).toStrictEqual({
-            'some-other-key': 4000,
-            events: new Date().getTime() + 150_000,
+            events: expectedRetry,
+            recordings: expectedRetry,
+            mystery: expectedRetry,
         })
     })
 
-    it('replaces matching keys on429Response and ignores unexpected ones', () => {
-        rateLimiter.limits = { 'some-other-key': 4000, events: 1000 }
-        rateLimiter.on429Response({
-            status: 429,
-            getResponseHeader: (name: string) => {
-                if (name === 'X-PostHog-Retry-After-Events') {
-                    return '150'
-                } else if (name === 'X-PostHog-Retry-After-Recordings') {
-                    return '200'
-                }
-                return null
-            },
-        } as unknown as XMLHttpRequest)
+    it('keeps existing batch keys checkForLimiting', () => {
+        rateLimiter.checkForLimiting({
+            responseText: JSON.stringify({ quota_limited: ['events'] }),
+        })
+
+        rateLimiter.checkForLimiting({
+            responseText: JSON.stringify({ quota_limited: ['recordings'] }),
+        })
 
         expect(rateLimiter.limits).toStrictEqual({
-            'some-other-key': 4000,
-            events: new Date().getTime() + 150_000,
-            sessionRecording: new Date().getTime() + 200_000,
+            events: expect.any(Number),
+            recordings: expect.any(Number),
         })
     })
 
-    it('does not set a limit if no Retry-After header is present', () => {
-        rateLimiter.on429Response({
-            status: 429,
-            getResponseHeader: () => null,
-        } as unknown as XMLHttpRequest)
+    it('replaces matching keys', () => {
+        rateLimiter.checkForLimiting({
+            responseText: JSON.stringify({ quota_limited: ['events'] }),
+        })
 
-        expect(rateLimiter.limits).toStrictEqual({})
+        const firstRetryValue = rateLimiter.limits.events
+        jest.advanceTimersByTime(1000)
+
+        rateLimiter.checkForLimiting({
+            responseText: JSON.stringify({ quota_limited: ['events'] }),
+        })
+
+        expect(rateLimiter.limits).toStrictEqual({
+            events: firstRetryValue + 1000,
+        })
     })
 
-    it('does not replace any limits if no Retry-After header is present', () => {
-        rateLimiter.limits = { 'some-other-key': 4000, events: 1000 }
+    it('does not set a limit if no limits are present', () => {
+        rateLimiter.checkForLimiting({
+            responseText: JSON.stringify({ quota_limited: [] }),
+        })
 
-        rateLimiter.on429Response({
-            status: 429,
-            getResponseHeader: () => null,
-        } as unknown as XMLHttpRequest)
+        expect(rateLimiter.limits).toStrictEqual({})
+
+        rateLimiter.checkForLimiting({
+            responseText: JSON.stringify({ status: 1 }),
+        })
 
-        expect(rateLimiter.limits).toStrictEqual({ 'some-other-key': 4000, events: 1000 })
+        expect(rateLimiter.limits).toStrictEqual({})
     })
 })

src/__tests__/send-request.js

@@ -16,7 +16,7 @@ describe('xhr', () => {
         status: 502,
     }))
     given('onXHRError', jest.fn)
-    given('on429Response', jest.fn)
+    given('checkForLimiting', jest.fn)
     given('xhrOptions', () => ({}))
     given('xhrParams', () => ({
         url: 'https://any.posthog-instance.com',
@@ -29,7 +29,7 @@ describe('xhr', () => {
             enqueue: () => {},
         },
         onXHRError: given.onXHRError,
-        onRateLimited: given.on429Response,
+        onResponse: given.checkForLimiting,
     }))
     given('subject', () => () => {
         xhr(given.xhrParams)
@@ -54,16 +54,10 @@ describe('xhr', () => {
             expect(requestFromError).toHaveProperty('status', 502)
         })
 
-        it('calls the injected 429 handler for 429 responses', () => {
-            given.mockXHR.status = 429
+        it('calls the on response handler - regardless of status', () => {
+            given.mockXHR.status = Math.floor(Math.random() * 100)
             given.subject()
-            expect(given.on429Response).toHaveBeenCalledWith(given.mockXHR)
-        })
-
-        it('does not call the injected 429 handler for non-429 responses', () => {
-            given.mockXHR.status = 404
-            given.subject()
-            expect(given.on429Response).not.toHaveBeenCalled()
+            expect(given.checkForLimiting).toHaveBeenCalledWith(given.mockXHR)
         })
     })
 })

src/extensions/sessionrecording.ts

@@ -24,7 +24,7 @@ const BASE_ENDPOINT = '/s/'
 export const RECORDING_IDLE_ACTIVITY_TIMEOUT_MS = 5 * 60 * 1000 // 5 minutes
 export const RECORDING_MAX_EVENT_SIZE = 1024 * 1024 * 0.9 // ~1mb (with some wiggle room)
 export const RECORDING_BUFFER_TIMEOUT = 2000 // 2 seconds
-export const SESSION_RECORDING_BATCH_KEY = 'sessionRecording'
+export const SESSION_RECORDING_BATCH_KEY = 'recordings'
 
 // NOTE: Importing this type is problematic as we can't safely bundle it to a TS definition so, instead we redefine.
 // import type { record } from 'rrweb2/typings'
@@ -348,7 +348,7 @@ export class SessionRecording {
 
         this.mutationRateLimiter =
             this.mutationRateLimiter ??
-            new MutationRateLimiter(this.rrwebRecord!, {
+            new MutationRateLimiter(this.rrwebRecord, {
                 onBlockedNode: (id, node) => {
                     const message = `Too many mutations on node '${id}'. Rate limiting. This could be due to SVG animations or something similar`
                     logger.log(message, {

src/posthog-core.ts

@@ -641,9 +641,6 @@ export class PostHog {
             return
         }
         if (this.rateLimiter.isRateLimited(options._batchKey)) {
-            if (this.get_config('debug')) {
-                console.warn('[PostHog SendRequest] is quota limited. Dropping request.')
-            }
             return
         }
 
@@ -688,7 +685,7 @@ export class PostHog {
                     retriesPerformedSoFar: 0,
                     retryQueue: this._retryQueue,
                     onXHRError: this.get_config('on_xhr_error'),
-                    onRateLimited: this.rateLimiter.on429Response,
+                    onResponse: this.rateLimiter.checkForLimiting,
                 })
             } catch (e) {
                 console.error(e)

src/rate-limiter.ts

@@ -1,23 +1,16 @@
-import { SESSION_RECORDING_BATCH_KEY } from './extensions/sessionrecording'
+import { logger } from './utils'
+import Config from './config'
 
-/**
- * Really a 429 response should have a `Retry-After` header which is either a date string,
- * or the number of seconds to wait before retrying
- *
- * But we can rate limit endpoints differently, so send custom header per endpoint
- * The endpoints are configurable, so we tie the headers/retries to specific batch keys
- *
- * And only support a number of seconds to wait before retrying
- */
-const supportedRetryHeaders = {
-    'X-PostHog-Retry-After-Recordings': SESSION_RECORDING_BATCH_KEY,
-    'X-PostHog-Retry-After-Events': 'events',
+const oneMinuteInMilliseconds = 60 * 1000
+
+interface CaptureResponse {
+    quota_limited?: string[]
 }
 
 export class RateLimiter {
     limits: Record<string, number> = {}
 
-    isRateLimited(batchKey: string | undefined): boolean {
+    public isRateLimited(batchKey: string | undefined): boolean {
         const retryAfter = this.limits[batchKey || 'events'] || false
 
         if (retryAfter === false) {
@@ -26,26 +19,17 @@ export class RateLimiter {
         return new Date().getTime() < retryAfter
     }
 
-    on429Response(response: XMLHttpRequest): void {
-        if (response.status !== 429) {
+    public checkForLimiting = (xmlHttpRequest: XMLHttpRequest): void => {
+        try {
+            const response: CaptureResponse = JSON.parse(xmlHttpRequest.responseText)
+            const quotaLimitedProducts = response.quota_limited || []
+            quotaLimitedProducts.forEach((batchKey) => {
+                logger.log(`[PostHog RateLimiter] ${batchKey || 'events'} is quota limited.`)
+                this.limits[batchKey] = new Date().getTime() + oneMinuteInMilliseconds
+            })
+        } catch (e) {
+            logger.error(e)
             return
         }
-
-        Object.entries(supportedRetryHeaders).forEach(([header, batchKey]) => {
-            const responseHeader = response.getResponseHeader(header)
-            if (!responseHeader) {
-                return
-            }
-
-            let retryAfterSeconds = parseInt(responseHeader, 10)
-            if (isNaN(retryAfterSeconds)) {
-                retryAfterSeconds = 60
-            }
-
-            if (retryAfterSeconds) {
-                const retryAfterMillis = retryAfterSeconds * 1000
-                this.limits[batchKey] = new Date().getTime() + retryAfterMillis
-            }
-        })
     }
 }

src/retry-queue.ts

@@ -122,9 +122,6 @@ export class RetryQueue extends RequestQueueScaffold {
 
     _executeXhrRequest({ url, data, options, headers, callback, retriesPerformedSoFar }: QueuedRequestData): void {
         if (this.rateLimiter.isRateLimited(options._batchKey)) {
-            if (Config.DEBUG) {
-                console.warn('[PostHog RetryQueue] in quota limited mode. Dropping request.')
-            }
             return
         }
 
@@ -137,7 +134,7 @@ export class RetryQueue extends RequestQueueScaffold {
             callback,
             retryQueue: this,
             onXHRError: this.onXHRError,
-            onRateLimited: this.rateLimiter.on429Response,
+            onResponse: this.rateLimiter.checkForLimiting,
         })
     }
 

src/send-request.ts

@@ -67,7 +67,7 @@ export const xhr = ({
     retryQueue,
     onXHRError,
     timeout = 10000,
-    onRateLimited,
+    onResponse,
 }: XHRParams) => {
     const req = new XMLHttpRequest()
     req.open(options.method || 'GET', url, true)
@@ -87,8 +87,9 @@ export const xhr = ({
     // withCredentials cannot be modified until after calling .open on Android and Mobile Safari
     req.withCredentials = true
     req.onreadystatechange = () => {
+        // XMLHttpRequest.DONE == 4, except in safari 4
         if (req.readyState === 4) {
-            // XMLHttpRequest.DONE == 4, except in safari 4
+            onResponse?.(req)
             if (req.status === 200) {
                 if (callback) {
                     let response
@@ -105,8 +106,8 @@ export const xhr = ({
                     onXHRError(req)
                 }
 
-                // don't retry certain errors
-                if ([401, 403, 404, 500].indexOf(req.status) < 0) {
+                // don't retry errors between 400 and 500 inclusive
+                if (req.status < 400 || req.status > 500) {
                     retryQueue.enqueue({
                         url,
                         data,
@@ -117,13 +118,7 @@ export const xhr = ({
                     })
                 }
 
-                if (req.status === 429) {
-                    onRateLimited?.(req)
-                }
-
-                if (callback) {
-                    callback({ status: 0 })
-                }
+                callback?.({ status: 0 })
             }
         }
     }

src/types.ts

@@ -200,7 +200,7 @@ export interface XHRParams extends QueuedRequestData {
     retryQueue: RetryQueue
     onXHRError: (req: XMLHttpRequest) => void
     timeout?: number
-    onRateLimited?: (req: XMLHttpRequest) => void
+    onResponse?: (req: XMLHttpRequest) => void
 }
 
 export interface DecideResponse {