fix(files): lock the workspace row before inserting a workspace file

apps/sim/ee/workspace-forking/lib/copy/copy-files.ts

@@ -6,6 +6,7 @@ import { generateId } from '@sim/utils/id'
 import { and, eq, inArray, isNull, or } from 'drizzle-orm'
 import {
   incrementStorageUsageForBillingContextInTx,
+  lockWorkspaceStorageForMutationInTx,
   resolveStorageBillingContext,
 } from '@/lib/billing/storage'
 import type { DbOrTx } from '@/lib/db/types'
@@ -360,6 +361,7 @@ export async function executeForkFileBlobCopies(
         assertForkCopyActive(control)
         await db.transaction(async (tx) => {
           assertForkCopyActive(control)
+          await lockWorkspaceStorageForMutationInTx(tx, task.workspaceId)
           const [inserted] = await tx
             .insert(workspaceFiles)
             .values({

apps/sim/lib/billing/storage/tracking.ts

@@ -430,7 +430,17 @@ export async function applyStorageUsageDeltasInTx(
   return destinationUpdatedUsage
 }
 
-/** Locks the canonical workspace before existing-file mutations that later change its storage ledger. */
+/**
+ * Locks the canonical workspace row before a file mutation that changes its storage ledger.
+ *
+ * Every transaction that writes `workspace_files` and charges storage takes its locks in one order:
+ * the folder advisory lock (only for operations that take one), then this workspace row, then file
+ * rows, then the file writes (whose search-index trigger upserts the one-per-workspace
+ * `workspace_file_search_dispatch_queue` row), then the payer row. Call this before the first file insert or update: an insert that reaches
+ * the queue row first and only then waits on the workspace row deadlocks with a content write
+ * that holds the workspace row and waits on the queue row. Re-locking later in the same
+ * transaction, as the storage increment does, never waits.
+ */
 export async function lockWorkspaceStorageForMutationInTx(
   tx: DbOrTx,
   workspaceId: string

apps/sim/lib/uploads/contexts/workspace/__integration__/storage-lock-order.integration.ts

@@ -0,0 +1,167 @@
+/**
+ * Real PostgreSQL lock ordering between new-file finalization and a concurrent content write in
+ * the same workspace. Both reach the workspace row and the per-workspace search dispatch queue row
+ * (through the `workspace_files` search-index trigger), so they must take them in the same order.
+ */
+import { mkdtempSync } from 'node:fs'
+import { rm } from 'node:fs/promises'
+import { tmpdir } from 'node:os'
+import path from 'node:path'
+import { db, dbFor } from '@sim/db'
+import {
+  organization,
+  user,
+  workspace,
+  workspaceFileSearchDispatchQueue,
+  workspaceFiles,
+} from '@sim/db/schema'
+import { deleteWorkspaceFixture } from '@sim/db/testing/workspace-fixtures'
+import { createDeferred } from '@sim/testing/helpers/deferred'
+import { generateId } from '@sim/utils/id'
+import { eq, inArray, sql } from 'drizzle-orm'
+import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'
+
+const fixtureStorage = vi.hoisted(() => ({ root: '' }))
+vi.mock('@/lib/uploads/core/setup.server', () => ({
+  get UPLOAD_DIR_SERVER() {
+    return fixtureStorage.root
+  },
+}))
+
+import { lockWorkspaceStorageForMutationInTx } from '@/lib/billing/storage'
+import {
+  createKnowledgeAclFixtureIds,
+  seedKnowledgeAclFixture,
+} from '@/lib/knowledge/__integration__/seed-source-access-fixture'
+import {
+  generateWorkspaceFileKey,
+  registerUploadedWorkspaceFile,
+  uploadWorkspaceFile,
+} from '@/lib/uploads/contexts/workspace/workspace-file-manager'
+import { uploadFile } from '@/lib/uploads/core/storage-service'
+
+type FinalizeNewFile = (workspaceId: string, userId: string) => Promise<unknown>
+
+const finalizeByUpload: FinalizeNewFile = (workspaceId, userId) => {
+  const name = `upload-${generateId()}.txt`
+  return uploadWorkspaceFile(workspaceId, userId, Buffer.from(name), name, 'text/plain', {
+    notifyWorkspaceChange: false,
+  })
+}
+
+const finalizeByRegistration: FinalizeNewFile = async (workspaceId, userId) => {
+  const name = `registered-${generateId()}.txt`
+  const key = generateWorkspaceFileKey(workspaceId, name)
+  await uploadFile({
+    file: Buffer.from(name),
+    fileName: key,
+    contentType: 'text/plain',
+    context: 'workspace',
+    preserveKey: true,
+    customKey: key,
+    persistMetadata: false,
+  })
+  return registerUploadedWorkspaceFile({
+    workspaceId,
+    userId,
+    key,
+    originalName: name,
+    contentType: 'text/plain',
+  })
+}
+
+/** `db:push` never installs the migration-only search-index trigger that reaches the queue row. */
+const [{ migrated }] = await db.execute<{ migrated: boolean }>(sql`SELECT EXISTS (
+  SELECT 1 FROM pg_trigger
+  WHERE tgrelid = 'workspace_files'::regclass AND tgname = 'workspace_files_search_index_pending'
+) AS migrated`)
+/** The suite's own `afterAll` never runs when it is skipped, so release the probe connection here. */
+if (!migrated) await db.$client.end()
+
+describe.skipIf(!migrated)('workspace file storage lock order in PostgreSQL', () => {
+  const fixtures: ReturnType<typeof createKnowledgeAclFixtureIds>[] = []
+
+  beforeAll(() => {
+    fixtureStorage.root = mkdtempSync(path.join(tmpdir(), 'sim-file-lock-order-'))
+  })
+
+  afterAll(async () => {
+    for (const ids of fixtures) {
+      await deleteWorkspaceFixture(db, eq(workspace.id, ids.workspaceId))
+      await db.delete(organization).where(eq(organization.id, ids.organizationId))
+      await db.delete(user).where(inArray(user.id, [ids.aliceId, ids.bobId]))
+    }
+    await rm(fixtureStorage.root, { recursive: true, force: true })
+    await Promise.all([db.$client.end(), dbFor('cleanup').$client.end()])
+  })
+
+  /** Seeds a workspace with one file, which also creates the workspace's dispatch queue row. */
+  async function seedWorkspaceWithFile() {
+    const ids = createKnowledgeAclFixtureIds()
+    fixtures.push(ids)
+    await seedKnowledgeAclFixture(ids)
+    const name = `existing-${generateId()}.txt`
+    const existing = await uploadWorkspaceFile(
+      ids.workspaceId,
+      ids.aliceId,
+      Buffer.from(name),
+      name,
+      'text/plain',
+      { notifyWorkspaceChange: false }
+    )
+    const queued = await db
+      .select({ workspaceId: workspaceFileSearchDispatchQueue.workspaceId })
+      .from(workspaceFileSearchDispatchQueue)
+      .where(eq(workspaceFileSearchDispatchQueue.workspaceId, ids.workspaceId))
+    expect(queued).toHaveLength(1)
+    return { ...ids, existingFileId: existing.id }
+  }
+
+  it.each([
+    ['uploadWorkspaceFile', finalizeByUpload],
+    ['registerUploadedWorkspaceFile', finalizeByRegistration],
+  ])(
+    '%s commits alongside a content write that already holds the workspace row',
+    async (_name, finalize) => {
+      const fixture = await seedWorkspaceWithFile()
+      const ready = createDeferred<number>()
+      const finalizeIsWaiting = createDeferred<void>()
+
+      // The lock sequence of `updateWorkspaceFileContent`: workspace row, then a content write.
+      const contentWrite = db.transaction(async (tx) => {
+        await lockWorkspaceStorageForMutationInTx(tx, fixture.workspaceId)
+        const [connection] = await tx.execute<{ pid: number }>(sql`SELECT pg_backend_pid() AS pid`)
+        ready.resolve(connection.pid)
+        await finalizeIsWaiting.promise
+        await tx
+          .update(workspaceFiles)
+          .set({ contentUpdatedAt: new Date(Date.now() + 1000) })
+          .where(eq(workspaceFiles.id, fixture.existingFileId))
+      })
+      const contentWritePid = await ready.promise
+
+      const finalized = finalize(fixture.workspaceId, fixture.aliceId)
+      const settled = Promise.allSettled([contentWrite, finalized])
+      try {
+        await expect
+          .poll(
+            async () =>
+              (
+                await db.execute(
+                  sql`SELECT pid FROM pg_stat_activity
+                    WHERE ${contentWritePid} = ANY(pg_blocking_pids(pid))`
+                )
+              ).length,
+            { timeout: 4000 }
+          )
+          .toBe(1)
+      } finally {
+        finalizeIsWaiting.resolve()
+        await settled
+      }
+
+      const results = await settled
+      expect(results.filter((result) => result.status === 'rejected')).toEqual([])
+    }
+  )
+})

apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts

@@ -310,11 +310,13 @@ interface WorkspaceFileMetadataInsert {
  * Attempts one active workspace-file insert and reports the row that this call
  * created. Conflict losers receive `undefined` and must inspect the active key
  * in the same transaction before deciding whether the operation is idempotent.
+ * Locks the workspace row first so the insert follows the storage lock order.
  */
 async function insertWorkspaceFileMetadataInTx(
   tx: DbOrTx,
   metadata: WorkspaceFileMetadataInsert
 ): Promise<WorkspaceFileRow | undefined> {
+  await lockWorkspaceStorageForMutationInTx(tx, metadata.workspaceId)
   const [inserted] = await tx
     .insert(workspaceFiles)
     .values({