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.

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(),
@@ -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

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

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