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

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,