Skip to content

Commit e13228a

Browse files
authored
fix(projects): fence file writers with a shared Project lock (#9025)
* 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. * fix(files): lock the workspace row before bulk archive and folder restore Their multi-row deleted_at updates fire the search-index trigger, which upserts the dispatch queue row before later file rows are locked. A concurrent content write holds the workspace row and a file row and then waits on that queue row, so the two could deadlock. The exclusive Project lock used to hide this; take the workspace row first, the order every file writer uses. * fix(files): lock the workspace row before path-based folder delete archives file rows
1 parent 2d0c050 commit e13228a

8 files changed

Lines changed: 176 additions & 177 deletions

File tree

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,10 @@ import {
6565
revokeUserSessionsTx,
6666
} from '@/lib/organizations/members/revocation'
6767
import { reassignOrganizationProjects } from '@/lib/projects/membership'
68-
import { handoffProjectsForOrganizationDepartureTx } from '@/lib/projects/resource-handoff'
68+
import {
69+
handoffProjectsForOrganizationDepartureTx,
70+
lockProjectsForResourceDepartureTx,
71+
} from '@/lib/projects/resource-handoff'
6972
import { removeWorkspaceSkillMembershipsTx } from '@/lib/skills/access'
7073
import { reassignSharedResourceOwnershipForWorkspaceMemberRemovalTx } from '@/lib/workspaces/resource-handoff'
7174
import { WorkspaceBillingAccountRemovalError } from '@/lib/workspaces/utils'
@@ -569,6 +572,8 @@ async function reassignOwnedOrganizationResourcesTx({
569572
organizationId: string
570573
workspaceIds: string[]
571574
}) {
575+
/** Project gates precede the workspace owner rewrite, the order file writers lock in. */
576+
if (workspaceIds.length > 0) await lockProjectsForResourceDepartureTx(tx, workspaceIds)
572577
const ownerId = await reassignOrganizationSharedResourcesTx(tx, organizationId, userId)
573578
if (!ownerId) return 0
574579

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

Lines changed: 66 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -67,10 +67,10 @@ vi.mock('drizzle-orm', () => ({
6767

6868
import {
6969
changeOrganizationWorkspaceBilledAccountsInTx,
70+
changeProjectAndWorkspaceStoragePayersInTx,
7071
changeWorkspaceStoragePayerInTx,
71-
changeWorkspaceStoragePayersInTx,
7272
} from '@/lib/billing/storage/payer-transfer'
73-
import type { DbOrTx } from '@/lib/db/types'
73+
import type { DbOrTx, DbTransaction } from '@/lib/db/types'
7474

7575
const { warn: loggerWarn } = getMockLogger('WorkspaceStoragePayerTransfer')
7676

@@ -277,7 +277,7 @@ function createFakeBatchTx(state: FakeBatchTransferState) {
277277
return {
278278
execute,
279279
locks,
280-
tx: { execute, select, update } as unknown as DbOrTx,
280+
tx: { execute, select, update } as unknown as DbTransaction,
281281
updates,
282282
}
283283
}
@@ -508,7 +508,7 @@ describe('changeWorkspaceStoragePayerInTx', () => {
508508
})
509509
})
510510

511-
describe('changeWorkspaceStoragePayersInTx', () => {
511+
describe('changeProjectAndWorkspaceStoragePayersInTx workspace changes', () => {
512512
it('uses the same payer lock order for opposite-direction moves', async () => {
513513
const personalToOrganization = createFakeBatchTx({
514514
exactBytes: { 'workspace-a': 10 },
@@ -537,41 +537,38 @@ describe('changeWorkspaceStoragePayersInTx', () => {
537537
],
538538
})
539539

540-
await changeWorkspaceStoragePayersInTx(personalToOrganization.tx, [
541-
{
542-
workspaceId: 'workspace-a',
543-
organizationId: 'org-a',
544-
billedAccountUserId: 'org-owner',
545-
},
546-
])
547-
await changeWorkspaceStoragePayersInTx(organizationToPersonal.tx, [
548-
{
549-
workspaceId: 'workspace-b',
550-
organizationId: null,
551-
billedAccountUserId: 'user-z',
552-
},
553-
])
540+
await changeProjectAndWorkspaceStoragePayersInTx(personalToOrganization.tx, {
541+
projectChanges: [],
542+
workspaceChanges: [
543+
{
544+
workspaceId: 'workspace-a',
545+
organizationId: 'org-a',
546+
billedAccountUserId: 'org-owner',
547+
},
548+
],
549+
})
550+
await changeProjectAndWorkspaceStoragePayersInTx(organizationToPersonal.tx, {
551+
projectChanges: [],
552+
workspaceChanges: [
553+
{
554+
workspaceId: 'workspace-b',
555+
organizationId: null,
556+
billedAccountUserId: 'user-z',
557+
},
558+
],
559+
})
554560

561+
const lockOrder = ['workspace', 'userStats', 'organization']
555562
expect(personalToOrganization.locks.map(({ table }) => table)).toEqual([
556-
'workspace',
557-
'userStats',
558-
'organization',
563+
...lockOrder,
564+
...lockOrder,
559565
])
560566
expect(organizationToPersonal.locks.map(({ table }) => table)).toEqual([
561-
'workspace',
562-
'userStats',
563-
'organization',
564-
])
565-
expect(personalToOrganization.updates.map(({ table }) => table)).toEqual([
566-
'userStats',
567-
'organization',
568-
'workspace',
569-
])
570-
expect(organizationToPersonal.updates.map(({ table }) => table)).toEqual([
571-
'userStats',
572-
'organization',
573-
'workspace',
567+
...lockOrder,
568+
...lockOrder,
574569
])
570+
expect(personalToOrganization.updates.map(({ table }) => table)).toEqual(lockOrder)
571+
expect(organizationToPersonal.updates.map(({ table }) => table)).toEqual(lockOrder)
575572
})
576573

577574
it('aggregates payer deltas once and clamps outgoing bytes before adding incoming bytes', async () => {
@@ -594,24 +591,28 @@ describe('changeWorkspaceStoragePayersInTx', () => {
594591
],
595592
})
596593

597-
await changeWorkspaceStoragePayersInTx(fake.tx, [
598-
{
599-
workspaceId: 'workspace-b',
600-
organizationId: null,
601-
billedAccountUserId: 'destination',
602-
},
603-
{
604-
workspaceId: 'workspace-a',
605-
organizationId: null,
606-
billedAccountUserId: 'destination',
607-
},
608-
])
594+
await changeProjectAndWorkspaceStoragePayersInTx(fake.tx, {
595+
projectChanges: [],
596+
workspaceChanges: [
597+
{
598+
workspaceId: 'workspace-b',
599+
organizationId: null,
600+
billedAccountUserId: 'destination',
601+
},
602+
{
603+
workspaceId: 'workspace-a',
604+
organizationId: null,
605+
billedAccountUserId: 'destination',
606+
},
607+
],
608+
})
609609

610610
expect(fake.execute).toHaveBeenCalledTimes(1)
611-
expect(fake.locks).toEqual([
611+
const sortedLocks = [
612612
{ ids: ['workspace-a', 'workspace-b'], table: 'workspace' },
613613
{ ids: ['destination', 'source'], table: 'userStats' },
614-
])
614+
]
615+
expect(fake.locks).toEqual([...sortedLocks, ...sortedLocks])
615616
const userStatsUpdates = fake.updates.filter(({ table }) => table === 'userStats')
616617
expect(userStatsUpdates).toHaveLength(1)
617618
expect(readCaseAssignments(userStatsUpdates[0]?.values.storageUsedBytes)).toEqual({
@@ -636,22 +637,29 @@ describe('changeWorkspaceStoragePayersInTx', () => {
636637
})
637638

638639
await expect(
639-
changeWorkspaceStoragePayersInTx(fake.tx, [
640-
{
641-
workspaceId: 'workspace-a',
642-
organizationId: null,
643-
billedAccountUserId: 'destination',
644-
expectedCurrentPayer: {
640+
changeProjectAndWorkspaceStoragePayersInTx(fake.tx, {
641+
projectChanges: [],
642+
workspaceChanges: [
643+
{
644+
workspaceId: 'workspace-a',
645645
organizationId: null,
646-
billedAccountUserId: 'stale',
646+
billedAccountUserId: 'destination',
647+
expectedCurrentPayer: {
648+
organizationId: null,
649+
billedAccountUserId: 'stale',
650+
},
647651
},
648-
},
649-
])
652+
],
653+
})
650654
).rejects.toThrow('Workspace workspace-a payer changed before the transaction lock')
651655

652656
expect(fake.execute).not.toHaveBeenCalled()
653657
expect(fake.updates).toEqual([])
654-
expect(fake.locks).toEqual([{ ids: ['workspace-a'], table: 'workspace' }])
658+
expect(fake.locks).toEqual([
659+
{ ids: ['workspace-a'], table: 'workspace' },
660+
{ ids: ['current', 'destination'], table: 'userStats' },
661+
{ ids: ['workspace-a'], table: 'workspace' },
662+
])
655663
})
656664
})
657665

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

Lines changed: 19 additions & 105 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ import { and, asc, eq, inArray, sql } from 'drizzle-orm'
1313
import type { BillingEntity } from '@/lib/billing/core/usage-log'
1414
import type { ProjectStorageOwnerSnapshot } from '@/lib/billing/storage/context'
1515
import type { DbOrTx, DbTransaction } from '@/lib/db/types'
16-
import { lockProject } from '@/lib/projects/membership'
16+
import { lockProjects } from '@/lib/projects/membership'
1717

1818
const logger = createLogger('WorkspaceStoragePayerTransfer')
1919

@@ -384,25 +384,13 @@ async function setStoragePayerUsagesBatch(
384384
}
385385

386386
/**
387-
* Changes multiple workspace payers in one short caller-owned transaction.
388-
*
389-
* Workspace rows are locked in ascending ID order, exact live byte totals are
390-
* grouped into one query, and distinct payer rows are locked with all users
391-
* before all organizations and ascending IDs within each type. Source and
392-
* destination deltas are aggregated separately so an underfunded source is
393-
* clamped before incoming bytes are added. No destination quota is enforced.
387+
* Locks workspace rows in ascending ID order and changes their payer metadata and exact
388+
* live byte totals. The caller already holds every affected payer row and applies the
389+
* returned contributions once, so no payer ledger is written here.
394390
*/
395-
export async function changeWorkspaceStoragePayersInTx(
396-
tx: DbOrTx,
397-
changes: ChangeWorkspaceStoragePayerParams[]
398-
): Promise<ChangeWorkspaceStoragePayerResult[]> {
399-
return changeWorkspaceStoragePayers(tx, changes, false)
400-
}
401-
402391
async function changeWorkspaceStoragePayers(
403392
tx: DbOrTx,
404-
changes: ChangeWorkspaceStoragePayerParams[],
405-
deferPayerUpdates: boolean
393+
changes: ChangeWorkspaceStoragePayerParams[]
406394
): Promise<ChangeWorkspaceStoragePayerResult[]> {
407395
if (changes.length === 0) return []
408396

@@ -442,98 +430,19 @@ async function changeWorkspaceStoragePayers(
442430
}
443431
}
444432

445-
const payerByKey = new Map<string, BillingEntity>()
446-
const payerChangingWorkspaceIds: string[] = []
447-
for (const workspaceId of workspaceIds) {
433+
const payerChangingWorkspaceIds = workspaceIds.filter((workspaceId) => {
448434
const lockedWorkspace = workspaceById.get(workspaceId)
449435
const change = changesByWorkspaceId.get(workspaceId)
450-
if (!lockedWorkspace || !change) continue
451-
const oldPayer = getWorkspacePayer(lockedWorkspace)
452-
const newPayer = getWorkspacePayer(change)
453-
if (getPayerKey(oldPayer) === getPayerKey(newPayer)) continue
454-
payerChangingWorkspaceIds.push(workspaceId)
455-
payerByKey.set(getPayerKey(oldPayer), oldPayer)
456-
payerByKey.set(getPayerKey(newPayer), newPayer)
457-
}
436+
if (!lockedWorkspace || !change) return false
437+
return (
438+
getPayerKey(getWorkspacePayer(lockedWorkspace)) !== getPayerKey(getWorkspacePayer(change))
439+
)
440+
})
458441

459442
const exactBytesByWorkspaceId = await getExactWorkspaceStorageBytesBatch(
460443
tx,
461444
payerChangingWorkspaceIds
462445
)
463-
if (!deferPayerUpdates) {
464-
const payerUsageByKey = await lockStoragePayers(tx, payerByKey)
465-
const deltaByPayerKey = new Map<string, PayerStorageDelta>()
466-
467-
for (const workspaceId of payerChangingWorkspaceIds) {
468-
const lockedWorkspace = workspaceById.get(workspaceId)
469-
const change = changesByWorkspaceId.get(workspaceId)
470-
if (!lockedWorkspace || !change) continue
471-
const exactBytes = exactBytesByWorkspaceId.get(workspaceId) ?? 0
472-
const oldPayer = getWorkspacePayer(lockedWorkspace)
473-
const newPayer = getWorkspacePayer(change)
474-
const oldPayerKey = getPayerKey(oldPayer)
475-
const newPayerKey = getPayerKey(newPayer)
476-
477-
if (payerUsageByKey.get(newPayerKey) === null) {
478-
throw new Error(`Storage destination payer ${newPayerKey} not found`)
479-
}
480-
481-
const sourceDelta = deltaByPayerKey.get(oldPayerKey) ?? {
482-
incomingBytes: 0,
483-
outgoingBytes: 0,
484-
}
485-
sourceDelta.outgoingBytes += exactBytes
486-
if (!Number.isSafeInteger(sourceDelta.outgoingBytes)) {
487-
throw new Error(`Storage source payer ${oldPayerKey} delta exceeds the safe integer range`)
488-
}
489-
deltaByPayerKey.set(oldPayerKey, sourceDelta)
490-
491-
const destinationDelta = deltaByPayerKey.get(newPayerKey) ?? {
492-
incomingBytes: 0,
493-
outgoingBytes: 0,
494-
}
495-
destinationDelta.incomingBytes += exactBytes
496-
if (!Number.isSafeInteger(destinationDelta.incomingBytes)) {
497-
throw new Error(
498-
`Storage destination payer ${newPayerKey} delta exceeds the safe integer range`
499-
)
500-
}
501-
deltaByPayerKey.set(newPayerKey, destinationDelta)
502-
}
503-
504-
const nextUsageByKey = new Map<string, number>()
505-
for (const [payerKey, delta] of [...deltaByPayerKey.entries()].sort(([left], [right]) =>
506-
comparePayerKeys(left, right)
507-
)) {
508-
const currentUsage = payerUsageByKey.get(payerKey)
509-
if (currentUsage === null || currentUsage === undefined) {
510-
logger.warn('Storage source payer is missing during workspace payer batch change', {
511-
sourcePayer: payerKey,
512-
outgoingBytes: delta.outgoingBytes,
513-
})
514-
continue
515-
}
516-
517-
const usageAfterOutgoing = Math.max(0, currentUsage - delta.outgoingBytes)
518-
if (currentUsage < delta.outgoingBytes) {
519-
logger.warn(
520-
'Clamping drifted source storage aggregate during workspace payer batch change',
521-
{
522-
sourcePayer: payerKey,
523-
sourceUsage: currentUsage,
524-
outgoingBytes: delta.outgoingBytes,
525-
}
526-
)
527-
}
528-
const nextUsage = usageAfterOutgoing + delta.incomingBytes
529-
if (!Number.isSafeInteger(nextUsage)) {
530-
throw new Error(`Storage payer ${payerKey} exceeds the safe integer range`)
531-
}
532-
nextUsageByKey.set(payerKey, nextUsage)
533-
}
534-
535-
await setStoragePayerUsagesBatch(tx, payerByKey, nextUsageByKey)
536-
}
537446

538447
const resultsByWorkspaceId = new Map<string, ChangeWorkspaceStoragePayerResult>()
539448
const workspaceUpdates = workspaceIds.map((workspaceId) => {
@@ -796,9 +705,14 @@ export async function changeProjectStoragePayersInTx(
796705
tx: DbTransaction,
797706
changes: readonly ChangeProjectStoragePayerParams[]
798707
): Promise<ChangeProjectStoragePayerResult[]> {
708+
await lockProjects(
709+
tx,
710+
changes.map((change) => change.projectId)
711+
)
799712
return changeProjectStoragePayers(tx, changes, false)
800713
}
801714

715+
/** The caller holds every changed Project's lock. */
802716
async function changeProjectStoragePayers(
803717
tx: DbTransaction,
804718
changes: readonly ChangeProjectStoragePayerParams[],
@@ -809,7 +723,6 @@ async function changeProjectStoragePayers(
809723
if (changesById.size !== changes.length)
810724
throw new Error('Storage payer batch contains duplicate Project IDs')
811725
const projectIds = [...changesById.keys()].sort()
812-
for (const id of projectIds) await lockProject(tx, id)
813726
const owners = await tx
814727
.select({ id: project.id, ownerId: project.ownerId, organizationId: project.organizationId })
815728
.from(project)
@@ -905,6 +818,7 @@ async function applyPayerContributionTransfers(
905818
const usageByKey = await lockStoragePayers(tx, payerByKey)
906819
const deltas = new Map<string, PayerStorageDelta>()
907820
for (const result of contributions) {
821+
if (result.billableBytes === 0) continue
908822
const sourceKey = getPayerKey(result.oldPayer)
909823
const destinationKey = result.newPayer ? getPayerKey(result.newPayer) : null
910824
if (sourceKey === destinationKey) continue
@@ -967,7 +881,7 @@ export async function changeProjectAndWorkspaceStoragePayersInTx(
967881
.sort()
968882
if (new Set(projectIds).size !== projectIds.length)
969883
throw new Error('Storage lifecycle batch contains duplicate Project IDs')
970-
for (const id of projectIds) await lockProject(tx, id)
884+
await lockProjects(tx, projectIds)
971885
const projects = projectIds.length
972886
? await tx
973887
.select({
@@ -1043,7 +957,7 @@ export async function changeProjectAndWorkspaceStoragePayersInTx(
1043957
removal.billableBytes = bytes
1044958
}
1045959
const projectResults = await changeProjectStoragePayers(tx, changes.projectChanges, true)
1046-
const workspaceResults = await changeWorkspaceStoragePayers(tx, changes.workspaceChanges, true)
960+
const workspaceResults = await changeWorkspaceStoragePayers(tx, changes.workspaceChanges)
1047961
await applyPayerContributionTransfers(tx, [
1048962
...projectResults,
1049963
...workspaceResults,

0 commit comments

Comments
 (0)