From eb4b9bbf7b4875440d335ca691749190b7ddb22f Mon Sep 17 00:00:00 2001 From: Paul D'Ambra Date: Thu, 27 Jul 2023 15:46:42 +0100 Subject: [PATCH 1/5] feat: different rate limiting handling --- src/__tests__/rate-limiter.ts | 180 +++++++++++++++++------------ src/extensions/sessionrecording.ts | 2 +- src/posthog-core.ts | 2 +- src/rate-limiter.ts | 51 ++++---- src/retry-queue.ts | 2 +- src/send-request.ts | 13 +-- 6 files changed, 133 insertions(+), 117 deletions(-) diff --git a/src/__tests__/rate-limiter.ts b/src/__tests__/rate-limiter.ts index 067424e5d..0ffc4fc97 100644 --- a/src/__tests__/rate-limiter.ts +++ b/src/__tests__/rate-limiter.ts @@ -5,105 +5,135 @@ 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) + } + expect(accessCounter).toBeLessThan(110) + 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({}) + }) }) }) diff --git a/src/extensions/sessionrecording.ts b/src/extensions/sessionrecording.ts index 44b1fef28..6737f0e65 100644 --- a/src/extensions/sessionrecording.ts +++ b/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' diff --git a/src/posthog-core.ts b/src/posthog-core.ts index d790dd7e4..51a00f017 100644 --- a/src/posthog-core.ts +++ b/src/posthog-core.ts @@ -702,7 +702,7 @@ export class PostHog { retriesPerformedSoFar: 0, retryQueue: this._retryQueue, onXHRError: this.get_config('on_xhr_error'), - onRateLimited: this.rateLimiter.on429Response, + onRateLimited: this.rateLimiter.checkForLimiting, }) } catch (e) { console.error(e) diff --git a/src/rate-limiter.ts b/src/rate-limiter.ts index 0bdef0f90..8ebf1eca0 100644 --- a/src/rate-limiter.ts +++ b/src/rate-limiter.ts @@ -1,22 +1,17 @@ import { SESSION_RECORDING_BATCH_KEY } from './extensions/sessionrecording' +import { logger } from './utils' -/** - * 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 = {} + constructor(private checkThreshold = 0.1) {} + isRateLimited(batchKey: string | undefined): boolean { const retryAfter = this.limits[batchKey || 'events'] || false @@ -26,26 +21,22 @@ export class RateLimiter { return new Date().getTime() < retryAfter } - on429Response(response: XMLHttpRequest): void { - if (response.status !== 429) { + 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) => { + this.limits[batchKey] = new Date().getTime() + oneMinuteInMilliseconds + }) + } catch (e) { + logger.error(e) + return + } } } diff --git a/src/retry-queue.ts b/src/retry-queue.ts index 56e08a206..e4dc1301c 100644 --- a/src/retry-queue.ts +++ b/src/retry-queue.ts @@ -145,7 +145,7 @@ export class RetryQueue extends RequestQueueScaffold { captureMetrics: this.captureMetrics, retryQueue: this, onXHRError: this.onXHRError, - onRateLimited: this.rateLimiter.on429Response, + onRateLimited: this.rateLimiter.checkForLimiting, }) } diff --git a/src/send-request.ts b/src/send-request.ts index 53273ad9f..f83ff220e 100644 --- a/src/send-request.ts +++ b/src/send-request.ts @@ -97,6 +97,7 @@ export const xhr = ({ captureMetrics.decr('_send_request_inflight') // XMLHttpRequest.DONE == 4, except in safari 4 + onRateLimited?.(req) if (req.status === 200) { if (callback) { let response @@ -113,8 +114,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, @@ -125,13 +126,7 @@ export const xhr = ({ }) } - if (req.status === 429) { - onRateLimited?.(req) - } - - if (callback) { - callback({ status: 0 }) - } + callback?.({ status: 0 }) } } } From 94283f4946147df0e7441e15135205fb08bdcdbb Mon Sep 17 00:00:00 2001 From: Paul D'Ambra Date: Thu, 27 Jul 2023 16:30:12 +0100 Subject: [PATCH 2/5] fix --- src/__tests__/extensions/sessionrecording.js | 6 +++--- src/__tests__/rate-limiter.ts | 3 ++- src/__tests__/send-request.js | 16 +++++----------- src/posthog-core.ts | 2 +- src/retry-queue.ts | 2 +- src/send-request.ts | 6 +++--- src/types.ts | 2 +- 7 files changed, 16 insertions(+), 21 deletions(-) diff --git a/src/__tests__/extensions/sessionrecording.js b/src/__tests__/extensions/sessionrecording.js index 77a231879..5108dbb51 100644 --- a/src/__tests__/extensions/sessionrecording.js +++ b/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(), } ) diff --git a/src/__tests__/rate-limiter.ts b/src/__tests__/rate-limiter.ts index 0ffc4fc97..859ae7894 100644 --- a/src/__tests__/rate-limiter.ts +++ b/src/__tests__/rate-limiter.ts @@ -24,7 +24,8 @@ describe('Rate Limiter', () => { for (let i = 0; i < 1000; i++) { rateLimiter.checkForLimiting(objectWithCounting as unknown as XMLHttpRequest) } - expect(accessCounter).toBeLessThan(110) + // 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) }) diff --git a/src/__tests__/send-request.js b/src/__tests__/send-request.js index 9c0e76e58..af4a3fb93 100644 --- a/src/__tests__/send-request.js +++ b/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) }) }) }) diff --git a/src/posthog-core.ts b/src/posthog-core.ts index 51a00f017..c00a9d579 100644 --- a/src/posthog-core.ts +++ b/src/posthog-core.ts @@ -702,7 +702,7 @@ export class PostHog { retriesPerformedSoFar: 0, retryQueue: this._retryQueue, onXHRError: this.get_config('on_xhr_error'), - onRateLimited: this.rateLimiter.checkForLimiting, + onResponse: this.rateLimiter.checkForLimiting, }) } catch (e) { console.error(e) diff --git a/src/retry-queue.ts b/src/retry-queue.ts index e4dc1301c..969c635fe 100644 --- a/src/retry-queue.ts +++ b/src/retry-queue.ts @@ -145,7 +145,7 @@ export class RetryQueue extends RequestQueueScaffold { captureMetrics: this.captureMetrics, retryQueue: this, onXHRError: this.onXHRError, - onRateLimited: this.rateLimiter.checkForLimiting, + onResponse: this.rateLimiter.checkForLimiting, }) } diff --git a/src/send-request.ts b/src/send-request.ts index f83ff220e..65ac5a910 100644 --- a/src/send-request.ts +++ b/src/send-request.ts @@ -68,7 +68,7 @@ export const xhr = ({ retryQueue, onXHRError, timeout = 10000, - onRateLimited, + onResponse, }: XHRParams) => { const req = new XMLHttpRequest() req.open(options.method || 'GET', url, true) @@ -91,13 +91,13 @@ 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) { captureMetrics.incr(`xhr-response`) captureMetrics.incr(`xhr-response-${req.status}`) captureMetrics.decr('_send_request_inflight') - // XMLHttpRequest.DONE == 4, except in safari 4 - onRateLimited?.(req) + onResponse?.(req) if (req.status === 200) { if (callback) { let response diff --git a/src/types.ts b/src/types.ts index 8b6e08c26..a7ea2411b 100644 --- a/src/types.ts +++ b/src/types.ts @@ -202,7 +202,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 { From 20810681d9dcfe80c0aa40a9a81991dddb3fd8e9 Mon Sep 17 00:00:00 2001 From: Paul D'Ambra Date: Thu, 27 Jul 2023 21:56:13 +0100 Subject: [PATCH 3/5] fix --- src/extensions/sessionrecording.ts | 2 +- src/posthog-core.ts | 3 --- src/rate-limiter.ts | 16 ++++++++++++---- src/retry-queue.ts | 3 --- 4 files changed, 13 insertions(+), 11 deletions(-) diff --git a/src/extensions/sessionrecording.ts b/src/extensions/sessionrecording.ts index 6737f0e65..adc5ee165 100644 --- a/src/extensions/sessionrecording.ts +++ b/src/extensions/sessionrecording.ts @@ -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, { diff --git a/src/posthog-core.ts b/src/posthog-core.ts index c00a9d579..124e3daf7 100644 --- a/src/posthog-core.ts +++ b/src/posthog-core.ts @@ -654,9 +654,6 @@ export class PostHog { _send_request(url: string, data: Record, 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 } diff --git a/src/rate-limiter.ts b/src/rate-limiter.ts index 8ebf1eca0..0b4c40660 100644 --- a/src/rate-limiter.ts +++ b/src/rate-limiter.ts @@ -1,5 +1,5 @@ -import { SESSION_RECORDING_BATCH_KEY } from './extensions/sessionrecording' import { logger } from './utils' +import Config from './config' const oneMinuteInMilliseconds = 60 * 1000 @@ -9,10 +9,13 @@ interface CaptureResponse { export class RateLimiter { limits: Record = {} + private checkThreshold: number - constructor(private checkThreshold = 0.1) {} + 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) { @@ -21,7 +24,9 @@ export class RateLimiter { return new Date().getTime() < retryAfter } - checkForLimiting(xmlHttpRequest: XMLHttpRequest): void { + // 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 @@ -32,6 +37,9 @@ export class RateLimiter { 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) { diff --git a/src/retry-queue.ts b/src/retry-queue.ts index 969c635fe..d7de73f87 100644 --- a/src/retry-queue.ts +++ b/src/retry-queue.ts @@ -129,9 +129,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 } From de8428e1606da76d436d282927606a3ec84593b3 Mon Sep 17 00:00:00 2001 From: Paul D'Ambra Date: Fri, 28 Jul 2023 11:44:42 +0100 Subject: [PATCH 4/5] Fix --- src/__tests__/rate-limiter.ts | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/__tests__/rate-limiter.ts b/src/__tests__/rate-limiter.ts index 859ae7894..7fce0151a 100644 --- a/src/__tests__/rate-limiter.ts +++ b/src/__tests__/rate-limiter.ts @@ -20,11 +20,10 @@ describe('Rate Limiter', () => { 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 + // it should be way less than 200 but the number varies and I don't want the test to flake expect(accessCounter).toBeLessThan(200) expect(accessCounter).toBeGreaterThanOrEqual(1) }) From 5b3c900454d5f47904daff31eff8b3b65fefda95 Mon Sep 17 00:00:00 2001 From: Paul D'Ambra Date: Tue, 1 Aug 2023 19:16:48 +0100 Subject: [PATCH 5/5] fix --- src/__tests__/rate-limiter.ts | 166 ++++++++++++++-------------------- src/rate-limiter.ts | 19 +--- 2 files changed, 72 insertions(+), 113 deletions(-) diff --git a/src/__tests__/rate-limiter.ts b/src/__tests__/rate-limiter.ts index 7fce0151a..2f7d1d9e3 100644 --- a/src/__tests__/rate-limiter.ts +++ b/src/__tests__/rate-limiter.ts @@ -5,135 +5,109 @@ describe('Rate Limiter', () => { beforeEach(() => { jest.useFakeTimers() + rateLimiter = new RateLimiter() }) - 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 '{}' - }, - }) - - rateLimiter = new RateLimiter(0.1) + it('is not rate limited with no batch key', () => { + expect(rateLimiter.isRateLimited(undefined)).toBe(false) + }) - for (let i = 0; i < 1000; i++) { - rateLimiter.checkForLimiting(objectWithCounting as unknown as XMLHttpRequest) - } - // it should be way less than 200 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 nothing in persistence', () => { + expect(rateLimiter.isRateLimited('the batch key')).toBe(false) }) - describe('with all requests checked for limiting', () => { - beforeEach(() => { - rateLimiter = new RateLimiter(1) - }) + it('is not rate limited if there is mo matching batch key in persistence', () => { + rateLimiter.limits = { 'a different batch key': 1000 } - it('is not rate limited with no batch key', () => { - expect(rateLimiter.isRateLimited(undefined)).toBe(false) - }) + expect(rateLimiter.isRateLimited('the batch key')).toBe(false) + }) - it('is not rate limited if there is nothing in persistence', () => { - expect(rateLimiter.isRateLimited('the batch key')).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 mo matching batch key in persistence', () => { - rateLimiter.limits = { 'a different batch key': 1000 } + expect(rateLimiter.isRateLimited('the batch key')).toBe(false) + }) - expect(rateLimiter.isRateLimited('the batch key')).toBe(false) - }) + it('is rate limited if the retryAfter is in the future', () => { + rateLimiter.limits = { 'the batch key': new Date(Date.now() + 1000).getTime() } - it('is not rate limited if the retryAfter is in the past', () => { - rateLimiter.limits = { 'the batch key': new Date(Date.now() - 1000).getTime() } + expect(rateLimiter.isRateLimited('the batch key')).toBe(true) + }) - expect(rateLimiter.isRateLimited('the batch key')).toBe(false) + it('sets the events retryAfter on checkForLimiting', () => { + rateLimiter.checkForLimiting({ + responseText: JSON.stringify({ quota_limited: ['events'] }), }) - it('is rate limited if the retryAfter is in the future', () => { - rateLimiter.limits = { 'the batch key': new Date(Date.now() + 1000).getTime() } + const expectedRetry = new Date().getTime() + 60_000 + expect(rateLimiter.limits).toStrictEqual({ events: expectedRetry }) + }) - expect(rateLimiter.isRateLimited('the batch key')).toBe(true) + it('sets the recordings retryAfter on checkForLimiting', () => { + rateLimiter.checkForLimiting({ + responseText: JSON.stringify({ quota_limited: ['recordings'] }), }) - it('sets the events retryAfter on checkForLimiting', () => { - rateLimiter.checkForLimiting({ - responseText: JSON.stringify({ quota_limited: ['events'] }), - } as unknown as XMLHttpRequest) + const expectedRetry = new Date().getTime() + 60_000 + expect(rateLimiter.limits).toStrictEqual({ recordings: expectedRetry }) + }) - const expectedRetry = new Date().getTime() + 60_000 - expect(rateLimiter.limits).toStrictEqual({ events: expectedRetry }) + it('sets multiple retryAfter on checkForLimiting', () => { + rateLimiter.checkForLimiting({ + responseText: JSON.stringify({ quota_limited: ['recordings', 'events', 'mystery'] }), }) - it('sets the recordings retryAfter on checkForLimiting', () => { - rateLimiter.checkForLimiting({ - responseText: JSON.stringify({ quota_limited: ['recordings'] }), - } as unknown as XMLHttpRequest) - - const expectedRetry = new Date().getTime() + 60_000 - expect(rateLimiter.limits).toStrictEqual({ recordings: expectedRetry }) + const expectedRetry = new Date().getTime() + 60_000 + expect(rateLimiter.limits).toStrictEqual({ + events: expectedRetry, + recordings: expectedRetry, + mystery: expectedRetry, }) + }) - 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'] }), }) - 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({ - events: expect.any(Number), - recordings: expect.any(Number), - }) + rateLimiter.checkForLimiting({ + responseText: JSON.stringify({ quota_limited: ['recordings'] }), }) - it('replaces matching keys', () => { - rateLimiter.checkForLimiting({ - responseText: JSON.stringify({ quota_limited: ['events'] }), - } as unknown as XMLHttpRequest) + expect(rateLimiter.limits).toStrictEqual({ + events: expect.any(Number), + recordings: expect.any(Number), + }) + }) - const firstRetryValue = rateLimiter.limits.events - jest.advanceTimersByTime(1000) + it('replaces matching keys', () => { + rateLimiter.checkForLimiting({ + responseText: JSON.stringify({ quota_limited: ['events'] }), + }) - rateLimiter.checkForLimiting({ - responseText: JSON.stringify({ quota_limited: ['events'] }), - } as unknown as XMLHttpRequest) + const firstRetryValue = rateLimiter.limits.events + jest.advanceTimersByTime(1000) - expect(rateLimiter.limits).toStrictEqual({ - events: firstRetryValue + 1000, - }) + rateLimiter.checkForLimiting({ + responseText: JSON.stringify({ quota_limited: ['events'] }), }) - 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({ + events: firstRetryValue + 1000, + }) + }) - expect(rateLimiter.limits).toStrictEqual({}) + it('does not set a limit if no limits are present', () => { + rateLimiter.checkForLimiting({ + responseText: JSON.stringify({ quota_limited: [] }), + }) - rateLimiter.checkForLimiting({ - responseText: JSON.stringify({ status: 1 }), - } as unknown as XMLHttpRequest) + expect(rateLimiter.limits).toStrictEqual({}) - expect(rateLimiter.limits).toStrictEqual({}) + rateLimiter.checkForLimiting({ + responseText: JSON.stringify({ status: 1 }), }) + + expect(rateLimiter.limits).toStrictEqual({}) }) }) diff --git a/src/rate-limiter.ts b/src/rate-limiter.ts index 0b4c40660..7b07f0802 100644 --- a/src/rate-limiter.ts +++ b/src/rate-limiter.ts @@ -9,11 +9,6 @@ interface CaptureResponse { export class RateLimiter { limits: Record = {} - private checkThreshold: number - - constructor(checkThreshold = 0.1) { - this.checkThreshold = checkThreshold - } public isRateLimited(batchKey: string | undefined): boolean { const retryAfter = this.limits[batchKey || 'events'] || false @@ -24,22 +19,12 @@ export class RateLimiter { return new Date().getTime() < retryAfter } - // 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 - } - - let response: CaptureResponse try { - response = JSON.parse(xmlHttpRequest.responseText) + const response: CaptureResponse = JSON.parse(xmlHttpRequest.responseText) const quotaLimitedProducts = response.quota_limited || [] quotaLimitedProducts.forEach((batchKey) => { - if (Config.DEBUG) { - console.warn(`[PostHog RateLimiter] ${batchKey || 'events'} is quota limited.`) - } + logger.log(`[PostHog RateLimiter] ${batchKey || 'events'} is quota limited.`) this.limits[batchKey] = new Date().getTime() + oneMinuteInMilliseconds }) } catch (e) {