Skip to content

Commit 77dfe17

Browse files
authored
fix(billing): serialize Gmail inbox allowance decisions so racing alias signups can't both keep it (#9431)
* fix(billing): serialize Gmail inbox allowance decisions so racing alias signups can't both keep it * fix(billing): restore lock_timeout after taking the Gmail inbox lock
1 parent 336ed5a commit 77dfe17

8 files changed

Lines changed: 227 additions & 85 deletions

File tree

‎apps/sim/lib/admin/dashboard-credit-grant.test.ts‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ const hoistedMocks = vi.hoisted(() => ({
2121
billingSubscriptions: [] as unknown[],
2222
idempotencyCalls: [] as { namespace: string; requestFingerprint: string }[],
2323
moveWorkspace: vi.fn(),
24-
userHasDeletedAccountIdentity: vi.fn(async () => false),
24+
identityHadFreeAllowance: vi.fn(async () => false),
2525
}))
2626

2727
vi.mock('@sim/audit', () => auditMock)
@@ -54,8 +54,8 @@ vi.mock('@/lib/billing/enterprise-outbox', () => ({
5454
resolveEnterpriseMetadataIntent: vi.fn(),
5555
}))
5656
vi.mock('@/lib/core/outbox/service', () => outboxServiceMock)
57-
vi.mock('@/lib/users/deleted-identity', () => ({
58-
userHasDeletedAccountIdentity: hoistedMocks.userHasDeletedAccountIdentity,
57+
vi.mock('@/lib/users/free-allowance', () => ({
58+
identityHadFreeAllowance: hoistedMocks.identityHadFreeAllowance,
5959
}))
6060

6161
import { grantDashboardOrganizationBalance, grantDashboardUserBalance } from '@/lib/admin/dashboard'
@@ -154,8 +154,8 @@ describe('grantDashboardUserBalance', () => {
154154
)
155155
})
156156

157-
it('creates a missing stats row without the free allowance for a returning deleted identity', async () => {
158-
mocks.userHasDeletedAccountIdentity.mockResolvedValueOnce(true)
157+
it('creates a missing stats row without the free allowance for an identity that already had it', async () => {
158+
mocks.identityHadFreeAllowance.mockResolvedValueOnce(true)
159159
queueTableRows(user, [{ id: 'user-1' }])
160160
queueTableRows(userStats, [
161161
{ creditBalance: '0', currentUsageLimit: '0', freeTierEligible: false },

‎apps/sim/lib/admin/dashboard.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ import { executeTransactionallyIdempotent } from '@/lib/core/idempotency/transac
7878
import { enqueueOutboxEvent } from '@/lib/core/outbox/service'
7979
import type { DbOrTx } from '@/lib/db/types'
8080
import { deliverRunStops } from '@/lib/mothership/request/application/stop-delivery'
81-
import { userHasDeletedAccountIdentity } from '@/lib/users/deleted-identity'
81+
import { identityHadFreeAllowance } from '@/lib/users/free-allowance'
8282
import { ownedAttachableWorkspacesWhere } from '@/lib/workspaces/organization-workspaces'
8383

8484
interface PaginationInput {
@@ -1503,7 +1503,7 @@ export async function grantDashboardUserBalance(
15031503
* the same email check as every other row creation. Looked up before the
15041504
* transaction so it never holds a second pooled connection inside it.
15051505
*/
1506-
const freeTierEligibleIfCreated = !(await userHasDeletedAccountIdentity(userId, {
1506+
const freeTierEligibleIfCreated = !(await identityHadFreeAllowance(userId, {
15071507
includeSignIns: false,
15081508
}))
15091509
const outcome = await db.transaction(async (tx) => {

‎apps/sim/lib/billing/core/usage.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,8 @@ vi.mock('@/lib/users/deleted-identity', () => ({
6060
}))
6161

6262
vi.mock('@/lib/users/gmail-alias', () => ({
63+
acquireGmailInboxLock: vi.fn(async () => {}),
64+
laterGmailAliasesWithAllowance: vi.fn(async () => []),
6365
userHasGmailAliasAccount: vi.fn(async () => false),
6466
}))
6567

‎apps/sim/lib/billing/core/usage.ts‎

Lines changed: 30 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -46,11 +46,9 @@ import { isBillingEnabled } from '@/lib/core/config/env-flags'
4646
import { getBaseUrl } from '@/lib/core/utils/urls'
4747
import type { DbClient, DbOrTx } from '@/lib/db/types'
4848
import { APP_ENTRY_PATH } from '@/lib/navigation/paths'
49-
import {
50-
isDeletedAccountIdentity,
51-
userHasDeletedAccountIdentity,
52-
} from '@/lib/users/deleted-identity'
53-
import { userHasGmailAliasAccount } from '@/lib/users/gmail-alias'
49+
import { isDeletedAccountIdentity } from '@/lib/users/deleted-identity'
50+
import { identityHadFreeAllowance } from '@/lib/users/free-allowance'
51+
import { acquireGmailInboxLock, laterGmailAliasesWithAllowance } from '@/lib/users/gmail-alias'
5452

5553
const logger = createLogger('UsageManagement')
5654

@@ -192,22 +190,6 @@ function calculateOrgUsageLimit(
192190
*/
193191
const usersCreatedInAuthRequest = new WeakMap<object, string>()
194192

195-
/**
196-
* Whether a new user's identity already had the free allowance: its email or a
197-
* linked sign-in (with `includeSignIns`) belonged to a recently deleted account,
198-
* or another account already holds its Gmail inbox.
199-
*/
200-
async function identityHadFreeAllowance(
201-
userId: string,
202-
{ includeSignIns }: { includeSignIns: boolean },
203-
executor: DbOrTx = db
204-
): Promise<boolean> {
205-
return (
206-
(await userHasDeletedAccountIdentity(userId, { includeSignIns }, executor)) ||
207-
(await userHasGmailAliasAccount(userId, executor))
208-
)
209-
}
210-
211193
function isUserCreatedInAuthRequest(
212194
userId: string,
213195
authRequest: object | null | undefined
@@ -219,27 +201,38 @@ function isUserCreatedInAuthRequest(
219201
* Handle new user setup when they join the platform.
220202
* Creates the userStats record with the free allowance, which is zero when the
221203
* user's email or an already-linked social sign-in belongs to an account
222-
* deleted within the retention window, or when another account holds its Gmail
223-
* inbox. Every user Better Auth creates comes through here, including SCIM- and
204+
* deleted within the retention window, or when an earlier account holds its
205+
* Gmail inbox; it decides under that inbox's lock (see `lib/users/gmail-alias`).
206+
* Every user Better Auth creates comes through here, including SCIM- and
224207
* SSO-provisioned users, so their email is always checked at creation.
225208
*
226209
* @param authRequest Better Auth's endpoint context for the signup request, when there is one
227210
*/
228211
export async function handleNewUser(userId: string, authRequest?: object | null): Promise<void> {
229212
if (authRequest) usersCreatedInAuthRequest.set(authRequest, userId)
230213
try {
231-
const freeTierEligible = !(await identityHadFreeAllowance(userId, { includeSignIns: true }))
232-
if (!freeTierEligible) {
233-
logger.info('Withholding the free allowance from an identity that already had it', {
214+
await db.transaction(async (tx) => {
215+
await acquireGmailInboxLock(tx, userId)
216+
const freeTierEligible = !(await identityHadFreeAllowance(
234217
userId,
218+
{ includeSignIns: true },
219+
tx
220+
))
221+
if (!freeTierEligible) {
222+
logger.info('Withholding the free allowance from an identity that already had it', {
223+
userId,
224+
})
225+
}
226+
await tx.insert(userStats).values({
227+
id: generateId(),
228+
userId: userId,
229+
currentUsageLimit: getFreeTierAllowance(freeTierEligible).toString(),
230+
freeTierEligible,
231+
usageLimitUpdatedAt: new Date(),
235232
})
236-
}
237-
await db.insert(userStats).values({
238-
id: generateId(),
239-
userId: userId,
240-
currentUsageLimit: getFreeTierAllowance(freeTierEligible).toString(),
241-
freeTierEligible,
242-
usageLimitUpdatedAt: new Date(),
233+
for (const aliasId of await laterGmailAliasesWithAllowance(userId, tx)) {
234+
await forfeitFreeTierAllowance(tx, aliasId)
235+
}
243236
})
244237

245238
logger.info('User stats record created for new user', { userId })
@@ -305,10 +298,10 @@ export async function ensureUserStatsExists(
305298

306299
/**
307300
* Creates a user's `user_stats` row with the free allowance when none exists, on `executor`
308-
* (the caller's transaction, or `db`). Only a row this call creates gets the deleted-identity
309-
* check signup makes: the email always, and the linked sign-ins only when `includeSignIns` is
310-
* set. An existing row costs only the conflicting insert. A failed check leaves the allowance in
311-
* place and never fails the caller or aborts its transaction.
301+
* (the caller's transaction, or `db`). Only a row this call creates gets the check signup makes:
302+
* the email always, and the linked sign-ins only when `includeSignIns` is set. An existing row
303+
* costs only the conflicting insert. A failed check leaves the allowance in place and never fails
304+
* the caller or aborts its transaction.
312305
*
313306
* @returns whether this call created the row
314307
*/
Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
import { db } from '@sim/db'
2+
import type { DbOrTx } from '@/lib/db/types'
3+
import { userHasDeletedAccountIdentity } from '@/lib/users/deleted-identity'
4+
import { userHasGmailAliasAccount } from '@/lib/users/gmail-alias'
5+
6+
/**
7+
* Whether a new user's identity already had the free allowance: its email or a
8+
* linked sign-in (with `includeSignIns`) belonged to a recently deleted account,
9+
* or an earlier account holds its Gmail inbox. Signup, the stats-row fallback,
10+
* and the admin balance grant decide a new row's allowance with this.
11+
*/
12+
export async function identityHadFreeAllowance(
13+
userId: string,
14+
{ includeSignIns }: { includeSignIns: boolean },
15+
executor: DbOrTx = db
16+
): Promise<boolean> {
17+
return (
18+
(await userHasDeletedAccountIdentity(userId, { includeSignIns }, executor)) ||
19+
(await userHasGmailAliasAccount(userId, executor))
20+
)
21+
}

‎apps/sim/lib/users/gmail-alias.integration.ts‎

Lines changed: 83 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,21 @@
11
/**
22
* A dot, `+tag`, or `googlemail.com` variant of an existing Gmail address must not get a second
3-
* free allowance. Runs against real PostgreSQL through the production signup-stats and
4-
* fallback-stats paths, and checks the lookup is served by `user_gmail_mailbox_idx`.
3+
* free allowance, even when two signups on one inbox race. Runs against real PostgreSQL through
4+
* the production signup-stats and fallback-stats paths, and checks the lookup is served by
5+
* `user_gmail_mailbox_idx`.
56
*/
67
import { db } from '@sim/db'
78
import { gmailMailbox, user, userStats } from '@sim/db/schema'
9+
import { createDeferred } from '@sim/testing'
810
import { authMock } from '@sim/testing/mocks/auth.mock'
911
import { envFlagsMock } from '@sim/testing/mocks/env-flags.mock'
1012
import { generateId } from '@sim/utils/id'
1113
import { eq, inArray, sql } from 'drizzle-orm'
14+
import { alias } from 'drizzle-orm/pg-core'
1215
import { afterAll, describe, expect, it, vi } from 'vitest'
1316
import { ensureUserStatsExists, handleNewUser } from '@/lib/billing/core/usage'
1417
import { getFreeTierLimit } from '@/lib/billing/subscriptions/utils'
18+
import { acquireGmailInboxLock } from '@/lib/users/gmail-alias'
1519

1620
vi.mock('@/lib/core/config/env-flags', () => ({
1721
...envFlagsMock,
@@ -60,6 +64,11 @@ async function statsOf(userId: string) {
6064
return { freeTierEligible: row?.freeTierEligible, limit: Number(row?.currentUsageLimit) }
6165
}
6266

67+
/** An account created a minute ago, so any account inserted after it counts as later. */
68+
function insertEarlierUser(email: string): Promise<string> {
69+
return insertUser(email, { createdAt: new Date(Date.now() - 60_000) })
70+
}
71+
6372
/** A Gmail local part no other test run can collide with. */
6473
function uniqueLocalPart(): string {
6574
return `alias${generateId().replaceAll('-', '')}`
@@ -127,18 +136,16 @@ describe('Gmail inboxes and the free allowance', () => {
127136
it('withholds it from a stats row the fallback creates after signup failed to', async () => {
128137
const local = uniqueLocalPart()
129138
await signUp(`${local}@gmail.com`)
130-
const alias = await insertUser(`${local.slice(0, 3)}.${local.slice(3)}@gmail.com`)
139+
const later = await insertUser(`${local.slice(0, 3)}.${local.slice(3)}@gmail.com`)
131140

132-
await ensureUserStatsExists(alias)
141+
await ensureUserStatsExists(later)
133142

134-
expect(await statsOf(alias)).toEqual({ freeTierEligible: false, limit: 0 })
143+
expect(await statsOf(later)).toEqual({ freeTierEligible: false, limit: 0 })
135144
})
136145

137146
it('keeps the allowance with the inbox’s first account when its stats row comes after an alias’s', async () => {
138147
const local = uniqueLocalPart()
139-
const original = await insertUser(`${local}@gmail.com`, {
140-
createdAt: new Date(Date.now() - 60_000),
141-
})
148+
const original = await insertEarlierUser(`${local}@gmail.com`)
142149
const alias = await signUp(`${local}+later@gmail.com`)
143150

144151
await ensureUserStatsExists(original)
@@ -164,11 +171,78 @@ describe('Gmail inboxes and the free allowance', () => {
164171
expect(eligible.sort()).toEqual([false, true])
165172
})
166173

174+
it('withdraws it from an alias whose signup decided before the earlier account was visible', async () => {
175+
const local = uniqueLocalPart()
176+
const later = await signUp(`${local}+fast@gmail.com`)
177+
const original = await insertEarlierUser(`${local}@gmail.com`)
178+
179+
await handleNewUser(original, {})
180+
181+
expect(await statsOf(original)).toEqual({ freeTierEligible: true, limit: getFreeTierLimit() })
182+
expect(await statsOf(later)).toEqual({ freeTierEligible: false, limit: 0 })
183+
})
184+
185+
it('never withdraws it from a later alias when the fallback creates the earlier account’s row', async () => {
186+
const local = uniqueLocalPart()
187+
const original = await insertEarlierUser(`${local}@gmail.com`)
188+
const later = await signUp(`${local}+later@gmail.com`)
189+
await db.update(userStats).set({ freeTierEligible: true }).where(eq(userStats.userId, later))
190+
191+
await ensureUserStatsExists(original)
192+
193+
expect(await statsOf(original)).toEqual({ freeTierEligible: true, limit: getFreeTierLimit() })
194+
expect((await statsOf(later)).freeTierEligible).toBe(true)
195+
})
196+
197+
it('makes a racing signup on the same inbox wait for the inbox lock', async () => {
198+
const local = uniqueLocalPart()
199+
const original = await insertEarlierUser(`${local}@gmail.com`)
200+
const racer = await insertUser(`${local.slice(0, 4)}.${local.slice(4)}@gmail.com`)
201+
const held = createDeferred<number>()
202+
const release = createDeferred<void>()
203+
const pending: Promise<unknown>[] = []
204+
try {
205+
const holder = db.transaction(async (tx) => {
206+
await acquireGmailInboxLock(tx, original)
207+
const [connection] = await tx.execute<{ pid: number }>(sql`SELECT pg_backend_pid() AS pid`)
208+
held.resolve(connection.pid)
209+
await release.promise
210+
})
211+
pending.push(holder)
212+
void holder.catch((error: unknown) => held.reject(error))
213+
const blockerPid = await held.promise
214+
const signup = handleNewUser(racer, {})
215+
pending.push(signup)
216+
217+
await vi.waitFor(
218+
async () => {
219+
const [waiting] = await db.execute<{ pid: number }>(sql`
220+
SELECT pid FROM pg_stat_activity
221+
WHERE datname = current_database()
222+
AND ${blockerPid}::int = ANY(pg_blocking_pids(pid))
223+
`)
224+
expect(waiting).toBeDefined()
225+
},
226+
{ timeout: 5_000, interval: 25 }
227+
)
228+
expect((await statsOf(racer)).freeTierEligible).toBeUndefined()
229+
230+
release.resolve()
231+
await holder
232+
await signup
233+
expect(await statsOf(racer)).toEqual({ freeTierEligible: false, limit: 0 })
234+
} finally {
235+
release.resolve()
236+
await Promise.allSettled(pending)
237+
}
238+
})
239+
167240
it('finds the inbox through user_gmail_mailbox_idx rather than a scan of user', async () => {
241+
const subject = alias(user, 'subject')
168242
const plan = await db.transaction(async (tx) => {
169243
await tx.execute(sql`SET LOCAL enable_seqscan = off`)
170244
const rows = await tx.execute<{ 'QUERY PLAN': string }>(
171-
sql`EXPLAIN SELECT ${user.id} FROM ${user} WHERE ${gmailMailbox(user.email)} = ${`${uniqueLocalPart()}@gmail.com`} LIMIT 1`
245+
sql`EXPLAIN SELECT ${user.id} FROM ${user} AS "subject" JOIN ${user} ON ${gmailMailbox(user.email)} = ${gmailMailbox(subject.email)} WHERE ${subject.id} = ${generateId()}`
172246
)
173247
return rows.map((row) => row['QUERY PLAN']).join('\n')
174248
})

0 commit comments

Comments
 (0)