fix(tables): lock a table's unique values under one table key so wide schemas stay within the lock budget

Author: waleedlatif1

apps/sim/lib/table/import-data.ts

@@ -121,7 +121,7 @@ export async function bulkInsertImportBatch(
   const inserted = await db.transaction(async (trx) => {
     await guardBatch(trx, data.tableId, revalidate)
     if (getUniqueColumns(table.schema).length > 0) {
-      // Whole-column locks, not per-value: a batch is far more values than the value-lock cap.
+      // The whole-table unique lock, not per-value: a batch is far more values than the value-lock cap.
       await lockUniqueColumns(trx, table)
       const uniqueResult = await checkBatchUniqueConstraintsDb(
         data.tableId,

apps/sim/lib/table/rows/row-writes.integration.ts

@@ -13,6 +13,7 @@ import { tableTriggerMock } from '@sim/testing/mocks/table-trigger.mock'
 import { tableWorkflowColumnsMock } from '@sim/testing/mocks/table-workflow-columns.mock'
 import { sleep } from '@sim/utils/helpers'
 import { generateId } from '@sim/utils/id'
+import { sql } from 'drizzle-orm'
 import postgres from 'postgres'
 import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
@@ -21,8 +22,11 @@ vi.mock('@/lib/table/trigger', () => tableTriggerMock)
 vi.mock('@/lib/table/workflow-columns', () => tableWorkflowColumnsMock)
 vi.mock('@/lib/uploads/core/storage-service', () => storageServiceMock)
 
-import { bulkInsertImportBatch } from '@/lib/table/import-data'
+import { bulkInsertImportBatch, importReplaceRows } from '@/lib/table/import-data'
+import type { DbTransaction } from '@/lib/table/planner'
+import { acquireRowOrderLock } from '@/lib/table/rows/ordering'
 import { batchInsertRows, batchUpdateRows, insertRow, updateRow } from '@/lib/table/rows/service'
+import { lockUniqueColumns, lockUniqueValues } from '@/lib/table/rows/unique-locks'
 import { getTableById } from '@/lib/table/service'
 import { getOrCreateTableSnapshot } from '@/lib/table/snapshot-cache'
 import type { ColumnDefinition, TableDefinition } from '@/lib/table/types'
@@ -271,8 +275,13 @@ describe('table row writes against real PostgreSQL', () => {
         table.id,
         [
           () => insertEmail(table, 'dup@example.test'),
-          () =>
-            updateRow(
+          async () => {
+            // Start the edit only once the insert has passed its check and holds its locks.
+            expect(await waitForLockWaiters(table.id, { onOrderLock: 1, onValueLock: 0 })).toEqual({
+              onOrderLock: 1,
+              onValueLock: 0,
+            })
+            return updateRow(
               {
                 tableId: table.id,
                 rowId: `${table.id}-a`,
@@ -283,7 +292,8 @@ describe('table row writes against real PostgreSQL', () => {
               },
               table,
               'unique-race'
-            ),
+            )
+          },
         ],
         { onOrderLock: 1, onValueLock: 1 }
       )
@@ -292,6 +302,69 @@ describe('table row writes against real PostgreSQL', () => {
       expect(await storedCount(table.id, { email: 'dup@example.test' })).toBe(1)
     })
 
+    it('lets a replace that adds the first unique column wait out a writer on an older schema', async () => {
+      const table = await createTable([{ id: 'name', name: 'name', type: 'string' }])
+      // The writer resolved the table while it still had a unique column, so it holds the unique
+      // lock shared and will want the row-order lock next.
+      const stale: TableDefinition = {
+        ...table,
+        schema: { columns: [...table.schema.columns, ...uniqueColumns] },
+      }
+      let holdsUniqueLock!: () => void
+      const holding = new Promise<void>((resolve) => {
+        holdsUniqueLock = resolve
+      })
+      let release!: () => void
+      const released = new Promise<void>((resolve) => {
+        release = resolve
+      })
+      const writer = db.transaction(async (trx) => {
+        await lockUniqueValues(trx, stale, [{ email: 'stale@example.test' }])
+        holdsUniqueLock()
+        await released
+        await acquireRowOrderLock(trx, table.id)
+      })
+      await holding
+
+      const replace = importReplaceRows(
+        table,
+        [{ id: 'code', name: 'code', type: 'string', unique: true }],
+        { rows: [{ name: 'fresh', code: 'c-1' }], workspaceId },
+        'unique-race'
+      )
+      expect(await waitForLockWaiters(table.id, { onOrderLock: 0, onValueLock: 1 })).toEqual({
+        onOrderLock: 0,
+        onValueLock: 1,
+      })
+      release()
+
+      const results = await Promise.allSettled([writer, replace])
+      expect(results.map((result) => result.status)).toEqual(['fulfilled', 'fulfilled'])
+    })
+
+    it('holds a bounded number of locks however many unique columns a table has', async () => {
+      const wide: ColumnDefinition[] = Array.from({ length: 100 }, (_, i) => ({
+        id: `code_${i}`,
+        name: `code_${i}`,
+        type: 'string',
+        unique: true,
+      }))
+      const table = await createTable(wide)
+      const row = Object.fromEntries(wide.map((_, i) => [`code_${i}`, `value-${i}`]))
+      const advisoryLocksHeld = async (lock: (trx: DbTransaction) => Promise<void>) =>
+        db.transaction(async (trx) => {
+          await lock(trx)
+          const [{ held }] = await trx.execute<{ held: number }>(sql`SELECT count(*)::int AS held
+            FROM pg_locks WHERE locktype = 'advisory' AND pid = pg_backend_pid()`)
+          return held
+        })
+
+      expect(await advisoryLocksHeld((trx) => lockUniqueValues(trx, table, [row]))).toBe(1)
+      expect(await advisoryLocksHeld((trx) => lockUniqueColumns(trx, table))).toBe(1)
+      const narrow = { code_0: 'value-0', code_1: 'value-1' }
+      expect(await advisoryLocksHeld((trx) => lockUniqueValues(trx, table, [narrow]))).toBe(3)
+    })
+
     it('serializes a batch over the value-lock cap against a single insert of the same value', async () => {
       const table = await createTable(uniqueColumns)
       // More values than the per-transaction value-lock cap, so the batch locks the column instead.
@@ -331,7 +404,10 @@ describe('table row writes against real PostgreSQL', () => {
           () => insertEmail(table, 'dup@example.test'),
           async () => {
             // Start the import only once the insert has passed its check and holds its locks.
-            await waitForLockWaiters(table.id, { onOrderLock: 1, onValueLock: 0 })
+            expect(await waitForLockWaiters(table.id, { onOrderLock: 1, onValueLock: 0 })).toEqual({
+              onOrderLock: 1,
+              onValueLock: 0,
+            })
             return bulkInsertImportBatch(
               {
                 tableId: table.id,

apps/sim/lib/table/rows/service.ts

@@ -508,7 +508,7 @@ export async function replaceTableRows(
  * Capacity is NOT checked here (it would mean a billing-pool read inside the tx).
  * Callers gate it before opening the tx — see `replaceTableRows` and `importReplaceRows`.
  *
- * Takes the table's unique-column locks before its row-order lock, so a caller already holding the
+ * Takes the table's unique lock before its row-order lock, so a caller already holding the
  * row-order lock must take `lockUniqueColumns` first.
  */
 export async function replaceTableRowsWithTx(

apps/sim/lib/table/rows/unique-locks.ts

@@ -7,8 +7,9 @@
  * values it writes, so the second writer's check runs after the first commits and sees its row.
  * Writers of different values never wait on each other.
  *
- * Lock order, everywhere: the table's schema lock (when taken), then these locks, taken all at once
- * in one sorted order, then the table's row-order lock, then the definition row.
+ * Lock order, everywhere: the table's schema lock (when taken), then the table's unique lock and its
+ * value locks in sorted order, in one statement, then the table's row-order lock, then the
+ * definition row.
  */
 
 import { compareStrings } from '@sim/utils/string'
@@ -21,77 +22,72 @@ import { getUniqueColumns, uniqueValueKey } from '@/lib/table/validation'
 const UNIQUE_LOCK_TAG = 'user_table_unique_value'
 
 /**
- * Most value locks one transaction takes before it locks whole columns instead. Every advisory lock
- * holds a slot in the server's shared lock table, which is sized at `max_locks_per_transaction`
- * (64 by default) × connections and shared by every transaction, so concurrent writers each holding
- * more than their per-connection share can exhaust it and fail unrelated queries with `out of
- * shared memory`. Capping at that per-connection budget keeps ordinary writes on per-value locks
- * while larger batches serialize on the column.
+ * Most value locks one transaction takes before it locks the table's unique values as a whole.
+ * Every advisory lock holds a slot in the server's shared lock table, which is sized at
+ * `max_locks_per_transaction` (64 by default) × connections and shared by every transaction, so
+ * concurrent writers each holding more than their per-connection share can exhaust it and fail
+ * unrelated queries with `out of shared memory`. Capping at that per-connection budget keeps
+ * ordinary writes on per-value locks while larger batches serialize on the table.
  */
 const MAX_VALUE_LOCKS = 64
 
-function columnLockKey(tableId: string, columnId: string): string {
-  return `user_table_unique:${tableId}:${columnId}`
+/**
+ * One lock per table over all its unique columns: value writers hold it shared, and writers that
+ * cannot lock by value hold it exclusively. A single key bounds a transaction at one lock plus its
+ * value locks, however many unique columns the schema has.
+ */
+function tableLockKey(tableId: string): string {
+  return `user_table_unique:${tableId}`
 }
 
 /**
  * Locks the unique-column values `rows` will write, before their unique check. Pass `columnIds` to
  * lock only the unique columns a patch changes; values it leaves alone are already stored.
  *
- * Each written column takes a shared column lock, then an exclusive lock per value. The value key
- * is `uniqueValueKey`, so two values the check treats as equal always share a key; null cells
- * never conflict and take none.
- * A column falls back to one exclusive column lock when a value is an object or array (a `json`
- * column's check matches by containment, which no single key can express), and every column does
- * when the transaction would exceed {@link MAX_VALUE_LOCKS}. An exclusive column lock waits for,
- * and blocks, every shared holder, so the fallback still serializes against per-value writers.
+ * The table lock is taken shared, then an exclusive lock per value. The value key includes
+ * `uniqueValueKey`, so two values the check treats as equal always share a key; null cells never
+ * conflict and take none. The table lock is taken exclusively instead, with no value locks, when a
+ * value is an object or array (a `json` column's check matches by containment, which no single key
+ * can express) or when the transaction would exceed {@link MAX_VALUE_LOCKS}. An exclusive holder
+ * waits for, and blocks, every shared holder, so it still serializes against per-value writers.
  */
 export async function lockUniqueValues(
   trx: DbTransaction,
   table: TableDefinition,
   rows: readonly RowData[],
   columnIds?: ReadonlySet<string>
 ): Promise<void> {
-  const columnLocks = new Map<string, boolean>()
+  const tableKey = tableLockKey(table.id)
   const valueKeys = new Set<string>()
   for (const column of getUniqueColumns(table.schema)) {
     const columnId = getColumnId(column)
     if (columnIds && !columnIds.has(columnId)) continue
-    const columnKey = columnLockKey(table.id, columnId)
-    const keys: string[] = []
-    let columnByValue = true
     for (const row of rows) {
       const value = row[columnId]
       if (value === null || value === undefined) continue
-      if (typeof value === 'object') {
-        columnByValue = false
-        break
+      if (typeof value === 'object') return lockUniqueColumns(trx, table)
+      const key = `${tableKey}:${columnId}:${uniqueValueKey(value, column)}`
+      if (!valueKeys.has(key) && valueKeys.size === MAX_VALUE_LOCKS) {
+        return lockUniqueColumns(trx, table)
       }
-      keys.push(`${columnKey}:${uniqueValueKey(value, column)}`)
+      valueKeys.add(key)
     }
-    if (columnByValue && keys.length === 0) continue
-    columnLocks.set(columnKey, columnByValue)
-    if (columnByValue) for (const key of keys) valueKeys.add(key)
   }
+  if (valueKeys.size === 0) return
 
-  const byValue = valueKeys.size <= MAX_VALUE_LOCKS
-  const locks: AdvisoryXactLockRequest[] = [...columnLocks.entries()]
-    .sort(([a], [b]) => compareStrings(a, b))
-    .map(([key, columnByValue]) => ({ key, shared: byValue && columnByValue }))
-  if (byValue) {
-    for (const key of [...valueKeys].sort(compareStrings)) locks.push({ key, shared: false })
-  }
+  const locks: AdvisoryXactLockRequest[] = [{ key: tableKey, shared: true }]
+  for (const key of [...valueKeys].sort(compareStrings)) locks.push({ key, shared: false })
   await acquireAdvisoryXactLocks(trx, UNIQUE_LOCK_TAG, locks)
 }
 
 /**
- * Locks every unique column of `table` exclusively, for writers that replace the table's rows
- * wholesale and so conflict with any concurrent write of a unique value.
+ * Locks every unique value of `table` exclusively, for writers that replace the table's rows
+ * wholesale and so conflict with any concurrent write of a unique value. It locks even a table with
+ * no unique columns yet: a writer holding an older schema may still hold the lock shared, and a
+ * whole-table writer that adds a unique column must already hold it before the row-order lock.
  */
 export async function lockUniqueColumns(trx: DbTransaction, table: TableDefinition): Promise<void> {
-  const locks = getUniqueColumns(table.schema)
-    .map((column) => columnLockKey(table.id, getColumnId(column)))
-    .sort(compareStrings)
-    .map((key) => ({ key, shared: false }))
-  await acquireAdvisoryXactLocks(trx, UNIQUE_LOCK_TAG, locks)
+  await acquireAdvisoryXactLocks(trx, UNIQUE_LOCK_TAG, [
+    { key: tableLockKey(table.id), shared: false },
+  ])
 }