Skip to content

Commit 22ff39f

Browse files
waleedlatif1waleedlatif1
authored
fix(knowledge): stop holding the connector table lock across the processing commit (#8191)
* fix(knowledge): check connector and knowledge base at the end of the processing commit * fix(knowledge): skip index writes for an inactive source and probe connector locks after the insert * test(knowledge): arm the pre-commit source check in the processing utils test
1 parent 230e42a commit 22ff39f

6 files changed

Lines changed: 405 additions & 101 deletions

File tree

‎.github/workflows/test-build.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,7 @@ jobs:
267267
lib/knowledge/__integration__/search-source-pagination.integration.ts
268268
lib/knowledge/__integration__/search-reference-batching.integration.ts
269269
lib/knowledge/__integration__/embedding-insert-batches.integration.ts
270+
lib/knowledge/__integration__/processing-lock-scope.integration.ts
270271
lib/knowledge/__integration__/connector-lifecycle-locks.integration.ts
271272
lib/knowledge/__integration__/connector-deferral.integration.ts
272273
lib/knowledge/__integration__/stored-document-recovery.integration.ts

‎apps/sim/app/api/knowledge/utils.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,8 @@ describe('Knowledge Utils', () => {
211211
secretProvenanceVersion: null,
212212
},
213213
])
214+
/** Pre-commit connector and knowledge-base activity check. */
215+
queueTableRows(schemaMock.document, [{ id: 'doc1' }])
214216
/** In-transaction active-document recheck. */
215217
queueTableRows(schemaMock.document, [{ id: 'doc1' }])
216218

‎apps/sim/lib/knowledge/__integration__/processing-lock-scope.integration.ts‎

Lines changed: 225 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,225 @@
1+
/** Relation locks held by the document processing commit while it writes embeddings. */
2+
import { mkdtempSync } from 'node:fs'
3+
import { rm } from 'node:fs/promises'
4+
import { tmpdir } from 'node:os'
5+
import path from 'node:path'
6+
import { db } from '@sim/db'
7+
import { document, embedding, knowledgeBase, organization, user, workspace } from '@sim/db/schema'
8+
import { generateId } from '@sim/utils/id'
9+
import { eq, inArray } from 'drizzle-orm'
10+
import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from 'vitest'
11+
12+
const fixtures = vi.hoisted(() => ({ root: '', process: vi.fn(), embeddings: vi.fn() }))
13+
vi.mock('@/lib/uploads/core/setup.server', () => ({
14+
get UPLOAD_DIR_SERVER() {
15+
return fixtures.root
16+
},
17+
}))
18+
vi.mock('@/lib/knowledge/documents/document-processor', () => ({
19+
processDocument: fixtures.process,
20+
}))
21+
vi.mock('@/lib/knowledge/embeddings', () => ({ generateEmbeddings: fixtures.embeddings }))
22+
23+
import { resolveBillingAttribution } from '@/lib/billing/core/billing-attribution'
24+
import * as embeddingClient from '@/lib/embeddings/client'
25+
import {
26+
createKnowledgeAclFixtureIds,
27+
seedKnowledgeAclFixture,
28+
} from '@/lib/knowledge/__integration__/seed-source-access-fixture'
29+
import { createContentSyncLease } from '@/lib/knowledge/connectors/sync-lock'
30+
import { addDocument } from '@/lib/knowledge/connectors/sync-persistence'
31+
import { processDocumentAsync } from '@/lib/knowledge/documents/service'
32+
33+
describe('document processing commit lock scope', () => {
34+
const ids = createKnowledgeAclFixtureIds()
35+
const probe = `fixture_lock_probe_${generateId().replaceAll('-', '')}`
36+
const chunks = Array.from({ length: 3 }, (_, index) => ({
37+
text: `Synthetic chunk ${index}`,
38+
metadata: { startIndex: index * 20, endIndex: index * 20 + 19 },
39+
}))
40+
const embeddingResult = {
41+
embeddings: chunks.map(() => Array(1536).fill(0.2)),
42+
billableTokens: 0,
43+
modelName: 'text-embedding-3-small',
44+
pricingId: 'text-embedding-3-small',
45+
}
46+
47+
beforeAll(async () => {
48+
fixtures.root = mkdtempSync(path.join(tmpdir(), 'sim-processing-lock-scope-'))
49+
await seedKnowledgeAclFixture(ids, { connectorType: 'google_drive' })
50+
vi.spyOn(embeddingClient, 'assertKnowledgeEmbeddingCapacity').mockResolvedValue(undefined)
51+
fixtures.process.mockResolvedValue({
52+
chunks,
53+
metadata: { chunkCount: chunks.length, tokenCount: 9, characterCount: 60 },
54+
})
55+
fixtures.embeddings.mockResolvedValue(embeddingResult)
56+
})
57+
58+
afterEach(async () => {
59+
await db.$client.unsafe(`DROP TRIGGER IF EXISTS ${probe} ON embedding`)
60+
await db.$client.unsafe(`DROP FUNCTION IF EXISTS ${probe}()`)
61+
})
62+
63+
afterAll(async () => {
64+
vi.restoreAllMocks()
65+
await db.delete(knowledgeBase).where(eq(knowledgeBase.id, ids.knowledgeBaseId))
66+
await db.delete(workspace).where(eq(workspace.id, ids.workspaceId))
67+
await db.delete(organization).where(eq(organization.id, ids.organizationId))
68+
await db.delete(user).where(inArray(user.id, [ids.aliceId, ids.bobId]))
69+
await rm(fixtures.root, { recursive: true, force: true })
70+
await db.$client.end()
71+
})
72+
73+
async function addConnectorDocument(externalId: string) {
74+
return addDocument(
75+
ids.knowledgeBaseId,
76+
ids.connectorId,
77+
'google_drive',
78+
{
79+
externalId,
80+
title: `${externalId}.txt`,
81+
content: 'Synthetic source text',
82+
mimeType: 'text/plain',
83+
contentHash: externalId,
84+
},
85+
{ userId: ids.aliceId, workspaceId: ids.workspaceId },
86+
undefined,
87+
'workspace',
88+
createContentSyncLease(ids.connectorId, ids.lockId)
89+
)
90+
}
91+
92+
/** Installs `body` as a BEFORE INSERT row trigger on embedding for one document. */
93+
async function installEmbeddingProbe(documentId: string, body: string) {
94+
await db.$client.unsafe(`CREATE FUNCTION ${probe}() RETURNS trigger LANGUAGE plpgsql AS $$
95+
BEGIN
96+
IF NEW.document_id = '${documentId}' THEN
97+
${body}
98+
END IF;
99+
RETURN NEW;
100+
END;
101+
$$`)
102+
await db.$client.unsafe(`CREATE TRIGGER ${probe} BEFORE INSERT ON embedding
103+
FOR EACH ROW EXECUTE FUNCTION ${probe}()`)
104+
}
105+
106+
/**
107+
* Fails any embedding insert statement whose backend holds a lock on
108+
* `knowledge_connector`. A statement-level AFTER trigger fires once the row
109+
* triggers and foreign-key checks of that statement have run, so it sees
110+
* every lock the insert itself took.
111+
*/
112+
async function installConnectorLockProbe() {
113+
await db.$client.unsafe(`CREATE FUNCTION ${probe}() RETURNS trigger LANGUAGE plpgsql AS $$
114+
BEGIN
115+
IF EXISTS (
116+
SELECT 1 FROM pg_locks
117+
WHERE pid = pg_backend_pid() AND relation = 'knowledge_connector'::regclass
118+
) THEN
119+
RAISE EXCEPTION 'embedding write holds a knowledge_connector lock';
120+
END IF;
121+
RETURN NULL;
122+
END;
123+
$$`)
124+
await db.$client.unsafe(`CREATE TRIGGER ${probe} AFTER INSERT ON embedding
125+
FOR EACH STATEMENT EXECUTE FUNCTION ${probe}()`)
126+
}
127+
128+
function billing() {
129+
return resolveBillingAttribution({ actorUserId: ids.aliceId, workspaceId: ids.workspaceId })
130+
}
131+
132+
it('holds no knowledge_connector lock while writing embeddings', async () => {
133+
const file = await addConnectorDocument('lock-scope-fixture')
134+
await installConnectorLockProbe()
135+
136+
const result = await processDocumentAsync(
137+
ids.knowledgeBaseId,
138+
file.documentId,
139+
file,
140+
{},
141+
await billing()
142+
)
143+
144+
expect(result).toEqual({ outcome: 'indexed' })
145+
expect(await db.select().from(document).where(eq(document.id, file.documentId))).toMatchObject([
146+
{ processingStatus: 'completed', chunkCount: 3 },
147+
])
148+
})
149+
150+
it.each([
151+
['connector', 'lock-scope-deleted-connector', 'knowledge_connector', () => ids.connectorId],
152+
['knowledge base', 'lock-scope-deleted-kb', 'knowledge_base', () => ids.knowledgeBaseId],
153+
])(
154+
'rolls back the embeddings when the %s is deleted during the embedding writes',
155+
async (_, externalId, table, id) => {
156+
const file = await addConnectorDocument(externalId)
157+
await installEmbeddingProbe(
158+
file.documentId,
159+
`UPDATE ${table} SET deleted_at = now() WHERE id = '${id()}';`
160+
)
161+
162+
const result = await processDocumentAsync(
163+
ids.knowledgeBaseId,
164+
file.documentId,
165+
file,
166+
{},
167+
await billing()
168+
)
169+
170+
expect(result).toEqual({ outcome: 'skipped', reason: 'superseded' })
171+
expect(
172+
await db
173+
.select({ id: embedding.id })
174+
.from(embedding)
175+
.where(eq(embedding.documentId, file.documentId))
176+
).toEqual([])
177+
for (const table of ['embedding_search', 'embedding_keyword_search']) {
178+
expect(
179+
await db.$client.unsafe(`SELECT id FROM ${table} WHERE document_id = $1`, [
180+
file.documentId,
181+
])
182+
).toEqual([])
183+
}
184+
expect(
185+
await db.select().from(document).where(eq(document.id, file.documentId))
186+
).toMatchObject([{ processingStatus: 'processing', chunkCount: 0 }])
187+
}
188+
)
189+
it.each([
190+
['connector', 'lock-scope-precheck-connector', 'knowledge_connector', () => ids.connectorId],
191+
['knowledge base', 'lock-scope-precheck-kb', 'knowledge_base', () => ids.knowledgeBaseId],
192+
])(
193+
'skips the index writes when the %s went inactive after the claim',
194+
async (_, externalId, table, id) => {
195+
const file = await addConnectorDocument(externalId)
196+
await installEmbeddingProbe(
197+
file.documentId,
198+
`RAISE EXCEPTION 'index writes ran for an inactive source';`
199+
)
200+
fixtures.embeddings.mockImplementationOnce(async () => {
201+
await db.$client.unsafe(`UPDATE ${table} SET deleted_at = now() WHERE id = $1`, [id()])
202+
return embeddingResult
203+
})
204+
205+
try {
206+
const result = await processDocumentAsync(
207+
ids.knowledgeBaseId,
208+
file.documentId,
209+
file,
210+
{},
211+
await billing()
212+
)
213+
expect(result).toEqual({ outcome: 'skipped', reason: 'superseded' })
214+
} finally {
215+
await db.$client.unsafe(`UPDATE ${table} SET deleted_at = NULL WHERE id = $1`, [id()])
216+
}
217+
expect(
218+
await db
219+
.select({ id: embedding.id })
220+
.from(embedding)
221+
.where(eq(embedding.documentId, file.documentId))
222+
).toEqual([])
223+
}
224+
)
225+
})

‎apps/sim/lib/knowledge/documents/document-indexing-usage.test.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -186,14 +186,15 @@ const DOC_DATA = {
186186

187187
/**
188188
* Re-arms the row sets one `processDocumentAsync` call consumes: the KB/document
189-
* context JOIN, the document secret-provenance row, and the in-transaction claim
190-
* re-check that lets the attempt commit.
189+
* context JOIN, the document secret-provenance row, the pre-commit source check,
190+
* and the in-transaction claim re-check that lets the attempt commit.
191191
*/
192192
function armDocumentReads(): void {
193193
dbChainMockFns.limit
194194
.mockResolvedValueOnce([PERSISTED_CONTEXT])
195195
.mockResolvedValueOnce([PERSISTED_PROVENANCE_ROW])
196196
.mockResolvedValueOnce([{ id: DOCUMENT_ID }])
197+
.mockResolvedValueOnce([{ id: DOCUMENT_ID }])
197198
}
198199

199200
/** The `sourceReference` of the single embedding charge recorded by call `index`. */

0 commit comments

Comments
 (0)