Skip to content

Commit cb8c936

Browse files
brandonkachenbrandonkachen
committed
feat(billing): use advisory locks for credit operations
Replace SERIALIZABLE transactions with advisory locks to serialize credit operations per user/org. This eliminates serialization failures by making concurrent transactions wait instead of failing and retrying. Changes: - Add withAdvisoryLockTransaction helper using pg_advisory_xact_lock - Update consumeCredits, consumeCreditsAndAddAgentStep to use advisory locks - Update grantCreditOperation, revokeGrantByOperationId to use advisory locks - Update triggerMonthlyResetAndGrant to use advisory locks - Update consumeOrganizationCredits, grantOrganizationCredits to use advisory locks - Add analytics events for lock contention and retry threshold exceeded - Add comprehensive tests for advisory lock behavior
1 parent 4f870f6 commit cb8c936

File tree

9 files changed

+2400
-304
lines changed

9 files changed

+2400
-304
lines changed

‎common/src/constants/analytics-events.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,10 @@ export enum AnalyticsEvent {
2626
UNKNOWN_TOOL_CALL = 'backend.unknown_tool_call',
2727
USER_INPUT = 'backend.user_input',
2828

29+
// Backend - Database Operations
30+
ADVISORY_LOCK_CONTENTION = 'backend.advisory_lock_contention',
31+
TRANSACTION_RETRY_THRESHOLD_EXCEEDED = 'backend.transaction_retry_threshold_exceeded',
32+
2933
// Web
3034
SIGNUP = 'web.signup',
3135

‎packages/billing/src/__tests__/grant-credits.test.ts‎

Lines changed: 81 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,40 @@ const logger: Logger = {
1818
const futureDate = new Date(Date.now() + 30 * 24 * 60 * 60 * 1000) // 30 days from now
1919
const pastDate = new Date(Date.now() - 30 * 24 * 60 * 60 * 1000) // 30 days ago
2020

21+
const createTxMock = (user: {
22+
next_quota_reset: Date | null
23+
auto_topup_enabled: boolean | null
24+
} | null) => ({
25+
query: {
26+
user: {
27+
findFirst: async () => user,
28+
},
29+
},
30+
update: () => ({
31+
set: () => ({
32+
where: () => Promise.resolve(),
33+
}),
34+
}),
35+
insert: () => ({
36+
values: () => ({
37+
onConflictDoNothing: () => ({
38+
returning: () => Promise.resolve([{ id: 'test-id' }]),
39+
}),
40+
}),
41+
}),
42+
select: () => ({
43+
from: () => ({
44+
where: () => ({
45+
orderBy: () => ({
46+
limit: () => [],
47+
}),
48+
}),
49+
then: (cb: any) => cb([]),
50+
}),
51+
}),
52+
execute: () => Promise.resolve([]),
53+
})
54+
2155
const createDbMock = (options: {
2256
user: {
2357
next_quota_reset: Date | null
@@ -27,34 +61,6 @@ const createDbMock = (options: {
2761
const { user } = options
2862

2963
return {
30-
transaction: async (callback: (tx: any) => Promise<any>) => {
31-
const tx = {
32-
query: {
33-
user: {
34-
findFirst: async () => user,
35-
},
36-
},
37-
update: () => ({
38-
set: () => ({
39-
where: () => Promise.resolve(),
40-
}),
41-
}),
42-
insert: () => ({
43-
values: () => Promise.resolve(),
44-
}),
45-
select: () => ({
46-
from: () => ({
47-
where: () => ({
48-
orderBy: () => ({
49-
limit: () => [],
50-
}),
51-
}),
52-
then: (cb: any) => cb([]),
53-
}),
54-
}),
55-
}
56-
return callback(tx)
57-
},
5864
select: () => ({
5965
from: () => ({
6066
where: () => ({
@@ -67,6 +73,17 @@ const createDbMock = (options: {
6773
}
6874
}
6975

76+
const createTransactionMock = (user: {
77+
next_quota_reset: Date | null
78+
auto_topup_enabled: boolean | null
79+
} | null) => ({
80+
withAdvisoryLockTransaction: async ({
81+
callback,
82+
}: {
83+
callback: (tx: any) => Promise<any>
84+
}) => await callback(createTxMock(user)),
85+
})
86+
7087
describe('grant-credits', () => {
7188
afterEach(() => {
7289
clearMockedModules()
@@ -75,14 +92,16 @@ describe('grant-credits', () => {
7592
describe('triggerMonthlyResetAndGrant', () => {
7693
describe('autoTopupEnabled return value', () => {
7794
it('should return autoTopupEnabled: true when user has auto_topup_enabled: true', async () => {
95+
const user = {
96+
next_quota_reset: futureDate,
97+
auto_topup_enabled: true,
98+
}
7899
await mockModule('@codebuff/internal/db', () => ({
79-
default: createDbMock({
80-
user: {
81-
next_quota_reset: futureDate,
82-
auto_topup_enabled: true,
83-
},
84-
}),
100+
default: createDbMock({ user }),
85101
}))
102+
await mockModule('@codebuff/internal/db/transaction', () =>
103+
createTransactionMock(user),
104+
)
86105

87106
// Need to re-import after mocking
88107
const { triggerMonthlyResetAndGrant: fn } = await import('../grant-credits')
@@ -97,14 +116,16 @@ describe('grant-credits', () => {
97116
})
98117

99118
it('should return autoTopupEnabled: false when user has auto_topup_enabled: false', async () => {
119+
const user = {
120+
next_quota_reset: futureDate,
121+
auto_topup_enabled: false,
122+
}
100123
await mockModule('@codebuff/internal/db', () => ({
101-
default: createDbMock({
102-
user: {
103-
next_quota_reset: futureDate,
104-
auto_topup_enabled: false,
105-
},
106-
}),
124+
default: createDbMock({ user }),
107125
}))
126+
await mockModule('@codebuff/internal/db/transaction', () =>
127+
createTransactionMock(user),
128+
)
108129

109130
const { triggerMonthlyResetAndGrant: fn } = await import('../grant-credits')
110131

@@ -117,14 +138,16 @@ describe('grant-credits', () => {
117138
})
118139

119140
it('should default autoTopupEnabled to false when user has auto_topup_enabled: null', async () => {
141+
const user = {
142+
next_quota_reset: futureDate,
143+
auto_topup_enabled: null,
144+
}
120145
await mockModule('@codebuff/internal/db', () => ({
121-
default: createDbMock({
122-
user: {
123-
next_quota_reset: futureDate,
124-
auto_topup_enabled: null,
125-
},
126-
}),
146+
default: createDbMock({ user }),
127147
}))
148+
await mockModule('@codebuff/internal/db/transaction', () =>
149+
createTransactionMock(user),
150+
)
128151

129152
const { triggerMonthlyResetAndGrant: fn } = await import('../grant-credits')
130153

@@ -138,10 +161,11 @@ describe('grant-credits', () => {
138161

139162
it('should throw error when user is not found', async () => {
140163
await mockModule('@codebuff/internal/db', () => ({
141-
default: createDbMock({
142-
user: null,
143-
}),
164+
default: createDbMock({ user: null }),
144165
}))
166+
await mockModule('@codebuff/internal/db/transaction', () =>
167+
createTransactionMock(null),
168+
)
145169

146170
const { triggerMonthlyResetAndGrant: fn } = await import('../grant-credits')
147171

@@ -156,14 +180,16 @@ describe('grant-credits', () => {
156180

157181
describe('quota reset behavior', () => {
158182
it('should return existing reset date when it is in the future', async () => {
183+
const user = {
184+
next_quota_reset: futureDate,
185+
auto_topup_enabled: false,
186+
}
159187
await mockModule('@codebuff/internal/db', () => ({
160-
default: createDbMock({
161-
user: {
162-
next_quota_reset: futureDate,
163-
auto_topup_enabled: false,
164-
},
165-
}),
188+
default: createDbMock({ user }),
166189
}))
190+
await mockModule('@codebuff/internal/db/transaction', () =>
191+
createTransactionMock(user),
192+
)
167193

168194
const { triggerMonthlyResetAndGrant: fn } = await import('../grant-credits')
169195

‎packages/billing/src/__tests__/org-billing.test.ts‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ const logger: Logger = {
5252

5353
const createDbMock = (options?: {
5454
grants?: typeof mockGrants | any[]
55-
insert?: () => { values: () => Promise<unknown> }
55+
insert?: () => { values: () => { onConflictDoNothing: () => { returning: () => Promise<unknown[]> } } }
5656
update?: () => { set: () => { where: () => Promise<unknown> } }
5757
}) => {
5858
const { grants = mockGrants, insert, update } = options ?? {}
@@ -68,7 +68,11 @@ const createDbMock = (options?: {
6868
insert:
6969
insert ??
7070
(() => ({
71-
values: () => Promise.resolve(),
71+
values: () => ({
72+
onConflictDoNothing: () => ({
73+
returning: () => Promise.resolve([{ id: 'test-id' }]),
74+
}),
75+
}),
7276
})),
7377
update:
7478
update ??
@@ -77,6 +81,7 @@ const createDbMock = (options?: {
7781
where: () => Promise.resolve(),
7882
}),
7983
})),
84+
execute: () => Promise.resolve([]),
8085
}
8186
}
8287

@@ -86,7 +91,7 @@ describe('Organization Billing', () => {
8691
default: createDbMock(),
8792
}))
8893
await mockModule('@codebuff/internal/db/transaction', () => ({
89-
withSerializableTransaction: async ({
94+
withAdvisoryLockTransaction: async ({
9095
callback,
9196
}: {
9297
callback: (tx: any) => Promise<unknown> | unknown
@@ -251,17 +256,15 @@ describe('Organization Billing', () => {
251256
})
252257

253258
it('should handle duplicate operation IDs gracefully', async () => {
254-
// Mock database constraint error
259+
// Mock database returning empty result for onConflictDoNothing (duplicate detected)
255260
await mockModule('@codebuff/internal/db', () => ({
256261
default: createDbMock({
257262
insert: () => ({
258-
values: () => {
259-
throw createPostgresError(
260-
'Duplicate key',
261-
'23505',
262-
'credit_ledger_pkey',
263-
)
264-
},
263+
values: () => ({
264+
onConflictDoNothing: () => ({
265+
returning: () => Promise.resolve([]), // Empty = duplicate, no insert
266+
}),
267+
}),
265268
}),
266269
}),
267270
}))
@@ -272,7 +275,7 @@ describe('Organization Billing', () => {
272275
const operationId = 'duplicate-operation'
273276
const description = 'Duplicate test'
274277

275-
// Should not throw, should handle gracefully
278+
// Should not throw, should handle gracefully via onConflictDoNothing
276279
await expect(
277280
grantOrganizationCredits({
278281
organizationId,

0 commit comments

Comments
 (0)