feat: different rate limiting handling

src/__tests__/extensions/sessionrecording.js

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

src/__tests__/rate-limiter.ts

@@ -5,105 +5,136 @@ describe('Rate Limiter', () => {
 
     beforeEach(() => {
         jest.useFakeTimers()
-
-        rateLimiter = new RateLimiter()
     })
 
-    it('is not rate limited with no batch key', () => {
-        expect(rateLimiter.isRateLimited(undefined)).toBe(false)
-    })
+    it('not every request is checked for limiting', () => {
+        // you can probably do this with a jest spy but jest's mocks and spies are too confusing for me
+        let accessCounter = 0
+        const fakeResponse = { responseText: '{}' }
+        const objectWithCounting = new Proxy(fakeResponse, {
+            get: function () {
+                accessCounter++
+                return '{}'
+            },
+        })
 
-    it('is not rate limited if there is nothing in persistence', () => {
-        expect(rateLimiter.isRateLimited('the batch key')).toBe(false)
+        rateLimiter = new RateLimiter(0.1)
+
+        // call a loop 1000 times
+        for (let i = 0; i < 1000; i++) {
+            rateLimiter.checkForLimiting(objectWithCounting as unknown as XMLHttpRequest)
+        }
+        // it should be way less than 300 but the number varies and I don't want the test to flake
+        expect(accessCounter).toBeLessThan(200)
+        expect(accessCounter).toBeGreaterThanOrEqual(1)
     })
 
-    it('is not rate limited if there is mo matching batch key in persistence', () => {
-        rateLimiter.limits = { 'a different batch key': 1000 }
+    describe('with all requests checked for limiting', () => {
+        beforeEach(() => {
+            rateLimiter = new RateLimiter(1)
+        })
 
-        expect(rateLimiter.isRateLimited('the batch key')).toBe(false)
-    })
+        it('is not rate limited with no batch key', () => {
+            expect(rateLimiter.isRateLimited(undefined)).toBe(false)
+        })
 
-    it('is not rate limited if the retryAfter is in the past', () => {
-        rateLimiter.limits = { 'the batch key': new Date(Date.now() - 1000).getTime() }
+        it('is not rate limited if there is nothing in persistence', () => {
+            expect(rateLimiter.isRateLimited('the batch key')).toBe(false)
+        })
 
-        expect(rateLimiter.isRateLimited('the batch key')).toBe(false)
-    })
+        it('is not rate limited if there is mo matching batch key in persistence', () => {
+            rateLimiter.limits = { 'a different batch key': 1000 }
 
-    it('is rate limited if the retryAfter is in the future', () => {
-        rateLimiter.limits = { 'the batch key': new Date(Date.now() + 1000).getTime() }
+            expect(rateLimiter.isRateLimited('the batch key')).toBe(false)
+        })
 
-        expect(rateLimiter.isRateLimited('the batch key')).toBe(true)
-    })
+        it('is not rate limited if the retryAfter is in the past', () => {
+            rateLimiter.limits = { 'the batch key': new Date(Date.now() - 1000).getTime() }
 
-    it('sets the retryAfter on429Response', () => {
-        rateLimiter.on429Response({
-            status: 429,
-            getResponseHeader: (name: string) => (name === 'X-PostHog-Retry-After-Events' ? '150' : null),
-        } as unknown as XMLHttpRequest)
+            expect(rateLimiter.isRateLimited('the batch key')).toBe(false)
+        })
 
-        expect(rateLimiter.limits).toStrictEqual({ events: new Date().getTime() + 150_000 })
-    })
+        it('is rate limited if the retryAfter is in the future', () => {
+            rateLimiter.limits = { 'the batch key': new Date(Date.now() + 1000).getTime() }
+
+            expect(rateLimiter.isRateLimited('the batch key')).toBe(true)
+        })
 
-    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 events retryAfter on checkForLimiting', () => {
+            rateLimiter.checkForLimiting({
+                responseText: JSON.stringify({ quota_limited: ['events'] }),
+            } as unknown as XMLHttpRequest)
 
-        expect(rateLimiter.limits).toStrictEqual({ events: new Date().getTime() + 60_000 })
-    })
+            const expectedRetry = new Date().getTime() + 60_000
+            expect(rateLimiter.limits).toStrictEqual({ events: 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 the recordings retryAfter on checkForLimiting', () => {
+            rateLimiter.checkForLimiting({
+                responseText: JSON.stringify({ quota_limited: ['recordings'] }),
+            } as unknown as XMLHttpRequest)
 
-        expect(rateLimiter.limits).toStrictEqual({
-            'some-other-key': 4000,
-            events: new Date().getTime() + 150_000,
+            const expectedRetry = new Date().getTime() + 60_000
+            expect(rateLimiter.limits).toStrictEqual({ recordings: 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('sets multiple retryAfter on checkForLimiting', () => {
+            rateLimiter.checkForLimiting({
+                responseText: JSON.stringify({ quota_limited: ['recordings', 'events', 'mystery'] }),
+            } as unknown as XMLHttpRequest)
+
+            const expectedRetry = new Date().getTime() + 60_000
+            expect(rateLimiter.limits).toStrictEqual({
+                events: expectedRetry,
+                recordings: expectedRetry,
+                mystery: expectedRetry,
+            })
+        })
+
+        it('keeps existing batch keys checkForLimiting', () => {
+            rateLimiter.checkForLimiting({
+                responseText: JSON.stringify({ quota_limited: ['events'] }),
+            } as unknown as XMLHttpRequest)
+
+            rateLimiter.checkForLimiting({
+                responseText: JSON.stringify({ quota_limited: ['recordings'] }),
+            } as unknown as XMLHttpRequest)
 
-        expect(rateLimiter.limits).toStrictEqual({
-            'some-other-key': 4000,
-            events: new Date().getTime() + 150_000,
-            sessionRecording: new Date().getTime() + 200_000,
+            expect(rateLimiter.limits).toStrictEqual({
+                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'] }),
+            } as unknown as XMLHttpRequest)
 
-        expect(rateLimiter.limits).toStrictEqual({})
-    })
+            const firstRetryValue = rateLimiter.limits.events
+            jest.advanceTimersByTime(1000)
 
-    it('does not replace any limits if no Retry-After header is present', () => {
-        rateLimiter.limits = { 'some-other-key': 4000, events: 1000 }
+            rateLimiter.checkForLimiting({
+                responseText: JSON.stringify({ quota_limited: ['events'] }),
+            } as unknown as XMLHttpRequest)
+
+            expect(rateLimiter.limits).toStrictEqual({
+                events: firstRetryValue + 1000,
+            })
+        })
 
-        rateLimiter.on429Response({
-            status: 429,
-            getResponseHeader: () => null,
-        } as unknown as XMLHttpRequest)
+        it('does not set a limit if no limits are present', () => {
+            rateLimiter.checkForLimiting({
+                responseText: JSON.stringify({ quota_limited: [] }),
+            } as unknown as XMLHttpRequest)
 
-        expect(rateLimiter.limits).toStrictEqual({ 'some-other-key': 4000, events: 1000 })
+            expect(rateLimiter.limits).toStrictEqual({})
+
+            rateLimiter.checkForLimiting({
+                responseText: JSON.stringify({ status: 1 }),
+            } as unknown as XMLHttpRequest)
+
+            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',
@@ -33,7 +33,7 @@ describe('xhr', () => {
             enqueue: () => {},
         },
         onXHRError: given.onXHRError,
-        onRateLimited: given.on429Response,
+        onResponse: given.checkForLimiting,
     }))
     given('subject', () => () => {
         xhr(given.xhrParams)
@@ -58,16 +58,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'
@@ -331,7 +331,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

@@ -654,9 +654,6 @@ export class PostHog {
 
     _send_request(url: string, data: Record<string, any>, options: CaptureOptions, callback?: RequestCallback): void {
         if (this.rateLimiter.isRateLimited(options._batchKey)) {
-            if (this.get_config('debug')) {
-                console.warn('[PostHog SendRequest] is quota limited. Dropping request.')
-            }
             return
         }
 
@@ -702,7 +699,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,21 @@
-import { SESSION_RECORDING_BATCH_KEY } from './extensions/sessionrecording'
-
-/**
- * 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',
+import { logger } from './utils'
+import Config from './config'
+
+const oneMinuteInMilliseconds = 60 * 1000
+
+interface CaptureResponse {
+    quota_limited?: string[]
 }
 
 export class RateLimiter {
     limits: Record<string, number> = {}
+    private checkThreshold: number
+
+    constructor(checkThreshold = 0.1) {
+        this.checkThreshold = checkThreshold
+    }
 
-    isRateLimited(batchKey: string | undefined): boolean {
+    public isRateLimited(batchKey: string | undefined): boolean {
         const retryAfter = this.limits[batchKey || 'events'] || false
 
         if (retryAfter === false) {
@@ -26,26 +24,27 @@ export class RateLimiter {
         return new Date().getTime() < retryAfter
     }
 
-    on429Response(response: XMLHttpRequest): void {
-        if (response.status !== 429) {
+    // this needs to be an arrow function so that it can be passed as a callback
+    // and have the correct `this` context in order to read `checkThreshold`
+    public checkForLimiting = (xmlHttpRequest: XMLHttpRequest): void => {
+        if (Math.random() >= this.checkThreshold) {
+            // we don't need to check this on every request
             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
-            }
-        })
+        let response: CaptureResponse
+        try {
+            response = JSON.parse(xmlHttpRequest.responseText)
+            const quotaLimitedProducts = response.quota_limited || []
+            quotaLimitedProducts.forEach((batchKey) => {
+                if (Config.DEBUG) {
+                    console.warn(`[PostHog RateLimiter] ${batchKey || 'events'} is quota limited.`)
+                }
+                this.limits[batchKey] = new Date().getTime() + oneMinuteInMilliseconds
+            })
+        } catch (e) {
+            logger.error(e)
+            return
+        }
     }
 }