From 3f2f72d811de78ada57f6be364adbf16c2a5eb2e Mon Sep 17 00:00:00 2001 From: "Reichenbach, Michael" <755327+Silthus@users.noreply.github.com> Date: Mon, 19 Oct 2026 09:39:53 +0000 Subject: [PATCH] feat(workflows): cap sandbox sends and refund ses throttles Co-Authored-By: Claude Opus 5.5 --- .../internal/workflows-email-sending-tiers.md | 9 +- nodejs/src/cdp/cdp-services.ts | 24 +- nodejs/src/cdp/config.ts | 5 + .../services/messaging/email.service.test.ts | 373 ++++++++++++++++-- .../cdp/services/messaging/email.service.ts | 77 +++- .../messaging/sandbox-email-sender.test.ts | 2 + .../messaging/sandbox-email-sender.ts | 93 ++++- .../rate-limiter/rate-limiter.service.test.ts | 127 +++++- .../rate-limiter/rate-limiter.service.ts | 151 +++++++ nodejs/src/server.ts | 23 +- 10 files changed, 820 insertions(+), 64 deletions(-) diff --git a/docs/internal/workflows-email-sending-tiers.md b/docs/internal/workflows-email-sending-tiers.md index d1378b2726b2..4efe19cd1105 100644 --- a/docs/internal/workflows-email-sending-tiers.md +++ b/docs/internal/workflows-email-sending-tiers.md @@ -70,7 +70,14 @@ Decay, suspension drops, admin recomputes, and the backfill stay silent. - The batch audience cap is decided when the batch is dispatched. Adding an email step to the workflow while a batch is queued does not re-cap it; the send-time buckets still cap every email at execution. This is why enforcement requires the worker caps to be deployed (see the rollout order). - Test-panel sends bypass the team buckets on purpose, matching the per-workflow rate limit. -- Sandbox sender sends bypass the team buckets without spending their tokens. Sandbox limits run independently. Keep `WORKFLOWS_SANDBOX_SENDER_ENABLED` off until the recipient membership checks, sandbox caps and sandbox pause gate are deployed. +- Sandbox sender sends bypass the team buckets without spending their tokens. Sandbox limits run independently. + The worker holds every sandbox send, test sends included, to a daily cap per project (`WORKFLOWS_SANDBOX_DAILY_TEAM_CAP`, one token per recipient copy) and per recipient address within that organization (`WORKFLOWS_SANDBOX_DAILY_RECIPIENT_CAP`). These are continuously refilling token buckets, not calendar-day counters; an idle bucket starts full and refills its capacity over 24 hours. + A reached cap skips the send instead of delaying it. A limiter error or an empty or invalid cap value skips every sandbox send. + Cap keys and the final membership check share one organization snapshot, valid for at most one minute. A stalled claim or member lookup that outlives that snapshot skips the send and returns the claim. + SES throttle responses return the sandbox claim before the existing delivery retry. Returns are deduplicated, and Valkey ignores a return command that arrives more than five minutes after it was issued. + Test sends run inline in the CDP API, so the API reads the same `SES_RATE_LIMITER_VALKEY_*` settings as the email worker. Without them, every sandbox test send is skipped. + CDP API opens the SES Valkey pool only when sandbox delivery is enabled; otherwise its MX validation uses the existing local cache and DNS. + Deploy membership checks, daily caps, and the shared pause gate first. Configure the SES identity and Valkey pool, enable `WORKFLOWS_SANDBOX_SENDER_ENABLED`, then enable the UI flag. - The buckets are token buckets: a full idle bucket plus refill allows up to roughly twice the stated cap in the very first period. The bucket TTLs exceed the refill periods so this does not recur from idling. - A denied send parks until every short bucket has refilled enough to cover it, instead of retrying on a fixed few-minute cadence. The computed wait is capped at one hour and then jittered 1x to 2x, so a parked send can wait just under two hours between attempts. diff --git a/nodejs/src/cdp/cdp-services.ts b/nodejs/src/cdp/cdp-services.ts index 3b60a0e40a33..41373d89b0f4 100644 --- a/nodejs/src/cdp/cdp-services.ts +++ b/nodejs/src/cdp/cdp-services.ts @@ -170,6 +170,8 @@ export type CdpCoreServicesConfig = Pick< | 'SES_TRACKED_CONFIGURATION_SET' | 'SES_UNTRACKED_CONFIGURATION_SET' | 'WORKFLOWS_SANDBOX_SENDER_ENABLED' + | 'WORKFLOWS_SANDBOX_DAILY_TEAM_CAP' + | 'WORKFLOWS_SANDBOX_DAILY_RECIPIENT_CAP' | 'SES_SANDBOX_TENANT_NAME' | 'SES_SANDBOX_CONFIGURATION_SET' | 'SES_SANDBOX_FROM_ADDRESS' @@ -214,13 +216,15 @@ export interface CdpCoreServicesDeps { /** * SES Valkey pool shared with the SES rate limiter, opened only on pods whose * capabilities actually execute email actions (hogflow/email cyclotron workers). - * `null` on every other CDP consumer and on cdp-api so idle pods don't hold - * open connections against the SES Valkey instance. + * Sandbox-enabled cdp-api opens it too, for the sandbox daily caps of test sends. + * `null` on every other CDP consumer so idle pods don't hold open connections + * against the SES Valkey instance. */ emailValidationValkey: RedisV2 | null /** - * Dedicated frequency cap Valkey pool, opened only on the same pods as `emailValidationValkey`. - * `null` elsewhere, and when no host is configured, which turns the cap off. + * Dedicated frequency cap Valkey pool, opened only on the email workers. `null` elsewhere, + * cdp-api included because its test sends skip the cap, and when no host is configured, + * which turns the cap off. */ frequencyCapValkey: RedisV2 | null } @@ -458,6 +462,9 @@ export function createCdpCoreServices( const teamEmailRateLimiter = deps.emailValidationValkey ? new RateLimiterService(deps.emailValidationValkey, { name: 'team-email' }) : null + const sandboxEmailRateLimiter = deps.emailValidationValkey + ? new RateLimiterService(deps.emailValidationValkey, { name: 'sandbox-email' }) + : null const emailService = new EmailService( { sesAccessKeyId: config.SES_ACCESS_KEY_ID, @@ -486,8 +493,11 @@ export function createCdpCoreServices( tenantName: config.SES_SANDBOX_TENANT_NAME, configurationSetName: config.SES_SANDBOX_CONFIGURATION_SET, fromAddress: config.SES_SANDBOX_FROM_ADDRESS, + dailyTeamCap: config.WORKFLOWS_SANDBOX_DAILY_TEAM_CAP, + dailyRecipientCap: config.WORKFLOWS_SANDBOX_DAILY_RECIPIENT_CAP, }, deps.teamManager, + sandboxEmailRateLimiter, redis ), new OrganizationMembersService(deps.postgres) @@ -556,9 +566,9 @@ export function createCdpCoreServices( : undefined ) // MX verdicts live on the dedicated SES Valkey (same instance as the SES rate - // limiter, separate pool). The pool is created by the server only on pods - // whose capabilities execute email actions; everywhere else this is null - // and EmailValidationService degrades to the local cache + DNS. + // limiter, separate pool). The server creates the pool only on pods whose + // capabilities execute email actions and on sandbox-enabled cdp-api; everywhere + // else this is null and EmailValidationService degrades to the local cache + DNS. const emailValidationService = new EmailValidationService(deps.emailValidationValkey) const cohortMembershipRepository = new PostgresCohortMembershipRepository(deps.postgres) // Observer writes to both stores; the read source decides which verdict drives the metric. diff --git a/nodejs/src/cdp/config.ts b/nodejs/src/cdp/config.ts index 383a6e6091a2..3e379233333d 100644 --- a/nodejs/src/cdp/config.ts +++ b/nodejs/src/cdp/config.ts @@ -167,6 +167,9 @@ export type CdpConfig = ClickhouseConfig & { // Empty means not provisioned: tracking-off sends fall back to the tracked set with a warning. SES_UNTRACKED_CONFIGURATION_SET: string WORKFLOWS_SANDBOX_SENDER_ENABLED: boolean + // An empty or invalid value blocks every sandbox send. + WORKFLOWS_SANDBOX_DAILY_TEAM_CAP: string + WORKFLOWS_SANDBOX_DAILY_RECIPIENT_CAP: string SES_SANDBOX_TENANT_NAME: string SES_SANDBOX_CONFIGURATION_SET: string SES_SANDBOX_FROM_ADDRESS: string @@ -412,6 +415,8 @@ export function getDefaultCdpConfig(): CdpConfig { SES_TRACKED_CONFIGURATION_SET: 'posthog-messaging', SES_UNTRACKED_CONFIGURATION_SET: '', WORKFLOWS_SANDBOX_SENDER_ENABLED: false, + WORKFLOWS_SANDBOX_DAILY_TEAM_CAP: '25', + WORKFLOWS_SANDBOX_DAILY_RECIPIENT_CAP: '5', SES_SANDBOX_TENANT_NAME: '', SES_SANDBOX_CONFIGURATION_SET: '', SES_SANDBOX_FROM_ADDRESS: '', diff --git a/nodejs/src/cdp/services/messaging/email.service.test.ts b/nodejs/src/cdp/services/messaging/email.service.test.ts index 3ca0fce8c2f5..4343e576e8bf 100644 --- a/nodejs/src/cdp/services/messaging/email.service.test.ts +++ b/nodejs/src/cdp/services/messaging/email.service.test.ts @@ -14,7 +14,7 @@ import { CyclotronInvocationQueueParametersEmailSchema, CyclotronInvocationQueueParametersEmailType, } from '~/cdp/schema/cyclotron' -import { CyclotronJobInvocationHogFunction } from '~/cdp/types' +import { CyclotronJobInvocationHogFunction, CyclotronJobInvocationResult } from '~/cdp/types' import { IngestionOutputs } from '~/common/outputs/ingestion-outputs' import { defineLuaTokenBucketV2 } from '~/common/redis/redis-token-bucket-v2.lua' import { defineLuaTokenBucketV3 } from '~/common/redis/redis-token-bucket-v3.lua' @@ -269,10 +269,11 @@ describe('EmailService', () => { extraOrganizationIds.push(id) return id } - let sandboxRedisClient: Redis.Redis | undefined + let sandboxRedisClients: Redis.Redis[] + let dailyCapLimiter: RateLimiterService const createSandboxLimiter = async (name: string): Promise => { - sandboxRedisClient = await hub.redisPool.acquire() - const client = sandboxRedisClient + const client = await hub.redisPool.acquire() + sandboxRedisClients.push(client) defineLuaTokenBucketV2(client) defineLuaTokenBucketV3(client) const redis: RedisV2 = { @@ -287,10 +288,23 @@ describe('EmailService', () => { } const createSandboxService = ( enabled: boolean, - tierLimiter: RateLimiterService | null = null, - messageAssetsService?: MessageAssetsService, - membersPostgres: PostgresRouter = hub.postgres, - routeSettings: Partial = {} + { + tierLimiter = null, + messageAssetsService, + membersPostgres = hub.postgres, + capLimiter = dailyCapLimiter, + dailyTeamCap = '100', + dailyRecipientCap = '100', + routeSettings = {}, + }: { + tierLimiter?: RateLimiterService | null + messageAssetsService?: MessageAssetsService + membersPostgres?: PostgresRouter + capLimiter?: RateLimiterService | null + dailyTeamCap?: string + dailyRecipientCap?: string + routeSettings?: Partial + } = {} ): EmailService => { const sandboxService = new EmailService( { @@ -320,18 +334,15 @@ describe('EmailService', () => { tenantName: 'sandbox-tenant', configurationSetName: 'sandbox-email', fromAddress: 'fixed-sandbox@example.com', + dailyTeamCap, + dailyRecipientCap, ...routeSettings, }, hub.teamManager, + capLimiter, { - useClient: async (_options, callback) => { - const client = await hub.redisPool.acquire() - try { - return await callback(client as unknown as RedisClient) - } finally { - await hub.redisPool.release(client) - } - }, + useClient: async (_options, callback) => + callback(sandboxRedisClients[0] as unknown as RedisClient), } ), new OrganizationMembersService(membersPostgres) @@ -343,6 +354,8 @@ describe('EmailService', () => { beforeEach(async () => { capture = jest.spyOn(posthog, 'captureTeamEvent').mockImplementation(() => {}) + sandboxRedisClients = [] + dailyCapLimiter = await createSandboxLimiter('sandbox-daily-cap-test') memberIds = [] extraOrganizationIds = [] memberId = await createMember(memberEmail()) @@ -390,6 +403,10 @@ describe('EmailService', () => { [extraOrganizationIds], 'test:delete-sandbox-organizations' ) + capture.mockRestore() + for (const client of sandboxRedisClients) { + await hub.redisPool.release(client) + } const client = await hub.redisPool.acquire() try { const keys = await client.keys(`@posthog/workflows-sandbox-blocked/${team.id}/*`) @@ -399,11 +416,6 @@ describe('EmailService', () => { } finally { await hub.redisPool.release(client) } - capture.mockRestore() - if (sandboxRedisClient) { - await hub.redisPool.release(sandboxRedisClient) - sandboxRedisClient = undefined - } }) it('skips a sandbox send to a person without an email address as a missing recipient', async () => { @@ -524,7 +536,7 @@ describe('EmailService', () => { if (checkFailed) { await membersPostgres.end() } - service = createSandboxService(true, null, undefined, membersPostgres) + service = createSandboxService(true, { membersPostgres }) invocation.queueParameters = createSandboxParams({ from: { integrationId: 4 }, cc: outside.slice(0, 100).join(', '), @@ -778,7 +790,7 @@ describe('EmailService', () => { const preparation = jest.spyOn(EmailSuppressionService.prototype, 'isSuppressed') const suppress = preparation.getMockImplementation()! try { - service = createSandboxService(true) + service = createSandboxService(true, { dailyTeamCap: '3' }) const first = await service.executeSendEmail(invocation, isTest) expect(first.error).toBeUndefined() expect(sendEmailSpy).toHaveBeenCalledTimes(1) @@ -808,6 +820,16 @@ describe('EmailService', () => { 'workflows sandbox email blocked', { reason: 'recipient_not_member', is_test: isTest, blocked_recipient_count: 1 } ) + + invocation.queueParameters = createSandboxParams({ + from: { integrationId: 4 }, + to: { email: memberEmail('cc') }, + bcc: memberEmail('bcc'), + }) + const afterBlock = await service.executeSendEmail(invocation, isTest) + expect(afterBlock.error).toBeUndefined() + expect(afterBlock.skipped).not.toBe(true) + expect(sendEmailSpy).toHaveBeenCalledTimes(2) } finally { clock.mockRestore() preparation.mockRestore() @@ -830,7 +852,7 @@ describe('EmailService', () => { [team.id], 'test:own-sandbox-team-lookup' ) - service = createSandboxService(true, null, undefined, membersPostgres) + service = createSandboxService(true, { membersPostgres }) if (warm) { const first = await service.executeSendEmail(invocation, isTest) expect(first.error).toBeUndefined() @@ -938,7 +960,7 @@ describe('EmailService', () => { } const memberQueries = jest.spyOn(membersPostgres, 'query') try { - service = createSandboxService(true, null, undefined, membersPostgres) + service = createSandboxService(true, { membersPostgres }) const result = await service.executeSendEmail(invocation, isTest) @@ -982,7 +1004,7 @@ describe('EmailService', () => { const clock = jest.spyOn(Date, 'now').mockReturnValue(now) let databaseClosed = false try { - service = createSandboxService(true, null, undefined, membersPostgres) + service = createSandboxService(true, { membersPostgres }) const params = createSandboxParams({ from: { integrationId: 4 } }) const first = await service.executeSendEmail(invocation) expect(first.error).toBeUndefined() @@ -1091,7 +1113,7 @@ describe('EmailService', () => { ] as const)( 'skips when the sandbox route is unavailable because %s (isTest=%s)', async (_reason, isTest, enabled: boolean, routeSettings: Partial) => { - service = createSandboxService(enabled, null, undefined, hub.postgres, routeSettings) + service = createSandboxService(enabled, { routeSettings }) const result = await service.executeSendEmail(invocation, isTest) expect(result).toMatchObject({ finished: true, skipped: true, metrics: [] }) @@ -1129,7 +1151,7 @@ describe('EmailService', () => { const outputs = { produce: jest.fn().mockResolvedValue(undefined), } as unknown as IngestionOutputs<'message_assets'> - service = createSandboxService(true, null, new MessageAssetsService(outputs)) + service = createSandboxService(true, { messageAssetsService: new MessageAssetsService(outputs) }) invocation.state.actionId = 'send-email' invocation.queueParameters = createSandboxParams({ from: { integrationId: 4, email: 'override@example.com', name: 'Custom sender' }, @@ -1260,7 +1282,10 @@ describe('EmailService', () => { produce: jest.fn().mockResolvedValue(undefined), } as unknown as IngestionOutputs<'message_assets'> const limiter = await createSandboxLimiter('sandbox-rejected-html-budget-test') - service = createSandboxService(true, limiter, new MessageAssetsService(outputs)) + service = createSandboxService(true, { + tierLimiter: limiter, + messageAssetsService: new MessageAssetsService(outputs), + }) invocation.state.actionId = 'send-email' const params = createSandboxParams({ from: { integrationId: 4 }, text, html }) invocation.queueParameters = params @@ -1363,7 +1388,7 @@ describe('EmailService', () => { it('bypasses an exhausted sending tier and leaves its budget for own senders', async () => { const limiter = await createSandboxLimiter('sandbox-tier-budget-test') - service = createSandboxService(true, limiter) + service = createSandboxService(true, { tierLimiter: limiter }) const sandbox = await service.executeSendEmail(invocation) expect(sandbox).toMatchObject({ finished: true }) expect(sandbox.error).toBeUndefined() @@ -1385,6 +1410,294 @@ describe('EmailService', () => { expect(afterExhaustion.error).toBeUndefined() expect(sendEmailSpy).toHaveBeenCalledTimes(3) }) + + describe('daily caps', () => { + it('skips and refunds when the organization snapshot expires during the cap claim', async () => { + const now = Date.now() + const clock = jest.spyOn(Date, 'now').mockReturnValue(now) + const claim = dailyCapLimiter.claimAllOrNothing.bind(dailyCapLimiter) + const delayedClaim = jest.spyOn(dailyCapLimiter, 'claimAllOrNothing') + try { + service = createSandboxService(true, { dailyTeamCap: '2', dailyRecipientCap: '2' }) + expectSent(await send(true, {})) + delayedClaim.mockImplementationOnce(async (buckets) => { + const grant = await claim(buckets) + clock.mockReturnValue(now + 61_000) + return grant + }) + + const expired = await send(true, {}) + expect(expired).toMatchObject({ finished: true, skipped: true }) + expect(expired.error).toBeUndefined() + expect(sendEmailSpy).toHaveBeenCalledTimes(1) + expect(expired.logs).toContainEqual( + expect.objectContaining({ + message: expect.stringContaining('Skipping send: could not check organization members'), + }) + ) + + clock.mockReturnValue(now + 122_000) + expectSent(await send(true, {})) + expect(sendEmailSpy).toHaveBeenCalledTimes(2) + } finally { + delayedClaim.mockRestore() + clock.mockRestore() + } + }) + + it('uses the membership organization for caps when the team cache is older', async () => { + const now = Date.now() + const clock = jest.spyOn(Date, 'now').mockReturnValue(now) + const nextOrganization = await createOtherOrganization() + try { + service = createSandboxService(true, { dailyRecipientCap: '1' }) + expectSent(await send(true, {})) + await createOrganizationMembership(hub.postgres, nextOrganization, memberId) + await hub.postgres.query( + PostgresUse.COMMON_WRITE, + 'UPDATE posthog_team SET organization_id = $1 WHERE id = $2', + [nextOrganization, team.id], + 'test:transfer-sandbox-project' + ) + clock.mockReturnValue(now + 61_000) + + expectSent(await send(true, {})) + expect(sendEmailSpy).toHaveBeenCalledTimes(2) + } finally { + clock.mockRestore() + await hub.postgres.query( + PostgresUse.COMMON_WRITE, + 'UPDATE posthog_team SET organization_id = $1 WHERE id = $2', + [team.organization_id, team.id], + 'test:restore-sandbox-project-organization' + ) + } + }) + + const TEAM_CAP_REACHED = + "Skipping send: this email would go over the sandbox sender's daily limit for this project. Verify your own domain to send more." + const recipientCapReached = (addresses: string): string => + `Skipping send: these addresses reached the sandbox sender's daily limit per address: ${addresses}. Verify your own domain to send more.` + const CAP_CHECK_FAILED = + "Skipping send: could not check the sandbox sender's daily limit. Try again later, or verify your own domain to send more." + + const send = async ( + isTest: boolean, + params: Partial + ): Promise> => { + invocation.state.vmState = { stack: [] } as any + invocation.queueParameters = createSandboxParams({ from: { integrationId: 4 }, ...params }) + return await service.executeSendEmail(invocation, isTest) + } + const expectSent = (result: CyclotronJobInvocationResult): void => { + expect(result.error).toBeUndefined() + expect(result.skipped).not.toBe(true) + expect(result.invocation.state.vmState?.stack).toEqual([{ success: true }]) + } + const expectSkipped = ( + result: CyclotronJobInvocationResult, + isTest: boolean, + message: string, + blocked: { reason: 'cap_reached' | 'check_failed'; blocked_recipient_count: number } + ): void => { + expect(result).toMatchObject({ finished: true, skipped: true, metrics: [], messageAssets: [] }) + expect(result.error).toBeUndefined() + expect(result.invocation.queueScheduledAt).toBeUndefined() + expect(result.invocation.state.vmState?.stack).toEqual([{ success: false }]) + expect(result.logs).toContainEqual(expect.objectContaining({ level: 'info', message })) + expect(capture).toHaveBeenCalledWith( + expect.objectContaining({ id: team.id }), + 'workflows sandbox email blocked', + expect.objectContaining({ reason: blocked.reason, is_test: isTest }) + ) + } + + it.each([false, true])( + 'holds the project to its daily cap without charging denied or own-sender sends (isTest=%s)', + async (isTest) => { + service = createSandboxService(true, { dailyTeamCap: '3', dailyRecipientCap: '1' }) + + expectSent(await send(isTest, { from: { integrationId: 1 } })) + expectSent(await send(isTest, { to: { email: memberEmail() } })) + expectSkipped( + await send(isTest, { + to: { email: memberEmail() }, + cc: memberEmail('cc'), + bcc: memberEmail('bcc'), + }), + isTest, + TEAM_CAP_REACHED, + { reason: 'cap_reached', blocked_recipient_count: 3 } + ) + expectSent(await send(isTest, { to: { email: memberEmail('cc') } })) + expectSent(await send(isTest, { to: { email: memberEmail('bcc') } })) + expectSkipped(await send(isTest, { to: { email: memberEmail() } }), isTest, TEAM_CAP_REACHED, { + reason: 'cap_reached', + blocked_recipient_count: 1, + }) + expect(sendEmailSpy).toHaveBeenCalledTimes(4) + } + ) + + it.each([false, true])( + 'skips a send with recipients at the daily cap and leaves the others their budget (isTest=%s)', + async (isTest) => { + service = createSandboxService(true, { dailyTeamCap: '5', dailyRecipientCap: '1' }) + + expectSent(await send(isTest, { to: { email: memberEmail('cc') }, cc: memberEmail('bcc') })) + expectSkipped( + await send(isTest, { + to: { email: memberEmail() }, + cc: memberEmail('cc').toUpperCase(), + bcc: memberEmail('bcc'), + }), + isTest, + recipientCapReached(`${memberEmail('cc').toUpperCase()}, ${memberEmail('bcc')}`), + { reason: 'cap_reached', blocked_recipient_count: 2 } + ) + expectSent(await send(isTest, { to: { email: memberEmail() } })) + expect(sendEmailSpy).toHaveBeenCalledTimes(2) + } + ) + + it.each([false, true])( + 'charges a repeated address once per address and once per copy for the project (isTest=%s)', + async (isTest) => { + service = createSandboxService(true, { dailyTeamCap: '3', dailyRecipientCap: '1' }) + + expectSent( + await send(isTest, { + to: { email: memberEmail() }, + cc: memberEmail().toUpperCase(), + bcc: ` ${memberEmail()} `, + }) + ) + expectSkipped( + await send(isTest, { to: { email: memberEmail('cc') } }), + isTest, + TEAM_CAP_REACHED, + { reason: 'cap_reached', blocked_recipient_count: 1 } + ) + expect(sendEmailSpy).toHaveBeenCalledTimes(1) + } + ) + + it('refills a spent sandbox allowance over a day rather than seconds', async () => { + service = createSandboxService(true, { dailyTeamCap: '1', dailyRecipientCap: '1' }) + expectSent(await send(true, {})) + const client = sandboxRedisClients[0] + const keys = ( + await client.keys(`@posthog/workflows-sandbox-daily/{${team.organization_id}}/*`) + ).filter((key) => !key.includes('/claimed/')) + const [seconds, microseconds] = await client.time() + const redisNow = Number(seconds) * 1000 + Math.floor(Number(microseconds) / 1000) + for (const key of keys) { + await client.hset(key, 'ts', redisNow - 12 * 60 * 60 * 1000) + } + expectSkipped(await send(true, {}), true, TEAM_CAP_REACHED, { + reason: 'cap_reached', + blocked_recipient_count: 1, + }) + for (const key of keys) { + await client.hset(key, 'ts', redisNow - 24 * 60 * 60 * 1000 - 1000) + } + expectSent(await send(true, {})) + expect(sendEmailSpy).toHaveBeenCalledTimes(2) + }) + + it.each([false, true])( + 'spends the daily caps once for a send that SES throttles before accepting it (isTest=%s)', + async (isTest) => { + service = createSandboxService(true, { dailyTeamCap: '2', dailyRecipientCap: '1' }) + sendEmailSpy.mockRejectedValueOnce(new ThrottlingException('Rate exceeded')) + + const throttled = await send(isTest, { cc: memberEmail('cc') }) + expect(throttled.finished).toBe(false) + expectSent(await service.executeSendEmail(throttled.invocation, isTest)) + expectSkipped( + await send(isTest, { to: { email: memberEmail('bcc') } }), + isTest, + TEAM_CAP_REACHED, + { reason: 'cap_reached', blocked_recipient_count: 1 } + ) + expect(sendEmailSpy).toHaveBeenCalledTimes(2) + } + ) + + it.each([false, true])( + 'rechecks membership after a delayed cap claim and returns a blocked grant (isTest=%s)', + async (isTest) => { + const now = Date.now() + const clock = jest.spyOn(Date, 'now').mockReturnValue(now) + service = createSandboxService(true, { dailyTeamCap: '1', dailyRecipientCap: '1' }) + const claim = dailyCapLimiter.claimAllOrNothing.bind(dailyCapLimiter) + const delayedClaim = jest + .spyOn(dailyCapLimiter, 'claimAllOrNothing') + .mockImplementationOnce(async (buckets) => { + const grant = await claim(buckets) + await hub.postgres.query( + PostgresUse.COMMON_WRITE, + 'UPDATE posthog_user SET is_email_verified = false WHERE id = $1', + [memberId], + 'test:revoke-sandbox-verification-during-cap-claim' + ) + clock.mockReturnValue(now + 61_000) + return grant + }) + try { + const blocked = await send(isTest, {}) + expect(blocked).toMatchObject({ finished: true, skipped: true, metrics: [] }) + expect(blocked.error).toBeUndefined() + expect(blocked.invocation.state.vmState?.stack).toEqual([{ success: false }]) + expect(sendEmailSpy).not.toHaveBeenCalled() + + await hub.postgres.query( + PostgresUse.COMMON_WRITE, + 'UPDATE posthog_user SET is_email_verified = true WHERE id = $1', + [memberId], + 'test:restore-sandbox-verification' + ) + clock.mockReturnValue(now + 122_000) + expectSent(await send(isTest, {})) + expect(sendEmailSpy).toHaveBeenCalledTimes(1) + } finally { + delayedClaim.mockRestore() + clock.mockRestore() + } + } + ) + + it.each([ + [ + 'the limiter is unreachable', + false, + { + capLimiter: new RateLimiterService( + { + useClient: () => Promise.reject(new Error('Connection is closed.')), + usePipeline: () => Promise.reject(new Error('Connection is closed.')), + }, + { name: 'sandbox-unreachable-test' } + ), + }, + ], + ['no limiter is configured', true, { capLimiter: null }], + ['the project cap is empty', false, { dailyTeamCap: '' }], + ['the recipient cap is empty', true, { dailyRecipientCap: ' ' }], + ['the project cap is zero', true, { dailyTeamCap: '0' }], + ['the recipient cap is fractional', false, { dailyRecipientCap: '1.5' }], + ['the project cap has trailing text', true, { dailyTeamCap: '25 per day' }], + ['the recipient cap is negative', false, { dailyRecipientCap: '-5' }], + ] as const)('skips every send when %s (isTest=%s)', async (_name, isTest, options) => { + service = createSandboxService(true, options) + + expectSkipped(await send(isTest, { cc: memberEmail('cc') }), isTest, CAP_CHECK_FAILED, { + reason: 'check_failed', + blocked_recipient_count: 2, + }) + expect(sendEmailSpy).not.toHaveBeenCalled() + }) + }) }) describe('integration validation', () => { beforeEach(async () => { diff --git a/nodejs/src/cdp/services/messaging/email.service.ts b/nodejs/src/cdp/services/messaging/email.service.ts index 2ff9f7e51116..8094da1eb54a 100644 --- a/nodejs/src/cdp/services/messaging/email.service.ts +++ b/nodejs/src/cdp/services/messaging/email.service.ts @@ -35,7 +35,7 @@ import { EmailTrackingCodeSigner, TRACKING_CODE_HEADER_NAME } from './helpers/tr import { addUtmTagsToEmail, renderUtmOverrides, resolveUtmTags } from './helpers/utm' import { MessageAssetsService } from './message-assets.service' import { RecipientTokensService } from './recipient-tokens.service' -import { SandboxEmailSender } from './sandbox-email-sender' +import { SandboxDailyCapGrant, SandboxEmailSender } from './sandbox-email-sender' const sesThrottleResponsesTotal = new Counter({ name: 'cdp_ses_throttle_responses_total', @@ -396,10 +396,10 @@ function sandboxAddressList(value?: string): string[] { .map((entry) => angleAddress(entry) ?? entry) } -type SandboxRecipientBlockReason = 'recipient_not_member' | 'recipient_unverified' | 'check_failed' +type SandboxAddressBlockReason = 'recipient_not_member' | 'recipient_unverified' | 'check_failed' | 'cap_reached' -const SANDBOX_RECIPIENT_BLOCK_COPY: Record< - SandboxRecipientBlockReason, +const SANDBOX_ADDRESS_BLOCK_COPY: Record< + SandboxAddressBlockReason, { explanation: string; label: string; guidance: string } > = { recipient_not_member: { @@ -419,8 +419,20 @@ const SANDBOX_RECIPIENT_BLOCK_COPY: Record< label: ' for these addresses: ', guidance: 'Try again, or verify your own domain to send to anyone.', }, + cap_reached: { + explanation: "Skipping send: these addresses reached the sandbox sender's daily limit per address", + label: ': ', + guidance: 'Verify your own domain to send more.', + }, } +const SANDBOX_CAP_SKIP_MESSAGES = { + project_cap_reached: + "Skipping send: this email would go over the sandbox sender's daily limit for this project. Verify your own domain to send more.", + check_failed: + "Skipping send: could not check the sandbox sender's daily limit. Try again later, or verify your own domain to send more.", +} as const + const BLOCKED_ADDRESSES_LABEL = 'Blocked addresses: ' function splitByLength(text: string, maxLength: number): string[] { @@ -1067,7 +1079,7 @@ export class EmailService { ): Promise { const invocation = result.invocation let blockedRecipients: string[] - let reason: SandboxRecipientBlockReason + let reason: Exclude try { if (!this.organizationMembers) { throw new Error('Organization member checks are not configured') @@ -1103,14 +1115,42 @@ export class EmailService { return false } + private async claimSandboxDailyCaps( + result: CyclotronJobInvocationResult, + recipients: string[], + isTest: boolean + ): Promise { + const teamId = result.invocation.teamId + const organization = await this.organizationMembers?.getOrganizationSnapshot(teamId).catch(() => null) + const claim = organization + ? await this.sandboxSender!.claimDailyCaps(teamId, recipients, organization) + : { type: 'check_failed' as const } + if (claim.type === 'granted') { + return claim + } + result.skipped = true + result.invocation.state.vmState?.stack.push({ success: false }) + if (claim.type === 'recipient_cap_reached') { + this.logSandboxRecipientBlock(result, claim.addresses, 'cap_reached') + } else { + createAddLogFunction(result.logs)('info', SANDBOX_CAP_SKIP_MESSAGES[claim.type]) + } + await this.sandboxSender!.capture(teamId, isTest, { + type: 'blocked', + reason: claim.type === 'check_failed' ? 'check_failed' : 'cap_reached', + blockedRecipientCount: claim.type === 'recipient_cap_reached' ? claim.addresses.length : recipients.length, + }) + return null + } + private logSandboxRecipientBlock( result: CyclotronJobInvocationResult, blockedRecipients: string[], - reason: SandboxRecipientBlockReason + reason: SandboxAddressBlockReason ): void { const addLog = createAddLogFunction(result.logs) const addresses = blockedRecipients.map((address) => address.trim()) - const { explanation, label, guidance } = SANDBOX_RECIPIENT_BLOCK_COPY[reason] + const { explanation, label, guidance } = SANDBOX_ADDRESS_BLOCK_COPY[reason] const message = `${explanation}${label}${addresses.join(', ')}. ${guidance}` if (sanitizeLogMessage([message]) === message) { addLog('info', message) @@ -1483,15 +1523,17 @@ export class EmailService { sendEmailParams.Destination!.BccAddresses = bccAddresses } - if ( - integration.config.provider === 'sandbox' && - !(await this.checkSandboxRecipients( - result, - [params.to.email, ...(ccAddresses ?? []), ...(bccAddresses ?? [])], - isTest - )) - ) { - return false + let dailyCapGrant: SandboxDailyCapGrant | null = null + if (integration.config.provider === 'sandbox') { + const sandboxRecipients = [params.to.email, ...(ccAddresses ?? []), ...(bccAddresses ?? [])] + dailyCapGrant = await this.claimSandboxDailyCaps(result, sandboxRecipients, isTest) + if (!dailyCapGrant) { + return false + } + if (!(await this.checkSandboxRecipients(result, sandboxRecipients, isTest, dailyCapGrant.organization))) { + await this.sandboxSender!.returnDailyCaps(dailyCapGrant) + return false + } } try { const response = await this.sesV2Client.send(new SendEmailCommand(sendEmailParams)) @@ -1501,6 +1543,9 @@ export class EmailService { } catch (error: unknown) { if (isSesThrottleError(error)) { sesThrottleResponsesTotal.inc({ error_code: error.name }) + if (dailyCapGrant) { + await this.sandboxSender!.returnDailyCaps(dailyCapGrant) + } throw new SESThrottleError(error.name, pickThrottleRetryDelayMs(), error.message) } const message = error instanceof Error ? error.message : String(error) diff --git a/nodejs/src/cdp/services/messaging/sandbox-email-sender.test.ts b/nodejs/src/cdp/services/messaging/sandbox-email-sender.test.ts index e14c85849250..a8aa11e26af3 100644 --- a/nodejs/src/cdp/services/messaging/sandbox-email-sender.test.ts +++ b/nodejs/src/cdp/services/messaging/sandbox-email-sender.test.ts @@ -13,6 +13,8 @@ const sender = new SandboxEmailSender( tenantName: 'sandbox-tenant', configurationSetName: 'sandbox-email', fromAddress: 'fixed-sandbox@example.com', + dailyTeamCap: '25', + dailyRecipientCap: '5', }, { getTeam: jest.fn().mockResolvedValue({ organization_id: '00000000-0000-0000-0000-000000000001' }), diff --git a/nodejs/src/cdp/services/messaging/sandbox-email-sender.ts b/nodejs/src/cdp/services/messaging/sandbox-email-sender.ts index 0186bcff2c23..c333dc97b603 100644 --- a/nodejs/src/cdp/services/messaging/sandbox-email-sender.ts +++ b/nodejs/src/cdp/services/messaging/sandbox-email-sender.ts @@ -1,3 +1,4 @@ +import { createHash, randomUUID } from 'crypto' import { type DefaultTreeAdapterMap, defaultTreeAdapter, parse, parseFragment, serialize } from 'parse5' import { CyclotronInvocationQueueParametersEmailType } from '~/cdp/schema/cyclotron' @@ -6,13 +7,20 @@ import { logger } from '~/common/utils/logger' import { captureTeamEvent } from '~/common/utils/posthog' import { TeamManager } from '~/common/utils/team-manager' +import { TeamOrganizationSnapshot } from '../managers/organization-members.service' +import { ClaimRequest, RateLimiterService } from '../rate-limiter/rate-limiter.service' import { maybeAddPreheaderToEmail } from './helpers/preheader' +const SECONDS_PER_DAY = 86400 +const DAILY_CAP_TTL_SECONDS = 2 * SECONDS_PER_DAY + export interface SandboxEmailSenderConfig { enabled: boolean tenantName: string configurationSetName: string fromAddress: string + dailyTeamCap: string + dailyRecipientCap: string } const ROUTE_SETTING_NAMES = { @@ -27,14 +35,57 @@ function emptyRouteSettings(config: SandboxEmailSenderConfig): string[] { .map((setting) => ROUTE_SETTING_NAMES[setting]) } +export type SandboxDailyCapGrant = { + type: 'granted' + claimId: string + buckets: ClaimRequest[] + organization: TeamOrganizationSnapshot +} + +export type SandboxDailyCapClaim = + | SandboxDailyCapGrant + | { type: 'project_cap_reached' } + | { type: 'recipient_cap_reached'; addresses: string[] } + | { type: 'check_failed' } + type SandboxEmailOutcome = | { type: 'sent'; recipientCount: number } | { type: 'blocked' - reason: 'switch_off' | 'recipient_not_member' | 'recipient_unverified' | 'check_failed' + reason: 'switch_off' | 'recipient_not_member' | 'recipient_unverified' | 'check_failed' | 'cap_reached' blockedRecipientCount: number } +function dailyBucket(key: string, requested: number, capacity: number): ClaimRequest { + return { key, requested, capacity, refillPerSecond: capacity / SECONDS_PER_DAY, ttlSeconds: DAILY_CAP_TTL_SECONDS } +} + +function parseDailyCap(value: string): number | null { + const cap = /^\d+$/.test(value) ? Number(value) : null + return cap !== null && Number.isSafeInteger(cap) && cap > 0 ? cap : null +} + +// The organization hash tag keeps shared recipient and project claims in one Valkey cluster slot. +function dailyCapKeyPrefix(organizationId: string): string { + return `@posthog/workflows-sandbox-daily/{${organizationId}}` +} + +function addressDigest(address: string): string { + return createHash('sha256').update(address.toLowerCase()).digest('hex') +} + +function distinctAddresses(recipients: string[]): string[] { + const addresses = new Map() + for (const recipient of recipients) { + const trimmed = recipient.trim() + const normalized = trimmed.toLowerCase() + if (!addresses.has(normalized)) { + addresses.set(normalized, trimmed) + } + } + return [...addresses.values()] +} + function htmlBody(document: DefaultTreeAdapterMap['document']): DefaultTreeAdapterMap['element'] | undefined { const root = document.childNodes.find(defaultTreeAdapter.isElementNode) return root?.childNodes.find( @@ -122,6 +173,7 @@ export class SandboxEmailSender { constructor( public readonly config: SandboxEmailSenderConfig, private teamManager: TeamManager, + private dailyCapLimiter: RateLimiterService | null = null, private blockedEventRedis: Pick | null = null ) { this.config = { @@ -139,6 +191,45 @@ export class SandboxEmailSender { this.canSend = config.enabled && !emptySettings.length } + public async claimDailyCaps( + teamId: number, + recipients: string[], + organization: TeamOrganizationSnapshot + ): Promise { + const dailyTeamCap = parseDailyCap(this.config.dailyTeamCap) + const dailyRecipientCap = parseDailyCap(this.config.dailyRecipientCap) + if (!this.dailyCapLimiter || dailyTeamCap === null || dailyRecipientCap === null) { + logger.error('Sandbox email daily caps are not configured correctly', { teamId }) + return { type: 'check_failed' } + } + if (Date.now() >= organization.expiresAt) { + return { type: 'check_failed' } + } + const prefix = dailyCapKeyPrefix(organization.organizationId) + const addresses = distinctAddresses(recipients) + const buckets = [ + dailyBucket(`${prefix}/project/${teamId}`, recipients.length, dailyTeamCap), + ...addresses.map((address) => + dailyBucket(`${prefix}/recipient/${addressDigest(address)}`, 1, dailyRecipientCap) + ), + ] + const claim = await this.dailyCapLimiter.claimAllOrNothing(buckets) + if (claim.granted) { + return { type: 'granted', claimId: randomUUID(), buckets, organization } + } + if (claim.deniedIndexes === null) { + return { type: 'check_failed' } + } + if (claim.deniedIndexes.includes(0)) { + return { type: 'project_cap_reached' } + } + return { type: 'recipient_cap_reached', addresses: claim.deniedIndexes.map((index) => addresses[index - 1]) } + } + + public async returnDailyCaps(grant: SandboxDailyCapGrant): Promise { + await this.dailyCapLimiter?.returnClaim(grant.claimId, grant.buckets) + } + public async withIdentificationFooter( params: CyclotronInvocationQueueParametersEmailType, senderName: string, diff --git a/nodejs/src/cdp/services/rate-limiter/rate-limiter.service.test.ts b/nodejs/src/cdp/services/rate-limiter/rate-limiter.service.test.ts index 490f0fee49ef..1b9ef5e4d58d 100644 --- a/nodejs/src/cdp/services/rate-limiter/rate-limiter.service.test.ts +++ b/nodejs/src/cdp/services/rate-limiter/rate-limiter.service.test.ts @@ -2,10 +2,13 @@ import { randomUUID } from 'crypto' import { register } from 'prom-client' import { deleteKeysWithPrefix } from '~/common/redis/_tests/redis' -import { RedisV2, createRedisV2PoolFromConfig } from '~/common/redis/redis-v2' +import { RedisClient, RedisClientPipeline, RedisV2 } from '~/common/redis/redis-v2' import { closeHub, createHub } from '~/common/utils/db/hub' -import { Hub } from '~/types' +import { createRedisPoolFromConfig } from '~/common/utils/db/redis' +import { TeamManager } from '~/common/utils/team-manager' +import { Hub, RedisPool } from '~/types' +import { SandboxEmailSender } from '../messaging/sandbox-email-sender' import { RateLimiterService } from './rate-limiter.service' const KEY = `@posthog-test/ses-rate-limiter/${randomUUID()}/bucket` @@ -15,13 +18,14 @@ describe('RateLimiterService', () => { let hub: Hub let redis: RedisV2 + let redisPool: RedisPool let limiter: RateLimiterService beforeEach(async () => { hub = await createHub() // Reuse the local CDP Redis for tests — same wire protocol, same atomicity. // In production this points at the dedicated SES rate-limiter Valkey. - redis = createRedisV2PoolFromConfig({ + redisPool = createRedisPoolFromConfig({ connection: hub.CDP_REDIS_HOST ? { url: hub.CDP_REDIS_HOST, @@ -31,11 +35,29 @@ describe('RateLimiterService', () => { poolMinSize: hub.REDIS_POOL_MIN_SIZE, poolMaxSize: hub.REDIS_POOL_MAX_SIZE, }) + redis = { + useClient: async (_options, callback) => { + const client = await redisPool.acquire() + try { + return await callback(client as RedisClient) + } finally { + await redisPool.release(client) + } + }, + usePipeline: async (options, callback) => + redis.useClient(options, async (client) => { + const pipeline = client.pipeline() as RedisClientPipeline + callback(pipeline) + return pipeline.exec() + }), + } limiter = new RateLimiterService(redis, { name: 'ses-rate-limiter' }) await deleteKeysWithPrefix(redis, KEY) }) afterEach(async () => { + await redisPool.drain() + await redisPool.clear() await closeHub(hub) }) @@ -413,4 +435,103 @@ describe('RateLimiterService', () => { expect([first.reserved, second.reserved, third.reserved]).toEqual([true, false, false]) }) }) + + it('shares recipient allowances across projects in an organization while keeping project allowances separate', async () => { + const organizations = { 1: randomUUID(), 3: randomUUID() } + const manager = { + getTeam: jest.fn((id: number) => + Promise.resolve({ organization_id: id === 3 ? organizations[3] : organizations[1] }) + ), + } as unknown as TeamManager + const sender = new SandboxEmailSender( + { + enabled: true, + tenantName: 'sandbox', + configurationSetName: 'sandbox', + fromAddress: 'sender@example.com', + dailyTeamCap: '1', + dailyRecipientCap: '1', + }, + manager, + limiter + ) + + try { + const firstOrganization = { organizationId: organizations[1], expiresAt: Infinity } + const otherOrganization = { organizationId: organizations[3], expiresAt: Infinity } + expect(await sender.claimDailyCaps(1, ['member@example.com'], firstOrganization)).toMatchObject({ + type: 'granted', + }) + expect(await sender.claimDailyCaps(2, ['member@example.com'], firstOrganization)).toEqual({ + type: 'recipient_cap_reached', + addresses: ['member@example.com'], + }) + expect(await sender.claimDailyCaps(2, ['other@example.com'], firstOrganization)).toMatchObject({ + type: 'granted', + }) + expect(await sender.claimDailyCaps(3, ['member@example.com'], otherOrganization)).toMatchObject({ + type: 'granted', + }) + expect(await sender.claimDailyCaps(2, ['third@example.com'], firstOrganization)).toEqual({ + type: 'project_cap_reached', + }) + } finally { + for (const organizationId of Object.values(organizations)) { + await deleteKeysWithPrefix(redis, `@posthog/workflows-sandbox-daily/{${organizationId}}`) + } + } + }) + + it('spends an unanswered claim once when the same EVAL is replayed', async () => { + const bucket = { key: `${KEY}/claim-replay`, requested: 1, capacity: 2, refillPerSecond: 0 } + const replayingRedis: RedisV2 = { + ...redis, + useClient: async (options, callback) => + redis.useClient(options, async (client) => { + const evaluate = client.eval.bind(client) + const replay = jest.spyOn(client, 'eval').mockImplementation(async (...args) => { + await evaluate(...args) + return evaluate(...args) + }) + try { + return await callback(client) + } finally { + replay.mockRestore() + } + }), + } + const replayingLimiter = new RateLimiterService(replayingRedis, { name: 'ses-rate-limiter' }) + + expect(await replayingLimiter.claimAllOrNothing([bucket])).toEqual({ granted: true }) + expect(await limiter.claimAllOrNothing([bucket])).toEqual({ granted: true }) + expect(await limiter.claimAllOrNothing([bucket])).toEqual({ granted: false, deniedIndexes: [0] }) + }) + + describe('returnClaim', () => { + it('gives a claim back once, even when the return is replayed', async () => { + const bucket = { key: `${KEY}/return`, requested: 1, capacity: 2, refillPerSecond: 0 } + expect(await limiter.claimAllOrNothing([bucket])).toEqual({ granted: true }) + expect(await limiter.claimAllOrNothing([bucket])).toEqual({ granted: true }) + + await limiter.returnClaim('claim-1', [bucket]) + await limiter.returnClaim('claim-1', [bucket]) + + expect(await limiter.claimAllOrNothing([bucket])).toEqual({ granted: true }) + expect(await limiter.claimAllOrNothing([bucket])).toEqual({ granted: false, deniedIndexes: [0] }) + }) + + it('ignores a return that reaches Valkey after its window', async () => { + const bucket = { key: `${KEY}/late-return`, requested: 1, capacity: 1, refillPerSecond: 0 } + expect(await limiter.claimAllOrNothing([bucket])).toEqual({ granted: true }) + + const issuedLongAgo = jest.spyOn(Date, 'now').mockReturnValue(Date.now() - 60 * 60 * 1000) + try { + await limiter.returnClaim('claim-1', [bucket]) + } finally { + issuedLongAgo.mockRestore() + } + + expect(await limiter.claimAllOrNothing([bucket])).toEqual({ granted: false, deniedIndexes: [0] }) + }) + }) }) diff --git a/nodejs/src/cdp/services/rate-limiter/rate-limiter.service.ts b/nodejs/src/cdp/services/rate-limiter/rate-limiter.service.ts index 1d4b178a0574..3e09da89b2cb 100644 --- a/nodejs/src/cdp/services/rate-limiter/rate-limiter.service.ts +++ b/nodejs/src/cdp/services/rate-limiter/rate-limiter.service.ts @@ -1,3 +1,4 @@ +import { randomUUID } from 'crypto' import { Counter, Histogram } from 'prom-client' import { RedisV2 } from '~/common/redis/redis-v2' @@ -249,6 +250,77 @@ end return {1, 0, 0, 0} ` +const CLAIM_ALL_OR_NOTHING_LUA = ` +local time = redis.call('TIME') +local now = tonumber(time[1]) * 1000 + math.floor(tonumber(time[2]) / 1000) +local claimKey = KEYS[#KEYS] +local previous = redis.call('GET', claimKey) +if previous then + return cjson.decode(previous) +end +if now > tonumber(ARGV[#ARGV - 1]) then + return {0} +end +local replayTtlMs = tonumber(ARGV[#ARGV]) +local available = {} +local denied = {0} +for i = 1, #KEYS - 1 do + local offset = (i - 1) * 4 + local requested = tonumber(ARGV[offset + 1]) + local capacity = tonumber(ARGV[offset + 2]) + local refillPerSecond = tonumber(ARGV[offset + 3]) + local existing = redis.call('hmget', KEYS[i], 'ts', 'pool') + local avail + if existing[1] == false then + avail = capacity + else + local elapsedMs = math.max(0, now - tonumber(existing[1])) + local currentTokens = capacity + if existing[2] ~= false then + currentTokens = tonumber(existing[2]) + end + avail = math.min(capacity, currentTokens + (elapsedMs / 1000.0) * refillPerSecond) + end + available[i] = avail + if avail < requested then + table.insert(denied, i) + end +end +if #denied > 1 then + redis.call('SET', claimKey, cjson.encode(denied), 'NX', 'PX', replayTtlMs) + return denied +end +for i = 1, #KEYS - 1 do + local offset = (i - 1) * 4 + redis.call('hset', KEYS[i], 'ts', now, 'pool', available[i] - tonumber(ARGV[offset + 1])) + redis.call('expire', KEYS[i], tonumber(ARGV[offset + 4])) +end +local granted = {1} +redis.call('SET', claimKey, cjson.encode(granted), 'NX', 'PX', replayTtlMs) +return granted +` + +// ioredis can replay unanswered commands, so returns need deduplication and an expiry. +const RETURN_CLAIM_LUA = ` +local time = redis.call('TIME') +local now = tonumber(time[1]) * 1000 + math.floor(tonumber(time[2]) / 1000) +local buckets = #KEYS - 1 +if now > tonumber(ARGV[buckets + 1]) then + return 0 +end +if not redis.call('set', KEYS[#KEYS], 1, 'NX', 'PX', ARGV[buckets + 2]) then + return 0 +end +for i = 1, buckets do + if redis.call('hexists', KEYS[i], 'pool') == 1 then + redis.call('hincrbyfloat', KEYS[i], 'pool', ARGV[i]) + end +end +return 1 +` + +const RETURN_WINDOW_MS = 5 * 60 * 1000 + export interface RateLimiterConfig { /** Logical name for metrics/logging only (e.g. 'ses'). */ name: string @@ -373,6 +445,61 @@ export class RateLimiterService { } } + // All bucket keys must share a hash tag because clustered Valkey rejects cross-slot Lua calls. + public async claimAllOrNothing( + buckets: ClaimRequest[] + ): Promise<{ granted: true } | { granted: false; deniedIndexes: number[] | null }> { + const endTimer = claimLatency.startTimer({ limiter: this.config.name }) + const keys = buckets.map((bucket) => bucket.key) + const claimId = randomUUID() + const deadline = Date.now() + RETURN_WINDOW_MS + try { + const result = await this.valkey.useClient( + { name: `rate-limiter:${this.config.name}:claimAllOrNothing`, timeout: 1000 }, + (client) => + client.eval( + CLAIM_ALL_OR_NOTHING_LUA, + buckets.length + 1, + ...keys, + `${keys[0]}/claimed/${claimId}`, + ...buckets.flatMap((bucket) => + [bucket.requested, bucket.capacity, bucket.refillPerSecond, bucket.ttlSeconds ?? 3600].map( + String + ) + ), + String(deadline), + String(2 * RETURN_WINDOW_MS) + ) + ) + if (Array.isArray(result) && result.length === 1 && result[0] === 1) { + claimCounter.inc({ limiter: this.config.name, result: 'granted_full' }) + return { granted: true } + } + if ( + Array.isArray(result) && + result[0] === 0 && + result.length > 1 && + result.slice(1).every((index) => Number.isInteger(index) && index >= 1 && index <= buckets.length) && + new Set(result.slice(1)).size === result.length - 1 + ) { + claimCounter.inc({ limiter: this.config.name, result: 'denied' }) + return { granted: false, deniedIndexes: result.slice(1).map((index) => index - 1) } + } + logger.warn('🪙', `RateLimiterService(${this.config.name}) claim returned invalid result`, { + keys, + raw: result, + }) + claimCounter.inc({ limiter: this.config.name, result: 'valkey_error' }) + return { granted: false, deniedIndexes: null } + } catch (err) { + logger.warn('🪙', `RateLimiterService(${this.config.name}) claim threw`, { keys, error: String(err) }) + claimCounter.inc({ limiter: this.config.name, result: 'valkey_error' }) + return { granted: false, deniedIndexes: null } + } finally { + endTimer() + } + } + /** * Atomically claim `requested` tokens from BOTH buckets, or neither. A denial consumes no * tokens, so a caller that retries a multi-token claim cannot drain the buckets while never @@ -449,4 +576,28 @@ export class RateLimiterService { endTimer() } } + + public async returnClaim(claimId: string, buckets: Pick[]): Promise { + const keys = buckets.map((bucket) => bucket.key) + try { + await this.valkey.useClient( + { name: `rate-limiter:${this.config.name}:returnClaim`, timeout: 1000 }, + (client) => + client.eval( + RETURN_CLAIM_LUA, + buckets.length + 1, + ...keys, + `${keys[0]}/returned/${claimId}`, + ...buckets.map((bucket) => String(bucket.requested)), + String(Date.now() + RETURN_WINDOW_MS), + String(2 * RETURN_WINDOW_MS) + ) + ) + } catch (err) { + logger.warn('🪙', `RateLimiterService(${this.config.name}) claim return threw`, { + keys, + error: String(err), + }) + } + } } diff --git a/nodejs/src/server.ts b/nodejs/src/server.ts index 81391973e04b..a95a45ab3c39 100644 --- a/nodejs/src/server.ts +++ b/nodejs/src/server.ts @@ -137,8 +137,9 @@ export class PluginServer implements NodeServer { // Build typed deps objects for consumers. `emailValidationValkey` and `frequencyCapValkey` // are null in the shared deps and set only by the cyclotron workers that run the hogflow - // email action (see `withEmailActionValkeys` at their loaders below) — no - // other consumer touches the SES or frequency cap Valkey. + // email action (see `withEmailActionValkeys` at their loaders below). Sandbox-enabled + // cdp-api gets only the SES pool, because its test sends skip the frequency cap. + // No other consumer touches the SES or frequency cap Valkey. const cdpDeps: CdpConsumerBaseDeps | undefined = needsCdp ? { postgres: this.postgres!, @@ -245,7 +246,18 @@ export class PluginServer implements NodeServer { : null const api = new CdpApi( this.config, - { ...cdpDeps!, idempotencyValkey: createIdempotencyValkeyPool(this.config) }, + { + ...cdpDeps!, + ...(this.config.WORKFLOWS_SANDBOX_SENDER_ENABLED + ? { + emailValidationValkey: createSesRateLimiterValkeyPool( + this.config, + 'email-mx-validation' + ), + } + : {}), + idempotencyValkey: createIdempotencyValkeyPool(this.config), + }, { hogQueue: kafkaQueue, hogflowQueue: postgresV2Queue }, batchResolverProducer ) @@ -456,9 +468,8 @@ export class PluginServer implements NodeServer { /** * Grants the SES Valkey pool that backs the shared MX-verdict cache to a worker's - * deps. Only the cyclotron workers that run the hogflow email action call this, so - * the pool is never opened on other CDP consumers or cdp-api — an idle Valkey sized - * for the SES rate limiter shouldn't hold connections from pods that never validate. + * deps. Only the email workers call this. Sandbox-enabled cdp-api opens the SES pool on its + * own, without the frequency cap pool, and the pool stays closed on other CDP consumers. * Null when no SES Valkey host is configured (local dev), in which case * EmailValidationService degrades to its local cache + DNS. *