From 53b6499a43a6518c7689aef0ab17139b2317e221 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sat, 10 Oct 2026 15:51:13 -0700 Subject: [PATCH 1/2] 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. --- apps/sim/lib/knowledge/documents/service.ts | 19 ++++++++++++++++--- .../lib/knowledge/orchestration/connectors.ts | 2 +- apps/sim/lib/knowledge/service.ts | 4 +++- 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/apps/sim/lib/knowledge/documents/service.ts b/apps/sim/lib/knowledge/documents/service.ts index 707c75b65c3..719274e2629 100644 --- a/apps/sim/lib/knowledge/documents/service.ts +++ b/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) diff --git a/apps/sim/lib/knowledge/orchestration/connectors.ts b/apps/sim/lib/knowledge/orchestration/connectors.ts index 1569e8d6431..c75ebcc6481 100644 --- a/apps/sim/lib/knowledge/orchestration/connectors.ts +++ b/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 }) diff --git a/apps/sim/lib/knowledge/service.ts b/apps/sim/lib/knowledge/service.ts index 7e34ad16f89..a9d17d09b18 100644 --- a/apps/sim/lib/knowledge/service.ts +++ b/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 From 09adc679b937ba1c1a40b42b4a988633906599bc Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Sat, 10 Oct 2026 21:02:40 -0700 Subject: [PATCH 2/2] test(knowledge): uploads and hard deletes commit while an embedding holds the KB key-share lock --- .../workspace-import.integration.ts | 70 ++++++++++++++++++- 1 file changed, 69 insertions(+), 1 deletion(-) diff --git a/apps/sim/lib/knowledge/__integration__/workspace-import.integration.ts b/apps/sim/lib/knowledge/__integration__/workspace-import.integration.ts index c27bdf44438..c68e4bc4aff 100644 --- a/apps/sim/lib/knowledge/__integration__/workspace-import.integration.ts +++ b/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(knowledgeBaseId: string, work: () => Promise) { + const release = createDeferred() + const acquired = createDeferred() + 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 =