feat: Add rate limit info to events (#1160)

PostHog · May 2, 2024 · c8a158e · c8a158e
1 parent dbb7c66
commit c8a158e
Show file tree

Hide file tree

Showing 5 changed files with 74 additions and 24 deletions.
diff --git a/src/__tests__/posthog-core.js b/src/__tests__/posthog-core.js
@@ -78,7 +78,7 @@ describe('posthog core', () => {
             __captureHooks: [],
             rateLimiter: {
                 isServerRateLimited: () => false,
-                isCaptureClientSideRateLimited: () => false,
+                clientRateLimitContext: () => false,
             },
         }))
 

diff --git a/src/__tests__/posthog-core.test.ts b/src/__tests__/posthog-core.test.ts
@@ -7,7 +7,7 @@ describe('posthog core', () => {
         const properties = {
             event: 'prop',
         }
-        const setup = (config: Partial<PostHogConfig>) => {
+        const setup = (config: Partial<PostHogConfig> = {}) => {
             const onCapture = jest.fn()
             const posthog = _posthog.init('testtoken', { ...config, _onCapture: onCapture }, uuidv7())!
             posthog.debug()
@@ -31,5 +31,41 @@ describe('posthog core', () => {
             expect(actual['$is_identified']).toBeUndefined()
             expect(actual['token']).toBeUndefined()
         })
+
+        describe('rate limiting', () => {
+            it('includes information about remaining rate limit', () => {
+                const { posthog, onCapture } = setup()
+
+                posthog.capture(eventName, properties)
+
+                expect(onCapture.mock.calls[0][1]).toMatchObject({
+                    properties: {
+                        $lib_rate_limit_remaining_tokens: 99,
+                    },
+                })
+            })
+
+            it('does not capture if rate limit is in place', () => {
+                jest.useFakeTimers()
+                jest.setSystemTime(Date.now())
+
+                console.error = jest.fn()
+                const { posthog, onCapture } = setup()
+
+                for (let i = 0; i < 100; i++) {
+                    posthog.capture(eventName, properties)
+                }
+                expect(onCapture).toHaveBeenCalledTimes(100)
+                onCapture.mockClear()
+                ;(console.error as any).mockClear()
+                posthog.capture(eventName, properties)
+                expect(onCapture).toHaveBeenCalledTimes(0)
+                expect(console.error).toHaveBeenCalledTimes(1)
+                expect(console.error).toHaveBeenCalledWith(
+                    '[PostHog.js]',
+                    'This capture call is ignored due to client rate limiting.'
+                )
+            })
+        })
     })
 })
diff --git a/src/__tests__/rate-limiter.test.ts b/src/__tests__/rate-limiter.test.ts
@@ -48,17 +48,17 @@ describe('Rate Limiter', () => {
 
     describe('client side', () => {
         it('starts with the max tokens', () => {
-            rateLimiter.isCaptureClientSideRateLimited(true)
+            rateLimiter.clientRateLimitContext(true)
             expect(persistedBucket['$capture_rate_limit']).toEqual({
                 tokens: 100,
                 last: systemTime,
             })
-            expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
+            expect(rateLimiter.clientRateLimitContext().isRateLimited).toBe(false)
         })
 
         it('subtracts a token with each call', () => {
             range(5).forEach(() => {
-                expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
+                expect(rateLimiter.clientRateLimitContext().isRateLimited).toBe(false)
             })
             expect(persistedBucket['$capture_rate_limit']).toEqual({
                 tokens: 95,
@@ -68,15 +68,15 @@ describe('Rate Limiter', () => {
 
         it('adds tokens if time has passed ', () => {
             range(50).forEach(() => {
-                expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
+                expect(rateLimiter.clientRateLimitContext().isRateLimited).toBe(false)
             })
             expect(persistedBucket['$capture_rate_limit']).toEqual({
                 tokens: 50,
                 last: systemTime,
             })
 
             moveTimeForward(2000) // 2 seconds = 20 tokens
-            expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
+            expect(rateLimiter.clientRateLimitContext().isRateLimited).toBe(false)
             expect(persistedBucket['$capture_rate_limit']).toEqual({
                 tokens: 69, // 50 + 20 - 1
                 last: systemTime,
@@ -85,18 +85,18 @@ describe('Rate Limiter', () => {
 
         it('rate limits when past the threshold ', () => {
             range(100).forEach(() => {
-                expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
+                expect(rateLimiter.clientRateLimitContext().isRateLimited).toBe(false)
             })
             range(200).forEach(() => {
-                expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(true)
+                expect(rateLimiter.clientRateLimitContext().isRateLimited).toBe(true)
             })
             expect(persistedBucket['$capture_rate_limit']).toEqual({
                 tokens: 0,
                 last: systemTime,
             })
 
             moveTimeForward(2000) // 2 seconds = 20 tokens
-            expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
+            expect(rateLimiter.clientRateLimitContext().isRateLimited).toBe(false)
             expect(persistedBucket['$capture_rate_limit']).toEqual({
                 tokens: 19, // 20 - 1
                 last: systemTime,
@@ -105,19 +105,19 @@ describe('Rate Limiter', () => {
 
         it('refills up to the maximum amount ', () => {
             range(100).forEach(() => {
-                expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
+                expect(rateLimiter.clientRateLimitContext().isRateLimited).toBe(false)
             })
-            expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(true)
+            expect(rateLimiter.clientRateLimitContext().isRateLimited).toBe(true)
             expect(persistedBucket['$capture_rate_limit'].tokens).toEqual(0)
 
             moveTimeForward(1000000)
-            expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
+            expect(rateLimiter.clientRateLimitContext().isRateLimited).toBe(false)
             expect(persistedBucket['$capture_rate_limit'].tokens).toEqual(99) // limit - 1
         })
 
         it('captures a rate limit event the first time it is rate limited', () => {
             range(200).forEach(() => {
-                rateLimiter.isCaptureClientSideRateLimited()
+                rateLimiter.clientRateLimitContext()
             })
 
             expect(mockPostHog.capture).toBeCalledTimes(1)
@@ -135,7 +135,7 @@ describe('Rate Limiter', () => {
 
         it('does not capture a rate limit event if the persisted config was already rate limited', () => {
             range(200).forEach(() => {
-                rateLimiter.isCaptureClientSideRateLimited()
+                rateLimiter.clientRateLimitContext()
             })
 
             expect(mockPostHog.capture).toBeCalledTimes(1)
@@ -144,7 +144,7 @@ describe('Rate Limiter', () => {
             const newRateLimiter = new RateLimiter(mockPostHog as any)
 
             range(200).forEach(() => {
-                newRateLimiter.isCaptureClientSideRateLimited()
+                newRateLimiter.clientRateLimitContext()
             })
 
             expect(mockPostHog.capture).toBeCalledTimes(0)

diff --git a/src/posthog-core.ts b/src/posthog-core.ts
@@ -731,11 +731,6 @@ export class PostHog {
             return
         }
 
-        if (!options?.skip_client_rate_limiting && this.rateLimiter.isCaptureClientSideRateLimited()) {
-            logger.critical('This capture call is ignored due to client rate limiting.')
-            return
-        }
-
         // typing doesn't prevent interesting data
         if (isUndefined(event_name) || !isString(event_name)) {
             logger.error('No event name provided to posthog.capture')
@@ -750,6 +745,15 @@ export class PostHog {
             return
         }
 
+        const clientRateLimitContext = !options?.skip_client_rate_limiting
+            ? this.rateLimiter.clientRateLimitContext()
+            : undefined
+
+        if (clientRateLimitContext?.isRateLimited) {
+            logger.critical('This capture call is ignored due to client rate limiting.')
+            return
+        }
+
         // update persistence
         this.sessionPersistence.update_search_keyword()
 
@@ -778,6 +782,10 @@ export class PostHog {
             }
         }
 
+        if (clientRateLimitContext) {
+            data.properties['$lib_rate_limit_remaining_tokens'] = clientRateLimitContext.remainingTokens
+        }
+
         const setProperties = options?.$set
         if (setProperties) {
             data.$set = options?.$set

diff --git a/src/rate-limiter.ts b/src/rate-limiter.ts
@@ -27,10 +27,13 @@ export class RateLimiter {
             this.captureEventsPerSecond
         )
 
-        this.lastEventRateLimited = this.isCaptureClientSideRateLimited(true)
+        this.lastEventRateLimited = this.clientRateLimitContext(true).isRateLimited
     }
 
-    public isCaptureClientSideRateLimited(checkOnly = false): boolean {
+    public clientRateLimitContext(checkOnly = false): {
+        isRateLimited: boolean
+        remainingTokens: number
+    } {
         // This is primarily to prevent runaway loops from flooding capture with millions of events for a single user.
         // It's as much for our protection as theirs.
         const now = new Date().getTime()
@@ -67,7 +70,10 @@ export class RateLimiter {
         this.lastEventRateLimited = isRateLimited
         this.instance.persistence?.set_property(CAPTURE_RATE_LIMIT, bucket)
 
-        return isRateLimited
+        return {
+            isRateLimited,
+            remainingTokens: bucket.tokens,
+        }
     }
 
     public isServerRateLimited(batchKey: string | undefined): boolean {