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..a7b6f481ab9 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(), @@ -120,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 }) }) @@ -138,9 +176,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 +222,39 @@ 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.each(PAYER_CASES)( + 'locks the workspace and its $label payer as FOR NO KEY UPDATE', + async ({ context, workspaceRow }) => { + mockWorkspaceRow.current = { ...workspaceRow } + + await incrementStorageUsageForBillingContextInTx(mockTx as unknown as DbOrTx, context, 100) + + expect(mockTxFor.mock.calls).toEqual([['no key update'], ['no key update']]) + } + ) + + 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, deltaBytes: 100 }], + legacyDeltas: [], + }) + + 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) 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..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 @@ -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 }