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

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 {