fix(storage): stop workspace ledger locks from deadlocking on FK key-share (#7187)

* 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.

* 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.

* 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.

Author: waleedlatif1

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,

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()
   })

apps/sim/lib/billing/storage/payer-transfer.ts

@@ -125,15 +125,17 @@ 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<number | null> {
   if (payer.type === 'organization') {
     const [row] = await tx
       .select({ storageUsedBytes: organization.storageUsedBytes })
       .from(organization)
       .where(eq(organization.id, payer.id))
-      .for('update')
+      .for('no key update')
       .limit(1)
     return row?.storageUsedBytes ?? null
   }
@@ -142,7 +144,7 @@ async function lockStoragePayer(tx: DbOrTx, payer: BillingEntity): Promise<numbe
     .select({ storageUsedBytes: userStats.storageUsedBytes })
     .from(userStats)
     .where(eq(userStats.userId, payer.id))
-    .for('update')
+    .for('no key update')
     .limit(1)
   return row?.storageUsedBytes ?? null
 }
@@ -257,7 +259,7 @@ async function lockStoragePayers(
       .from(userStats)
       .where(inArray(userStats.userId, userIds))
       .orderBy(asc(userStats.userId))
-      .for('update')
+      .for('no key update')
     for (const row of rows) {
       usageByKey.set(getPayerKey({ type: 'user', id: row.id }), row.storageUsedBytes)
     }
@@ -269,7 +271,7 @@ async function lockStoragePayers(
       .from(organization)
       .where(inArray(organization.id, organizationIds))
       .orderBy(asc(organization.id))
-      .for('update')
+      .for('no key update')
     for (const row of rows) {
       usageByKey.set(getPayerKey({ type: 'organization', id: row.id }), row.storageUsedBytes)
     }
@@ -371,7 +373,7 @@ export async function changeWorkspaceStoragePayersInTx(
     .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]))
   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) {

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

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) {

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)

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
 

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() })

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) {

apps/sim/lib/workspaces/organization-workspaces.ts

@@ -94,15 +94,20 @@ 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<void> {
   if (workspaceIds.length === 0) return
   await tx
     .select({ id: workspace.id })
     .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 {