fix(projects): fence file writers with a shared Project lock

lockWorkspaceProject now takes the Project key shared, so file writers in
one Project no longer serialize against each other while exclusive holders
(fork, split, archive, payer transfer, departure, account deletion,
retention) still fence them. lockWorkspaceProjectExclusive keeps the old
semantics for requireForkProject, splitForkProject, and archiveWorkspace.

Organization departure takes the departing workspaces' Project locks before
rewriting workspace owners, matching the writers' Project -> workspace row
-> payer order instead of inverting it.

Payer transfer skips zero-byte contributions before the destination payer
check, locks Projects in one ordered call instead of one-by-one and twice,
and drops the test-only changeWorkspaceStoragePayersInTx entry point.

Author: waleedlatif1

apps/sim/lib/billing/organizations/membership.ts

@@ -65,7 +65,10 @@ import {
   revokeUserSessionsTx,
 } from '@/lib/organizations/members/revocation'
 import { reassignOrganizationProjects } from '@/lib/projects/membership'
-import { handoffProjectsForOrganizationDepartureTx } from '@/lib/projects/resource-handoff'
+import {
+  handoffProjectsForOrganizationDepartureTx,
+  lockProjectsForResourceDepartureTx,
+} from '@/lib/projects/resource-handoff'
 import { removeWorkspaceSkillMembershipsTx } from '@/lib/skills/access'
 import { reassignSharedResourceOwnershipForWorkspaceMemberRemovalTx } from '@/lib/workspaces/resource-handoff'
 import { WorkspaceBillingAccountRemovalError } from '@/lib/workspaces/utils'
@@ -569,6 +572,8 @@ async function reassignOwnedOrganizationResourcesTx({
   organizationId: string
   workspaceIds: string[]
 }) {
+  /** Project gates precede the workspace owner rewrite, the order file writers lock in. */
+  if (workspaceIds.length > 0) await lockProjectsForResourceDepartureTx(tx, workspaceIds)
   const ownerId = await reassignOrganizationSharedResourcesTx(tx, organizationId, userId)
   if (!ownerId) return 0
 

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

@@ -67,10 +67,10 @@ vi.mock('drizzle-orm', () => ({
 
 import {
   changeOrganizationWorkspaceBilledAccountsInTx,
+  changeProjectAndWorkspaceStoragePayersInTx,
   changeWorkspaceStoragePayerInTx,
-  changeWorkspaceStoragePayersInTx,
 } from '@/lib/billing/storage/payer-transfer'
-import type { DbOrTx } from '@/lib/db/types'
+import type { DbOrTx, DbTransaction } from '@/lib/db/types'
 
 const { warn: loggerWarn } = getMockLogger('WorkspaceStoragePayerTransfer')
 
@@ -277,7 +277,7 @@ function createFakeBatchTx(state: FakeBatchTransferState) {
   return {
     execute,
     locks,
-    tx: { execute, select, update } as unknown as DbOrTx,
+    tx: { execute, select, update } as unknown as DbTransaction,
     updates,
   }
 }
@@ -508,7 +508,7 @@ describe('changeWorkspaceStoragePayerInTx', () => {
   })
 })
 
-describe('changeWorkspaceStoragePayersInTx', () => {
+describe('changeProjectAndWorkspaceStoragePayersInTx workspace changes', () => {
   it('uses the same payer lock order for opposite-direction moves', async () => {
     const personalToOrganization = createFakeBatchTx({
       exactBytes: { 'workspace-a': 10 },
@@ -537,41 +537,38 @@ describe('changeWorkspaceStoragePayersInTx', () => {
       ],
     })
 
-    await changeWorkspaceStoragePayersInTx(personalToOrganization.tx, [
-      {
-        workspaceId: 'workspace-a',
-        organizationId: 'org-a',
-        billedAccountUserId: 'org-owner',
-      },
-    ])
-    await changeWorkspaceStoragePayersInTx(organizationToPersonal.tx, [
-      {
-        workspaceId: 'workspace-b',
-        organizationId: null,
-        billedAccountUserId: 'user-z',
-      },
-    ])
+    await changeProjectAndWorkspaceStoragePayersInTx(personalToOrganization.tx, {
+      projectChanges: [],
+      workspaceChanges: [
+        {
+          workspaceId: 'workspace-a',
+          organizationId: 'org-a',
+          billedAccountUserId: 'org-owner',
+        },
+      ],
+    })
+    await changeProjectAndWorkspaceStoragePayersInTx(organizationToPersonal.tx, {
+      projectChanges: [],
+      workspaceChanges: [
+        {
+          workspaceId: 'workspace-b',
+          organizationId: null,
+          billedAccountUserId: 'user-z',
+        },
+      ],
+    })
 
+    const lockOrder = ['workspace', 'userStats', 'organization']
     expect(personalToOrganization.locks.map(({ table }) => table)).toEqual([
-      'workspace',
-      'userStats',
-      'organization',
+      ...lockOrder,
+      ...lockOrder,
     ])
     expect(organizationToPersonal.locks.map(({ table }) => table)).toEqual([
-      'workspace',
-      'userStats',
-      'organization',
-    ])
-    expect(personalToOrganization.updates.map(({ table }) => table)).toEqual([
-      'userStats',
-      'organization',
-      'workspace',
-    ])
-    expect(organizationToPersonal.updates.map(({ table }) => table)).toEqual([
-      'userStats',
-      'organization',
-      'workspace',
+      ...lockOrder,
+      ...lockOrder,
     ])
+    expect(personalToOrganization.updates.map(({ table }) => table)).toEqual(lockOrder)
+    expect(organizationToPersonal.updates.map(({ table }) => table)).toEqual(lockOrder)
   })
 
   it('aggregates payer deltas once and clamps outgoing bytes before adding incoming bytes', async () => {
@@ -594,24 +591,28 @@ describe('changeWorkspaceStoragePayersInTx', () => {
       ],
     })
 
-    await changeWorkspaceStoragePayersInTx(fake.tx, [
-      {
-        workspaceId: 'workspace-b',
-        organizationId: null,
-        billedAccountUserId: 'destination',
-      },
-      {
-        workspaceId: 'workspace-a',
-        organizationId: null,
-        billedAccountUserId: 'destination',
-      },
-    ])
+    await changeProjectAndWorkspaceStoragePayersInTx(fake.tx, {
+      projectChanges: [],
+      workspaceChanges: [
+        {
+          workspaceId: 'workspace-b',
+          organizationId: null,
+          billedAccountUserId: 'destination',
+        },
+        {
+          workspaceId: 'workspace-a',
+          organizationId: null,
+          billedAccountUserId: 'destination',
+        },
+      ],
+    })
 
     expect(fake.execute).toHaveBeenCalledTimes(1)
-    expect(fake.locks).toEqual([
+    const sortedLocks = [
       { ids: ['workspace-a', 'workspace-b'], table: 'workspace' },
       { ids: ['destination', 'source'], table: 'userStats' },
-    ])
+    ]
+    expect(fake.locks).toEqual([...sortedLocks, ...sortedLocks])
     const userStatsUpdates = fake.updates.filter(({ table }) => table === 'userStats')
     expect(userStatsUpdates).toHaveLength(1)
     expect(readCaseAssignments(userStatsUpdates[0]?.values.storageUsedBytes)).toEqual({
@@ -636,22 +637,29 @@ describe('changeWorkspaceStoragePayersInTx', () => {
     })
 
     await expect(
-      changeWorkspaceStoragePayersInTx(fake.tx, [
-        {
-          workspaceId: 'workspace-a',
-          organizationId: null,
-          billedAccountUserId: 'destination',
-          expectedCurrentPayer: {
+      changeProjectAndWorkspaceStoragePayersInTx(fake.tx, {
+        projectChanges: [],
+        workspaceChanges: [
+          {
+            workspaceId: 'workspace-a',
             organizationId: null,
-            billedAccountUserId: 'stale',
+            billedAccountUserId: 'destination',
+            expectedCurrentPayer: {
+              organizationId: null,
+              billedAccountUserId: 'stale',
+            },
           },
-        },
-      ])
+        ],
+      })
     ).rejects.toThrow('Workspace workspace-a payer changed before the transaction lock')
 
     expect(fake.execute).not.toHaveBeenCalled()
     expect(fake.updates).toEqual([])
-    expect(fake.locks).toEqual([{ ids: ['workspace-a'], table: 'workspace' }])
+    expect(fake.locks).toEqual([
+      { ids: ['workspace-a'], table: 'workspace' },
+      { ids: ['current', 'destination'], table: 'userStats' },
+      { ids: ['workspace-a'], table: 'workspace' },
+    ])
   })
 })
 

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

@@ -13,7 +13,7 @@ import { and, asc, eq, inArray, sql } from 'drizzle-orm'
 import type { BillingEntity } from '@/lib/billing/core/usage-log'
 import type { ProjectStorageOwnerSnapshot } from '@/lib/billing/storage/context'
 import type { DbOrTx, DbTransaction } from '@/lib/db/types'
-import { lockProject } from '@/lib/projects/membership'
+import { lockProjects } from '@/lib/projects/membership'
 
 const logger = createLogger('WorkspaceStoragePayerTransfer')
 
@@ -384,25 +384,13 @@ async function setStoragePayerUsagesBatch(
 }
 
 /**
- * Changes multiple workspace payers in one short caller-owned transaction.
- *
- * Workspace rows are locked in ascending ID order, exact live byte totals are
- * grouped into one query, and distinct payer rows are locked with all users
- * before all organizations and ascending IDs within each type. Source and
- * destination deltas are aggregated separately so an underfunded source is
- * clamped before incoming bytes are added. No destination quota is enforced.
+ * Locks workspace rows in ascending ID order and changes their payer metadata and exact
+ * live byte totals. The caller already holds every affected payer row and applies the
+ * returned contributions once, so no payer ledger is written here.
  */
-export async function changeWorkspaceStoragePayersInTx(
-  tx: DbOrTx,
-  changes: ChangeWorkspaceStoragePayerParams[]
-): Promise<ChangeWorkspaceStoragePayerResult[]> {
-  return changeWorkspaceStoragePayers(tx, changes, false)
-}
-
 async function changeWorkspaceStoragePayers(
   tx: DbOrTx,
-  changes: ChangeWorkspaceStoragePayerParams[],
-  deferPayerUpdates: boolean
+  changes: ChangeWorkspaceStoragePayerParams[]
 ): Promise<ChangeWorkspaceStoragePayerResult[]> {
   if (changes.length === 0) return []
 
@@ -442,98 +430,19 @@ async function changeWorkspaceStoragePayers(
     }
   }
 
-  const payerByKey = new Map<string, BillingEntity>()
-  const payerChangingWorkspaceIds: string[] = []
-  for (const workspaceId of workspaceIds) {
+  const payerChangingWorkspaceIds = workspaceIds.filter((workspaceId) => {
     const lockedWorkspace = workspaceById.get(workspaceId)
     const change = changesByWorkspaceId.get(workspaceId)
-    if (!lockedWorkspace || !change) continue
-    const oldPayer = getWorkspacePayer(lockedWorkspace)
-    const newPayer = getWorkspacePayer(change)
-    if (getPayerKey(oldPayer) === getPayerKey(newPayer)) continue
-    payerChangingWorkspaceIds.push(workspaceId)
-    payerByKey.set(getPayerKey(oldPayer), oldPayer)
-    payerByKey.set(getPayerKey(newPayer), newPayer)
-  }
+    if (!lockedWorkspace || !change) return false
+    return (
+      getPayerKey(getWorkspacePayer(lockedWorkspace)) !== getPayerKey(getWorkspacePayer(change))
+    )
+  })
 
   const exactBytesByWorkspaceId = await getExactWorkspaceStorageBytesBatch(
     tx,
     payerChangingWorkspaceIds
   )
-  if (!deferPayerUpdates) {
-    const payerUsageByKey = await lockStoragePayers(tx, payerByKey)
-    const deltaByPayerKey = new Map<string, PayerStorageDelta>()
-
-    for (const workspaceId of payerChangingWorkspaceIds) {
-      const lockedWorkspace = workspaceById.get(workspaceId)
-      const change = changesByWorkspaceId.get(workspaceId)
-      if (!lockedWorkspace || !change) continue
-      const exactBytes = exactBytesByWorkspaceId.get(workspaceId) ?? 0
-      const oldPayer = getWorkspacePayer(lockedWorkspace)
-      const newPayer = getWorkspacePayer(change)
-      const oldPayerKey = getPayerKey(oldPayer)
-      const newPayerKey = getPayerKey(newPayer)
-
-      if (payerUsageByKey.get(newPayerKey) === null) {
-        throw new Error(`Storage destination payer ${newPayerKey} not found`)
-      }
-
-      const sourceDelta = deltaByPayerKey.get(oldPayerKey) ?? {
-        incomingBytes: 0,
-        outgoingBytes: 0,
-      }
-      sourceDelta.outgoingBytes += exactBytes
-      if (!Number.isSafeInteger(sourceDelta.outgoingBytes)) {
-        throw new Error(`Storage source payer ${oldPayerKey} delta exceeds the safe integer range`)
-      }
-      deltaByPayerKey.set(oldPayerKey, sourceDelta)
-
-      const destinationDelta = deltaByPayerKey.get(newPayerKey) ?? {
-        incomingBytes: 0,
-        outgoingBytes: 0,
-      }
-      destinationDelta.incomingBytes += exactBytes
-      if (!Number.isSafeInteger(destinationDelta.incomingBytes)) {
-        throw new Error(
-          `Storage destination payer ${newPayerKey} delta exceeds the safe integer range`
-        )
-      }
-      deltaByPayerKey.set(newPayerKey, destinationDelta)
-    }
-
-    const nextUsageByKey = new Map<string, number>()
-    for (const [payerKey, delta] of [...deltaByPayerKey.entries()].sort(([left], [right]) =>
-      comparePayerKeys(left, right)
-    )) {
-      const currentUsage = payerUsageByKey.get(payerKey)
-      if (currentUsage === null || currentUsage === undefined) {
-        logger.warn('Storage source payer is missing during workspace payer batch change', {
-          sourcePayer: payerKey,
-          outgoingBytes: delta.outgoingBytes,
-        })
-        continue
-      }
-
-      const usageAfterOutgoing = Math.max(0, currentUsage - delta.outgoingBytes)
-      if (currentUsage < delta.outgoingBytes) {
-        logger.warn(
-          'Clamping drifted source storage aggregate during workspace payer batch change',
-          {
-            sourcePayer: payerKey,
-            sourceUsage: currentUsage,
-            outgoingBytes: delta.outgoingBytes,
-          }
-        )
-      }
-      const nextUsage = usageAfterOutgoing + delta.incomingBytes
-      if (!Number.isSafeInteger(nextUsage)) {
-        throw new Error(`Storage payer ${payerKey} exceeds the safe integer range`)
-      }
-      nextUsageByKey.set(payerKey, nextUsage)
-    }
-
-    await setStoragePayerUsagesBatch(tx, payerByKey, nextUsageByKey)
-  }
 
   const resultsByWorkspaceId = new Map<string, ChangeWorkspaceStoragePayerResult>()
   const workspaceUpdates = workspaceIds.map((workspaceId) => {
@@ -796,9 +705,14 @@ export async function changeProjectStoragePayersInTx(
   tx: DbTransaction,
   changes: readonly ChangeProjectStoragePayerParams[]
 ): Promise<ChangeProjectStoragePayerResult[]> {
+  await lockProjects(
+    tx,
+    changes.map((change) => change.projectId)
+  )
   return changeProjectStoragePayers(tx, changes, false)
 }
 
+/** The caller holds every changed Project's lock. */
 async function changeProjectStoragePayers(
   tx: DbTransaction,
   changes: readonly ChangeProjectStoragePayerParams[],
@@ -809,7 +723,6 @@ async function changeProjectStoragePayers(
   if (changesById.size !== changes.length)
     throw new Error('Storage payer batch contains duplicate Project IDs')
   const projectIds = [...changesById.keys()].sort()
-  for (const id of projectIds) await lockProject(tx, id)
   const owners = await tx
     .select({ id: project.id, ownerId: project.ownerId, organizationId: project.organizationId })
     .from(project)
@@ -905,6 +818,7 @@ async function applyPayerContributionTransfers(
   const usageByKey = await lockStoragePayers(tx, payerByKey)
   const deltas = new Map<string, PayerStorageDelta>()
   for (const result of contributions) {
+    if (result.billableBytes === 0) continue
     const sourceKey = getPayerKey(result.oldPayer)
     const destinationKey = result.newPayer ? getPayerKey(result.newPayer) : null
     if (sourceKey === destinationKey) continue
@@ -967,7 +881,7 @@ export async function changeProjectAndWorkspaceStoragePayersInTx(
     .sort()
   if (new Set(projectIds).size !== projectIds.length)
     throw new Error('Storage lifecycle batch contains duplicate Project IDs')
-  for (const id of projectIds) await lockProject(tx, id)
+  await lockProjects(tx, projectIds)
   const projects = projectIds.length
     ? await tx
         .select({
@@ -1043,7 +957,7 @@ export async function changeProjectAndWorkspaceStoragePayersInTx(
     removal.billableBytes = bytes
   }
   const projectResults = await changeProjectStoragePayers(tx, changes.projectChanges, true)
-  const workspaceResults = await changeWorkspaceStoragePayers(tx, changes.workspaceChanges, true)
+  const workspaceResults = await changeWorkspaceStoragePayers(tx, changes.workspaceChanges)
   await applyPayerContributionTransfers(tx, [
     ...projectResults,
     ...workspaceResults,