From aea143a1448cde8b9bf2bbfbed9a15afa1a44cd6 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 27 Aug 2026 18:05:55 -0700 Subject: [PATCH 1/3] fix(storage): stop workspace ledger locks from deadlocking on FK key-share Workspace storage accounting locked the workspace, organization, and user_stats rows with SELECT ... FOR UPDATE. Those rows are foreign-key parents, so a transaction that has already written a billable child row holds an implicit FOR KEY SHARE on the parent, and the stronger lock is an upgrade that two concurrent uploads take on each other. Take FOR NO KEY UPDATE instead. It does not conflict with FOR KEY SHARE, still conflicts with itself, and is the lock a plain UPDATE of these non-key counters takes anyway, so the ledgers stay serialized. --- .../workspaces/[id]/permissions/route.test.ts | 2 +- .../billing/storage/payer-transfer.test.ts | 2 +- .../sim/lib/billing/storage/payer-transfer.ts | 18 ++++---- apps/sim/lib/billing/storage/tracking.test.ts | 46 ++++++++++++++++++- apps/sim/lib/billing/storage/tracking.ts | 28 ++++++++--- .../tools/handlers/materialize-file.ts | 3 +- apps/sim/lib/credentials/environment.ts | 3 +- apps/sim/lib/table/service.ts | 10 ++-- apps/sim/lib/workspaces/admin-move.ts | 8 +++- .../lib/workspaces/organization-workspaces.ts | 11 +++-- apps/sim/lib/workspaces/permissions/utils.ts | 9 +++- 11 files changed, 113 insertions(+), 27 deletions(-) diff --git a/apps/sim/app/api/workspaces/[id]/permissions/route.test.ts b/apps/sim/app/api/workspaces/[id]/permissions/route.test.ts index b163b2bb428..36683982480 100644 --- a/apps/sim/app/api/workspaces/[id]/permissions/route.test.ts +++ b/apps/sim/app/api/workspaces/[id]/permissions/route.test.ts @@ -73,7 +73,7 @@ function queuePersonalWorkspace( ) { const workspaceRow = { ownerId: OWNER_ID, billedAccountUserId, organizationId: null } queueTableRows(schemaMock.workspace, [workspaceRow]) - /** The in-transaction re-read of the same row, taken `FOR UPDATE`. */ + /** The in-transaction re-read of the same row, taken `FOR NO KEY UPDATE`. */ permissionsMockFns.mockGetWorkspaceWithOwner.mockResolvedValue({ id: WORKSPACE_ID, ...workspaceRow, diff --git a/apps/sim/lib/billing/storage/payer-transfer.test.ts b/apps/sim/lib/billing/storage/payer-transfer.test.ts index 8113e236e5f..a9518c25add 100644 --- a/apps/sim/lib/billing/storage/payer-transfer.test.ts +++ b/apps/sim/lib/billing/storage/payer-transfer.test.ts @@ -769,7 +769,7 @@ describe('changeOrganizationWorkspaceBilledAccountsInTx', () => { expect(returning).toHaveBeenCalledWith({ id: 'workspace.id' }) expect(select).toHaveBeenCalledWith({ id: 'workspace.id' }) expect(orderBy).toHaveBeenCalledTimes(1) - expect(lock).toHaveBeenCalledWith('update') + expect(lock).toHaveBeenCalledWith('no key update') expect(lock.mock.invocationCallOrder[0]).toBeLessThan(update.mock.invocationCallOrder[0]) expect(execute).not.toHaveBeenCalled() }) diff --git a/apps/sim/lib/billing/storage/payer-transfer.ts b/apps/sim/lib/billing/storage/payer-transfer.ts index 14687a4df56..049070dfd8b 100644 --- a/apps/sim/lib/billing/storage/payer-transfer.ts +++ b/apps/sim/lib/billing/storage/payer-transfer.ts @@ -125,7 +125,9 @@ async function getExactWorkspaceStorageBytes(tx: DbOrTx, workspaceId: string): P /** * Locks a payer row and returns its current aggregate. A missing source can be * historical drift and is represented as `null`; callers must reject a - * missing destination. + * missing destination. `FOR NO KEY UPDATE` avoids upgrading the implicit + * foreign-key `FOR KEY SHARE` this transaction may already hold; see the + * module header of `lib/billing/storage/tracking.ts`. */ async function lockStoragePayer(tx: DbOrTx, payer: BillingEntity): Promise { if (payer.type === 'organization') { @@ -133,7 +135,7 @@ async function lockStoragePayer(tx: DbOrTx, payer: BillingEntity): Promise [row.id, row])) for (const workspaceId of workspaceIds) { @@ -562,7 +564,7 @@ export async function changeOrganizationWorkspaceBilledAccountsInTx( ) ) .orderBy(asc(workspace.id)) - .for('update') + .for('no key update') const rows = await tx .update(workspace) @@ -604,7 +606,7 @@ export async function changeWorkspaceStoragePayerInTx( }) .from(workspace) .where(eq(workspace.id, params.workspaceId)) - .for('update') + .for('no key update') .limit(1) if (!lockedWorkspace) { diff --git a/apps/sim/lib/billing/storage/tracking.test.ts b/apps/sim/lib/billing/storage/tracking.test.ts index f8d36e8202d..b72dfc6efc4 100644 --- a/apps/sim/lib/billing/storage/tracking.test.ts +++ b/apps/sim/lib/billing/storage/tracking.test.ts @@ -13,6 +13,7 @@ const { mockMaybeNotifyLimit, mockOrderedLockRows, mockSql, + mockTxFor, mockTxFrom, mockTxLimit, mockTxOrderBy, @@ -32,6 +33,7 @@ const { mockMaybeNotifyLimit: vi.fn(), mockOrderedLockRows: { queue: [] as unknown[][] }, mockSql: vi.fn((strings: TemplateStringsArray, ...values: unknown[]) => ({ strings, values })), + mockTxFor: vi.fn(), mockTxFrom: vi.fn(), mockTxLimit: vi.fn(), mockTxOrderBy: vi.fn(), @@ -138,9 +140,10 @@ describe('workspace storage counter mutations', () => { mockOrderedLockRows.queue = [] mockTxSelect.mockReturnValue({ from: mockTxFrom }) + mockTxFor.mockReturnValue({ limit: mockTxLimit }) mockTxFrom.mockReturnValue({ where: vi.fn(() => ({ - for: vi.fn(() => ({ limit: mockTxLimit })), + for: mockTxFor, limit: mockTxLimit, orderBy: mockTxOrderBy, })), @@ -183,6 +186,47 @@ describe('workspace storage counter mutations', () => { expect(mockMaybeNotifyLimit).not.toHaveBeenCalled() }) + /** + * `FOR UPDATE` on these rows deadlocked in production: `workspace`, + * `organization`, and `user_stats` are foreign-key parents, so the calling + * transaction already holds an implicit `FOR KEY SHARE` on them from the + * billable child row it just wrote, and the stronger lock is an upgrade that + * two concurrent uploads take on each other. `FOR NO KEY UPDATE` still + * conflicts with itself, so the ledgers stay serialized. + */ + it('takes every ledger lock as FOR NO KEY UPDATE so it never upgrades a key-share lock', async () => { + await incrementStorageUsageForBillingContextInTx(mockTx as unknown as DbOrTx, ORG_CONTEXT, 100) + + expect(mockTxFor).toHaveBeenCalled() + for (const call of mockTxFor.mock.calls) { + expect(call).toEqual(['no key update']) + } + }) + + it('takes batch ledger locks as FOR NO KEY UPDATE', async () => { + mockOrderedLockRows.queue = [ + [ + { + id: 'workspace-1', + billedAccountUserId: 'workspace-owner', + organizationId: 'workspace-org', + storageUsedBytes: 1_000, + }, + ], + [{ id: 'workspace-org', storageUsedBytes: 1_000 }], + ] + + await applyStorageUsageDeltasInTx(mockTx as unknown as DbOrTx, { + workspaceDeltas: [{ context: ORG_CONTEXT, deltaBytes: 100 }], + legacyDeltas: [], + }) + + expect(mockTxOrderedFor).toHaveBeenCalled() + for (const call of mockTxOrderedFor.mock.calls) { + expect(call).toEqual(['no key update']) + } + }) + it('serializes quota admission on the locked payer ledger', async () => { mockGetStorageLimitForBillingContext.mockReturnValue(1_050) mockTxLimit diff --git a/apps/sim/lib/billing/storage/tracking.ts b/apps/sim/lib/billing/storage/tracking.ts index 0de9b03507b..7ae702a94c3 100644 --- a/apps/sim/lib/billing/storage/tracking.ts +++ b/apps/sim/lib/billing/storage/tracking.ts @@ -1,5 +1,20 @@ /** * Storage usage tracking for durable workspace and payer ledgers. + * + * Every row lock here is `FOR NO KEY UPDATE`, never `FOR UPDATE`. The + * `workspace`, `organization`, and `user_stats` rows these transactions lock + * are foreign-key parents (49 tables reference `workspace` alone), so any + * insert or update of a child row — a `workspace_files` row in this very + * transaction — implicitly takes `FOR KEY SHARE` on the parent first. A later + * `FOR UPDATE` on the same row is then a lock upgrade, and two concurrent + * uploads or deletes in one workspace deadlock on it. `FOR NO KEY UPDATE` + * does not conflict with `FOR KEY SHARE`, yet still conflicts with itself and + * with `FOR UPDATE`, so writers remain serialized against each other and + * against payer transfers. It is exactly the lock a plain `UPDATE` of these + * non-key counters takes anyway. The only key columns on these tables are + * `workspace.id`, `workspace.inbox_provider_id`, `organization.id`, + * `user_stats.id`, and `user_stats.user_id`, and no path under these locks + * writes any of them or deletes a locked row. */ import { organization, userStats, workspace } from '@sim/db/schema' @@ -124,6 +139,7 @@ async function mutateStorageUsage( /** * Locks and reads the payer ledger after the workspace row has been locked. + * `FOR NO KEY UPDATE` for the reason documented at the top of this module. */ async function lockStorageUsageForMutation( tx: DbOrTx, @@ -134,7 +150,7 @@ async function lockStorageUsageForMutation( .select({ storageUsedBytes: organization.storageUsedBytes }) .from(organization) .where(eq(organization.id, billingEntity.id)) - .for('update') + .for('no key update') .limit(1) if (!row) throw new Error(`Storage payer organization:${billingEntity.id} not found`) return row.storageUsedBytes @@ -144,7 +160,7 @@ async function lockStorageUsageForMutation( .select({ storageUsedBytes: userStats.storageUsedBytes }) .from(userStats) .where(eq(userStats.userId, billingEntity.id)) - .for('update') + .for('no key update') .limit(1) if (!row) throw new Error(`Storage payer user:${billingEntity.id} not found`) return row.storageUsedBytes @@ -242,7 +258,7 @@ export async function applyStorageUsageDeltasInTx( .from(workspace) .where(inArray(workspace.id, workspaceIds)) .orderBy(asc(workspace.id)) - .for('update') + .for('no key update') : [] const workspaceById = new Map(lockedWorkspaces.map((row) => [row.id, row])) @@ -318,7 +334,7 @@ export async function applyStorageUsageDeltasInTx( .from(userStats) .where(inArray(userStats.userId, userIds)) .orderBy(asc(userStats.userId)) - .for('update') + .for('no key update') for (const row of rows) { payerUsageByKey.set(getPayerKey({ type: 'user', id: row.id }), row.storageUsedBytes) } @@ -329,7 +345,7 @@ export async function applyStorageUsageDeltasInTx( .from(organization) .where(inArray(organization.id, organizationIds)) .orderBy(asc(organization.id)) - .for('update') + .for('no key update') for (const row of rows) { payerUsageByKey.set(getPayerKey({ type: 'organization', id: row.id }), row.storageUsedBytes) } @@ -439,7 +455,7 @@ async function mutateWorkspaceStorageUsage( }) .from(workspace) .where(eq(workspace.id, workspaceId)) - .for('update') + .for('no key update') .limit(1) if (!workspacePayer) { diff --git a/apps/sim/lib/copilot/tools/handlers/materialize-file.ts b/apps/sim/lib/copilot/tools/handlers/materialize-file.ts index 755c8b718d8..3ebfb5e2941 100644 --- a/apps/sim/lib/copilot/tools/handlers/materialize-file.ts +++ b/apps/sim/lib/copilot/tools/handlers/materialize-file.ts @@ -140,7 +140,8 @@ async function executeSave( try { transition = await db.transaction(async (tx) => { - await tx.execute(sql`SELECT 1 FROM workspace WHERE id = ${workspaceId} FOR UPDATE`) + /** `FOR NO KEY UPDATE`: see the module header of `lib/billing/storage/tracking.ts`. */ + await tx.execute(sql`SELECT 1 FROM workspace WHERE id = ${workspaceId} FOR NO KEY UPDATE`) const [updated] = await tx .update(workspaceFiles) diff --git a/apps/sim/lib/credentials/environment.ts b/apps/sim/lib/credentials/environment.ts index ad808306d30..9756c4659e7 100644 --- a/apps/sim/lib/credentials/environment.ts +++ b/apps/sim/lib/credentials/environment.ts @@ -78,8 +78,9 @@ export async function getCredentialCreationWorkspaceContext(params: { }) .from(workspace) .where(and(eq(workspace.id, params.workspaceId), isNull(workspace.archivedAt))) + /** `FOR NO KEY UPDATE`: see the module header of `lib/billing/storage/tracking.ts`. */ const [workspaceRow] = params.forUpdate - ? await workspaceQuery.for('update').limit(1) + ? await workspaceQuery.for('no key update').limit(1) : await workspaceQuery.limit(1) if (!workspaceRow) return null diff --git a/apps/sim/lib/table/service.ts b/apps/sim/lib/table/service.ts index 6063ec9be9c..8365aac490b 100644 --- a/apps/sim/lib/table/service.ts +++ b/apps/sim/lib/table/service.ts @@ -613,12 +613,16 @@ export async function createTable( }) } - // Wrap count check, duplicate check, and insert in a transaction with FOR UPDATE - // to prevent TOCTOU race on the table count limit + // Wrap count check, duplicate check, and insert in a transaction with FOR NO KEY UPDATE + // to prevent TOCTOU race on the table count limit. The weaker lock still conflicts with + // itself, so table creations stay serialized, but it does not block unrelated inserts + // into the workspace's other child tables. See lib/billing/storage/tracking.ts. try { await db.transaction(async (trx) => { await setTableTxTimeouts(trx) - await trx.execute(sql`SELECT 1 FROM workspace WHERE id = ${data.workspaceId} FOR UPDATE`) + await trx.execute( + sql`SELECT 1 FROM workspace WHERE id = ${data.workspaceId} FOR NO KEY UPDATE` + ) const [{ count: existingCount }] = await trx .select({ count: count() }) diff --git a/apps/sim/lib/workspaces/admin-move.ts b/apps/sim/lib/workspaces/admin-move.ts index aeb795682f9..401fb6ea787 100644 --- a/apps/sim/lib/workspaces/admin-move.ts +++ b/apps/sim/lib/workspaces/admin-move.ts @@ -502,6 +502,12 @@ export async function moveWorkspaceToOrganization(params: { throw new InvitationSetChangedError(currentInvitationIds) } + /** + * `FOR NO KEY UPDATE`, not `FOR UPDATE`: the workspace row is a + * foreign-key parent, so concurrent writers hold an implicit + * `FOR KEY SHARE` on it. See the module header of + * `lib/billing/storage/tracking.ts`. + */ const [workspaceRow] = await tx .select({ id: workspace.id, @@ -513,7 +519,7 @@ export async function moveWorkspaceToOrganization(params: { }) .from(workspace) .where(eq(workspace.id, params.workspaceId)) - .for('update') + .for('no key update') .limit(1) if (!workspaceRow) { diff --git a/apps/sim/lib/workspaces/organization-workspaces.ts b/apps/sim/lib/workspaces/organization-workspaces.ts index 49c261f0d5a..ac36c8efbb9 100644 --- a/apps/sim/lib/workspaces/organization-workspaces.ts +++ b/apps/sim/lib/workspaces/organization-workspaces.ts @@ -94,7 +94,12 @@ export function ownedAttachableWorkspacesWhere({ ) } -/** Locks workspace rows before any payer or membership mutation. */ +/** + * Locks workspace rows before any payer or membership mutation. `FOR NO KEY + * UPDATE` keeps this compatible with the implicit foreign-key `FOR KEY SHARE` + * concurrent writers hold; see the module header of + * `lib/billing/storage/tracking.ts`. + */ async function lockWorkspaceRowsForPayerChanges(tx: DbOrTx, workspaceIds: string[]): Promise { if (workspaceIds.length === 0) return await tx @@ -102,7 +107,7 @@ async function lockWorkspaceRowsForPayerChanges(tx: DbOrTx, workspaceIds: string .from(workspace) .where(inArray(workspace.id, [...workspaceIds].sort())) .orderBy(asc(workspace.id)) - .for('update') + .for('no key update') } interface AttachOwnedWorkspacesToOrganizationParams { @@ -243,7 +248,7 @@ export async function attachOwnedWorkspacesToOrganizationTx( ) ) .orderBy(asc(workspace.id)) - .for('update') + .for('no key update') if (ownedWorkspaces.length === 0) { return { diff --git a/apps/sim/lib/workspaces/permissions/utils.ts b/apps/sim/lib/workspaces/permissions/utils.ts index cf023292aa7..19fb541f3ca 100644 --- a/apps/sim/lib/workspaces/permissions/utils.ts +++ b/apps/sim/lib/workspaces/permissions/utils.ts @@ -78,6 +78,13 @@ export async function getWorkspaceById( return exists ? { id: workspaceId } : null } +/** + * Reads one workspace row, optionally locking it. The lock is `FOR NO KEY + * UPDATE`: the workspace row is a foreign-key parent, and `FOR UPDATE` would + * both block every concurrent insert into its child tables and deadlock + * against callers that write a child row first. See the module header of + * `lib/billing/storage/tracking.ts`. + */ async function selectWorkspaceWithOwner( workspaceId: string, includeArchived: boolean, @@ -101,7 +108,7 @@ async function selectWorkspaceWithOwner( ? eq(workspace.id, workspaceId) : and(eq(workspace.id, workspaceId), isNull(workspace.archivedAt)) ) - const [ws] = forUpdate ? await query.for('update').limit(1) : await query.limit(1) + const [ws] = forUpdate ? await query.for('no key update').limit(1) : await query.limit(1) return ws || null } From ba3dd81e79d4f98247d7dbe1c8dbe09c911c378c Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 27 Aug 2026 18:16:51 -0700 Subject: [PATCH 2/3] test(storage): cover both payer kinds in the ledger lock-mode assertions The lock-mode regression tests only exercised the organization payer, so the user_stats lock branches were never asserted and a revert of just those would have passed. Parameterize both tests over both payer kinds and assert the exact call list, so a lock that stops being taken at all fails too. --- apps/sim/lib/billing/storage/tracking.test.ts | 80 +++++++++++++------ 1 file changed, 54 insertions(+), 26 deletions(-) diff --git a/apps/sim/lib/billing/storage/tracking.test.ts b/apps/sim/lib/billing/storage/tracking.test.ts index b72dfc6efc4..a7b6f481ab9 100644 --- a/apps/sim/lib/billing/storage/tracking.test.ts +++ b/apps/sim/lib/billing/storage/tracking.test.ts @@ -122,6 +122,42 @@ const ORG_CONTEXT: StorageBillingContext = { customStorageLimitGB: null, } +const USER_CONTEXT: StorageBillingContext = { + workspaceId: 'workspace-1', + billedAccountUserId: 'workspace-owner', + billingEntity: { type: 'user', id: 'workspace-owner' }, + plan: 'pro', + customStorageLimitGB: null, +} + +/** + * Both payer kinds. The workspace lock is shared, but the payer lock branches + * to a different table per kind, so a lock-mode regression on only one of them + * has to fail a test. + */ +const PAYER_CASES = [ + { + label: 'organization', + context: ORG_CONTEXT, + workspaceRow: { + billedAccountUserId: 'workspace-owner', + organizationId: 'workspace-org' as string | null, + storageUsedBytes: 1_000, + }, + payerLockRows: [{ id: 'workspace-org', storageUsedBytes: 1_000 }], + }, + { + label: 'user', + context: USER_CONTEXT, + workspaceRow: { + billedAccountUserId: 'workspace-owner', + organizationId: null as string | null, + storageUsedBytes: 1_000, + }, + payerLockRows: [{ id: 'workspace-owner', storageUsedBytes: 1_000 }], + }, +] as const + beforeAll(() => { setEnvFlags({ isBillingEnabled: true }) }) @@ -194,38 +230,30 @@ describe('workspace storage counter mutations', () => { * two concurrent uploads take on each other. `FOR NO KEY UPDATE` still * conflicts with itself, so the ledgers stay serialized. */ - it('takes every ledger lock as FOR NO KEY UPDATE so it never upgrades a key-share lock', async () => { - await incrementStorageUsageForBillingContextInTx(mockTx as unknown as DbOrTx, ORG_CONTEXT, 100) + it.each(PAYER_CASES)( + 'locks the workspace and its $label payer as FOR NO KEY UPDATE', + async ({ context, workspaceRow }) => { + mockWorkspaceRow.current = { ...workspaceRow } - expect(mockTxFor).toHaveBeenCalled() - for (const call of mockTxFor.mock.calls) { - expect(call).toEqual(['no key update']) + await incrementStorageUsageForBillingContextInTx(mockTx as unknown as DbOrTx, context, 100) + + expect(mockTxFor.mock.calls).toEqual([['no key update'], ['no key update']]) } - }) + ) - it('takes batch ledger locks as FOR NO KEY UPDATE', async () => { - mockOrderedLockRows.queue = [ - [ - { - id: 'workspace-1', - billedAccountUserId: 'workspace-owner', - organizationId: 'workspace-org', - storageUsedBytes: 1_000, - }, - ], - [{ id: 'workspace-org', storageUsedBytes: 1_000 }], - ] + it.each(PAYER_CASES)( + 'locks batched workspace and $label payer ledgers as FOR NO KEY UPDATE', + async ({ context, workspaceRow, payerLockRows }) => { + mockOrderedLockRows.queue = [[{ id: 'workspace-1', ...workspaceRow }], [...payerLockRows]] - await applyStorageUsageDeltasInTx(mockTx as unknown as DbOrTx, { - workspaceDeltas: [{ context: ORG_CONTEXT, deltaBytes: 100 }], - legacyDeltas: [], - }) + await applyStorageUsageDeltasInTx(mockTx as unknown as DbOrTx, { + workspaceDeltas: [{ context, deltaBytes: 100 }], + legacyDeltas: [], + }) - expect(mockTxOrderedFor).toHaveBeenCalled() - for (const call of mockTxOrderedFor.mock.calls) { - expect(call).toEqual(['no key update']) + expect(mockTxOrderedFor.mock.calls).toEqual([['no key update'], ['no key update']]) } - }) + ) it('serializes quota admission on the locked payer ledger', async () => { mockGetStorageLimitForBillingContext.mockReturnValue(1_050) From 737ff44ba4481c1774a67b6f444a57fde58dc0ee Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 27 Aug 2026 18:33:27 -0700 Subject: [PATCH 3/3] docs(table): correct the stale FOR UPDATE reference in the quota note The advisory quota check describes createTable's count as the authoritative FOR UPDATE read; that lock is now FOR NO KEY UPDATE. --- apps/sim/lib/table/service.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/sim/lib/table/service.ts b/apps/sim/lib/table/service.ts index 8365aac490b..cc35089085c 100644 --- a/apps/sim/lib/table/service.ts +++ b/apps/sim/lib/table/service.ts @@ -506,7 +506,7 @@ function workspaceTableLimitReached(maxTables: number): ForbiddenOperationError * Advisory table-quota check for a caller that is about to make the user pay * for work before {@link createTable} would run. * - * The authoritative check is the `FOR UPDATE` count inside `createTable`'s + * The authoritative check is the `FOR NO KEY UPDATE` count inside `createTable`'s * transaction and stays there — this one races, by construction, because the * ceiling can be reached (or cleared) during whatever the caller does next. It * exists so that "next" is not a multi-gigabyte upload: the CSV import used to