perf(knowledge): lock the knowledge base row FOR NO KEY UPDATE on document writes (#8909)

* perf(knowledge): lock the knowledge base row FOR NO KEY UPDATE on document writes

Document uploads, document hard deletes, connector creation, and knowledge
base restore locked the parent knowledge_base row FOR UPDATE. Every embedding
insert takes FOR KEY SHARE on that row through its foreign key and holds it
for the whole indexing transaction, and FOR UPDATE conflicts with KEY SHARE,
so these writes queued behind any indexing pass in the same knowledge base
(and vice versa). A hard delete could also deadlock against an indexing pass
on one of its own documents.

FOR NO KEY UPDATE still conflicts with itself, FOR SHARE (connector sync
saves and completion), FOR UPDATE (KB moves, tag mutations), and the
UPDATE/DELETE of the row, so every serialization these locks provided is
kept. No path under these locks writes knowledge_base.id, the only key
column (the other unique indexes are partial). This matches soft deletion
and connector cleanup, which already use FOR NO KEY UPDATE.

* test(knowledge): uploads and hard deletes commit while an embedding holds the KB key-share lock

Author: waleedlatif1

apps/sim/lib/knowledge/__integration__/workspace-import.integration.ts

@@ -17,6 +17,7 @@ import {
   workspace,
   workspaceFiles,
 } from '@sim/db/schema'
+import { createDeferred } from '@sim/testing/helpers/deferred'
 import { generateId } from '@sim/utils/id'
 import { and, eq, inArray, sql } from 'drizzle-orm'
 import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'
@@ -50,7 +51,11 @@ import { listKnowledgeChunks } from '@/lib/knowledge/application/chunks'
 import { searchKnowledge } from '@/lib/knowledge/application/search'
 import { KNOWLEDGE_DOCUMENT_PROCESSING_OUTBOX_EVENT } from '@/lib/knowledge/documents/processing-events'
 import { knowledgeDocumentProcessingOutboxHandlers } from '@/lib/knowledge/documents/processing-outbox-handler'
-import { createDocumentRecords, createSingleDocument } from '@/lib/knowledge/documents/service'
+import {
+  createDocumentRecords,
+  createSingleDocument,
+  hardDeleteDocuments,
+} from '@/lib/knowledge/documents/service'
 import { KNOWLEDGE_STORAGE_CLEANUP_EVENT } from '@/lib/knowledge/documents/storage-cleanup-event'
 import { uploadKnowledgeArtifact } from '@/lib/knowledge/documents/storage-upload'
 import {
@@ -88,6 +93,42 @@ async function sourceFile(
   )
 }
 
+/**
+ * Runs `work` while another transaction holds `FOR KEY SHARE` on the knowledge base row, as an
+ * in-flight embedding insert does through its foreign key, and requires it to settle well before
+ * that holder's idle timeout would release the lock.
+ */
+async function settlesWhileKeyShareHeld<T>(knowledgeBaseId: string, work: () => Promise<T>) {
+  const release = createDeferred<void>()
+  const acquired = createDeferred<void>()
+  const holding = db.transaction(async (tx) => {
+    await tx.execute(sql`SET LOCAL idle_in_transaction_session_timeout = '15s'`)
+    await tx
+      .select({ id: knowledgeBase.id })
+      .from(knowledgeBase)
+      .where(eq(knowledgeBase.id, knowledgeBaseId))
+      .for('key share')
+    acquired.resolve()
+    await release.promise
+  })
+  await acquired.promise
+  const running = Promise.allSettled([work()])
+  let settled = false
+  void running.then(() => {
+    settled = true
+  })
+  try {
+    await expect.poll(() => settled, { interval: 5, timeout: 5000 }).toBe(true)
+    const [outcome] = await running
+    if (outcome.status === 'rejected') throw outcome.reason
+    return outcome.value
+  } finally {
+    release.resolve()
+    await holding
+    await running
+  }
+}
+
 /** Drains real audit writes before asserting persisted rows, including duplicates from one operation. */
 function observeAbsenceAudits(workspaceId: string) {
   let submitted = 0
@@ -229,6 +270,33 @@ describe('durable workspace file import', () => {
     }
   })
 
+  it.each(['single', 'bulk'] as const)(
+    '%s admission and hard delete commit while an embedding holds the KB key-share lock',
+    async (mode) => {
+      const ids = await seed()
+      const source = await sourceFile(ids, 'ordinary source bytes')
+      const input = {
+        filename: source.name,
+        fileUrl: `/api/files/serve/${encodeURIComponent(source.key)}?context=workspace`,
+        fileSize: source.size,
+        mimeType: source.type,
+      }
+      const documentIds = await settlesWhileKeyShareHeld(ids.knowledgeBaseId, async () =>
+        mode === 'single'
+          ? [(await createSingleDocument(input, ids.knowledgeBaseId, generateId(), ids.aliceId)).id]
+          : (
+              await createDocumentRecords([input], ids.knowledgeBaseId, generateId(), ids.aliceId)
+            ).map((created) => created.documentId)
+      )
+      expect(documentIds).toHaveLength(1)
+      expect(
+        await settlesWhileKeyShareHeld(ids.knowledgeBaseId, () =>
+          hardDeleteDocuments(documentIds, generateId())
+        )
+      ).toBe(1)
+    }
+  )
+
   it('indexes the admitted snapshot after the source was deleted and temporary access would have expired', async () => {
     const ids = await seed()
     const content =

apps/sim/lib/knowledge/documents/service.ts

@@ -2546,7 +2546,15 @@ export async function createDocumentRecords(
   const { returnData, storageNotification, unrecordedCount } = await db.transaction(async (tx) => {
     let storageNotification: DocumentStorageNotification | null = null
 
-    await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`)
+    /**
+     * `FOR NO KEY UPDATE` still serializes this insert with other document writes,
+     * KB moves, and archive, but not with the `FOR KEY SHARE` that each embedding
+     * insert takes through its foreign key, so uploads never queue behind a
+     * long-running indexing transaction in the same knowledge base.
+     */
+    await tx.execute(
+      sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR NO KEY UPDATE`
+    )
 
     const kb = await tx
       .select({
@@ -3229,7 +3237,9 @@ export async function createSingleDocument(
       await claimKnowledgeUploadForAttachment(tx, options.uploadedArtifact.cleanupEventId)
     }
 
-    await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`)
+    await tx.execute(
+      sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR NO KEY UPDATE`
+    )
 
     const kb = await tx
       .select({
@@ -4332,6 +4342,9 @@ async function hardDeleteDocumentBatch(
      * Lock every parent KB in stable ID order before deleting document rows.
      * Normal inserts and KB moves take the same parent lock first, so the
      * workspace snapshots used for accounting cannot change mid-delete.
+     * `no key update` lets embedding inserts keep taking their foreign-key
+     * `KEY SHARE` lock, so a delete neither waits on indexing elsewhere in the
+     * KB nor deadlocks against an indexing pass on one of its own documents.
      */
     const knowledgeBaseIds = [
       ...new Set(documentsToDelete.map((doc) => doc.knowledgeBaseId)),
@@ -4351,7 +4364,7 @@ async function hardDeleteDocumentBatch(
         )
       )
       .orderBy(asc(knowledgeBase.id))
-      .for('update')
+      .for('no key update')
     const lockedKnowledgeBaseById = new Map(lockedKnowledgeBases.map((kb) => [kb.id, kb]))
     for (const doc of documentsToDelete) {
       const lockedKb = lockedKnowledgeBaseById.get(doc.knowledgeBaseId)

apps/sim/lib/knowledge/orchestration/connectors.ts

@@ -471,7 +471,7 @@ export async function performCreateKnowledgeConnector(
   try {
     created = await db.transaction(async (tx) => {
       if (owner.organizationId) await lockOrganizationSearchApproval(tx, owner.organizationId)
-      await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${kb.id} FOR UPDATE`)
+      await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${kb.id} FOR NO KEY UPDATE`)
 
       const activeKb = await tx
         .select({ id: knowledgeBase.id })

apps/sim/lib/knowledge/service.ts

@@ -1136,7 +1136,9 @@ export async function restoreKnowledgeBase(
     try {
       await db.transaction(async (tx) => {
         if (kb.organizationId) await lockOrganizationSearchApproval(tx, kb.organizationId)
-        await tx.execute(sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR UPDATE`)
+        await tx.execute(
+          sql`SELECT 1 FROM knowledge_base WHERE id = ${knowledgeBaseId} FOR NO KEY UPDATE`
+        )
 
         attemptedRestoreName = await generateRestoreName(kb.name, async (candidate) => {
           if (!kb.workspaceId) return false