Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Add rate limit info to events #1160

Merged
merged 3 commits into from
May 2, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 34 additions & 1 deletion src/__tests__/posthog-core.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -31,5 +31,38 @@ 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_tokens: 99,
},
})
})

it('does not capture if rate limit is in place', () => {
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.'
)
})
})
})
})
20 changes: 10 additions & 10 deletions src/__tests__/rate-limiter.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -53,12 +53,12 @@ describe('Rate Limiter', () => {
tokens: 100,
last: systemTime,
})
expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
expect(rateLimiter.isCaptureClientSideRateLimited().isRateLimited).toBe(false)
})

it('subtracts a token with each call', () => {
range(5).forEach(() => {
expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
expect(rateLimiter.isCaptureClientSideRateLimited().isRateLimited).toBe(false)
})
expect(persistedBucket['$capture_rate_limit']).toEqual({
tokens: 95,
Expand All @@ -68,15 +68,15 @@ describe('Rate Limiter', () => {

it('adds tokens if time has passed ', () => {
range(50).forEach(() => {
expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
expect(rateLimiter.isCaptureClientSideRateLimited().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.isCaptureClientSideRateLimited().isRateLimited).toBe(false)
expect(persistedBucket['$capture_rate_limit']).toEqual({
tokens: 69, // 50 + 20 - 1
last: systemTime,
Expand All @@ -85,18 +85,18 @@ describe('Rate Limiter', () => {

it('rate limits when past the threshold ', () => {
range(100).forEach(() => {
expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
expect(rateLimiter.isCaptureClientSideRateLimited().isRateLimited).toBe(false)
})
range(200).forEach(() => {
expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(true)
expect(rateLimiter.isCaptureClientSideRateLimited().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.isCaptureClientSideRateLimited().isRateLimited).toBe(false)
benjackwhite marked this conversation as resolved.
Show resolved Hide resolved
expect(persistedBucket['$capture_rate_limit']).toEqual({
tokens: 19, // 20 - 1
last: systemTime,
Expand All @@ -105,13 +105,13 @@ describe('Rate Limiter', () => {

it('refills up to the maximum amount ', () => {
range(100).forEach(() => {
expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
expect(rateLimiter.isCaptureClientSideRateLimited().isRateLimited).toBe(false)
})
expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(true)
expect(rateLimiter.isCaptureClientSideRateLimited().isRateLimited).toBe(true)
expect(persistedBucket['$capture_rate_limit'].tokens).toEqual(0)

moveTimeForward(1000000)
expect(rateLimiter.isCaptureClientSideRateLimited()).toBe(false)
expect(rateLimiter.isCaptureClientSideRateLimited().isRateLimited).toBe(false)
expect(persistedBucket['$capture_rate_limit'].tokens).toEqual(99) // limit - 1
})

Expand Down
18 changes: 13 additions & 5 deletions src/posthog-core.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand All @@ -750,6 +745,15 @@ export class PostHog {
return
}

const clientRateLimitContext = !options?.skip_client_rate_limiting
? this.rateLimiter.isCaptureClientSideRateLimited()
: undefined

if (clientRateLimitContext?.isRateLimited) {
logger.critical('This capture call is ignored due to client rate limiting.')
return
}

// update persistence
this.sessionPersistence.update_search_keyword()

Expand Down Expand Up @@ -778,6 +782,10 @@ export class PostHog {
}
}

if (clientRateLimitContext) {
data.properties['$lib_rate_limit_tokens'] = clientRateLimitContext.remainingTokens
benjackwhite marked this conversation as resolved.
Show resolved Hide resolved
}

const setProperties = options?.$set
if (setProperties) {
data.$set = options?.$set
Expand Down
12 changes: 9 additions & 3 deletions src/rate-limiter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,10 +27,13 @@ export class RateLimiter {
this.captureEventsPerSecond
)

this.lastEventRateLimited = this.isCaptureClientSideRateLimited(true)
this.lastEventRateLimited = this.isCaptureClientSideRateLimited(true).isRateLimited
}

public isCaptureClientSideRateLimited(checkOnly = false): boolean {
public isCaptureClientSideRateLimited(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()
Expand Down Expand Up @@ -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 {
Expand Down
Loading