diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index 591a7414b05..487c28ef1b7 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -277,6 +277,7 @@ jobs: lib/knowledge/__integration__/listing-continuation.integration.ts lib/knowledge/__integration__/member-scope-renewal.integration.ts lib/knowledge/__integration__/member-document-lifecycle.integration.ts + lib/knowledge/__integration__/connector-lease-pages.integration.ts lib/knowledge/__integration__/slack-empty-threads.integration.ts lib/knowledge/__integration__/kb-block-search.integration.ts lib/knowledge/__integration__/gitlab-workspace.integration.ts diff --git a/apps/sim/app/api/knowledge/connectors/member-sync/route.test.ts b/apps/sim/app/api/knowledge/connectors/member-sync/route.test.ts index c045b61d8c9..0dba1d286f2 100644 --- a/apps/sim/app/api/knowledge/connectors/member-sync/route.test.ts +++ b/apps/sim/app/api/knowledge/connectors/member-sync/route.test.ts @@ -101,6 +101,21 @@ describe('member sync scheduler owner routing', () => { ).toBe(true) }) + it('still dispatches due connectors when the stale observation sweep fails', async () => { + mocks.sweep.mockRejectedValue( + Object.assign(new Error('canceling statement due to lock timeout'), { code: '55P03' }) + ) + queueTableRows(schemaMock.knowledgeConnector, [ + { id: 'workspace-source', workspaceId: 'workspace-a', organizationId: null }, + ]) + const response = await GET(createMockRequest('GET')) + expect(response.status).toBe(200) + expect(mocks.dispatch).toHaveBeenCalledExactlyOnceWith( + 'workspace-source', + expect.objectContaining({ requireRunnable: true }) + ) + }) + it('preserves workspace dispatch and refuses absent or ambiguous ownership', async () => { queueTableRows(schemaMock.knowledgeConnector, [ { id: 'missing', workspaceId: null, organizationId: null }, diff --git a/apps/sim/app/api/knowledge/connectors/member-sync/route.ts b/apps/sim/app/api/knowledge/connectors/member-sync/route.ts index 10088e0cb14..6d1403f0106 100644 --- a/apps/sim/app/api/knowledge/connectors/member-sync/route.ts +++ b/apps/sim/app/api/knowledge/connectors/member-sync/route.ts @@ -1,6 +1,7 @@ import { db } from '@sim/db' import { knowledgeBase, knowledgeConnector, knowledgeConnectorMemberSyncLog } from '@sim/db/schema' import { createLogger } from '@sim/logger' +import { getErrorMessage } from '@sim/utils/errors' import { and, asc, eq, inArray, isNull, lte, type SQL, sql } from 'drizzle-orm' import { type NextRequest, NextResponse } from 'next/server' import { verifyCronAuth } from '@/lib/auth/internal' @@ -157,9 +158,16 @@ export const GET = withRouteHandler(async (request: NextRequest) => { logger.warn(`[${requestId}] Closed ${closedLogs.length} orphaned member sync log(s)`) } - const sweep = await sweepStaleMemberObservations(now) - if (sweep.members > 0) { - logger.warn(`[${requestId}] Swept observations of ${sweep.members} stale member(s)`, sweep) + /** Observation hygiene never holds back dispatch; an unfinished sweep resumes next tick. */ + try { + const sweep = await sweepStaleMemberObservations(now) + if (sweep.members > 0) { + logger.warn(`[${requestId}] Swept observations of ${sweep.members} stale member(s)`, sweep) + } + } catch (error) { + logger.error(`[${requestId}] Stale member observation sweep failed`, { + error: getErrorMessage(error), + }) } const dueConnectors = await db diff --git a/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts new file mode 100644 index 00000000000..b8d61052f05 --- /dev/null +++ b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts @@ -0,0 +1,1310 @@ +/** + * Real PostgreSQL coverage for the connector-lease ACL writers: every transaction that holds the + * connector row assigns at most one page of ACLs, a lease lost between pages stops the writes that + * follow, and an interrupted member-sync disable resumes to a disabled connector with every ACL + * revoked. The `document` ACL trigger installed by the migrations fires on every page. + */ +import { db } from '@sim/db' +import { + document, + embedding, + knowledgeConnector, + knowledgeConnectorMember, + knowledgeConnectorMemberSyncLog, + knowledgeDocumentObservation, + organization, + resourcePolicy, + user, + workspace, +} from '@sim/db/schema' +import { installProjectionSourceAcl } from '@sim/db/script-migrations/0021_embedding_search_connector' +import { generateId } from '@sim/utils/id' +import { and, eq, inArray, sql } from 'drizzle-orm' +import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const provider = vi.hoisted(() => ({ list: vi.fn(), get: vi.fn(), changes: vi.fn() })) +vi.mock('@/connectors/registry.server', () => ({ + CONNECTOR_REGISTRY: { + google_drive: { + id: 'google_drive', + name: 'Fixture Drive', + auth: { mode: 'oauth', provider: 'google-drive' }, + permissionScopedListing: { capFieldIds: [] }, + listDocuments: provider.list, + getDocument: provider.get, + getChangeCursor: async () => 'fixture-start', + listChanges: provider.changes, + }, + }, +})) + +import { resolveBillingAttribution } from '@/lib/billing/core/billing-attribution' +import { compileCredentialGroupWorkflowAccessPolicy } from '@/lib/credential-groups/application/workflow-access-policy' +import { + createKnowledgeAclFixtureIds, + seedKnowledgeAclFixture, + seedKnowledgeMemberFixture, +} from '@/lib/knowledge/__integration__/seed-source-access-fixture' +import * as connectorTokens from '@/lib/knowledge/connectors/access-token' +import * as memberAccess from '@/lib/knowledge/connectors/member-access' +import * as memberObservations from '@/lib/knowledge/connectors/member-observations' +import { + materializeDocumentAcls, + recordMemberObservations, + rewriteConnectorAcls, + sweepStaleMemberObservations, +} from '@/lib/knowledge/connectors/member-observations' +import { + executeMemberSync, + resumeMembershipRewrites, +} from '@/lib/knowledge/connectors/member-sync-engine' +import { executeSync } from '@/lib/knowledge/connectors/sync-engine' +import { PROJECTION_ROW_BATCH_SIZE } from '@/lib/knowledge/connectors/sync-limits' +import { + createMemberSyncLease, + type LeaseTransaction, + leaseTransaction, + SyncLockLostException, + stillHoldsMemberSyncLock, + stillHoldsSyncLock, +} from '@/lib/knowledge/connectors/sync-lock' +import * as syncPersistence from '@/lib/knowledge/connectors/sync-persistence' +import { + persistDocumentAcls, + restoreWorkspaceDocumentAcls, +} from '@/lib/knowledge/connectors/sync-persistence' + +const PAGE = 25 +const DOCUMENTS = 60 + +describe('connector lease ACL pages in PostgreSQL', () => { + let ids: ReturnType + let members: Awaited> + const alice = () => `u:${ids.aliceId}@fixture.test` + const bob = () => `u:${ids.bobId}@fixture.test` + + beforeAll(async () => { + vi.stubGlobal('fetch', async () => { + throw new Error('Unexpected provider request in the lease page fixture') + }) + vi.spyOn(memberAccess, 'mintKnowledgeConnectorMemberToken').mockResolvedValue({ + accessToken: 'fixture-token', + refreshed: false, + }) + /** Records, for every ACL assignment, the transaction that made it. */ + await db.execute(sql`CREATE TABLE IF NOT EXISTS lease_page_acl_writes ( + document_id text NOT NULL, connector_id text, xact text NOT NULL, + lock_timeout text NOT NULL, statement_timeout text NOT NULL + )`) + await db.execute( + sql.raw(`CREATE OR REPLACE FUNCTION log_lease_page_acl_write() RETURNS trigger + LANGUAGE plpgsql AS $$ BEGIN + INSERT INTO lease_page_acl_writes VALUES (NEW.id, NEW.connector_id, pg_current_xact_id()::text, + current_setting('lock_timeout'), current_setting('statement_timeout')); + RETURN NEW; + END $$`) + ) + await db.execute(sql`DROP TRIGGER IF EXISTS log_lease_page_acl_write ON document`) + await db.execute(sql`CREATE TRIGGER log_lease_page_acl_write AFTER UPDATE OF acl ON document + FOR EACH ROW EXECUTE FUNCTION log_lease_page_acl_write()`) + /** Records, for every projection row the document trigger rewrites, its table and transaction. */ + await db.execute(sql`CREATE TABLE IF NOT EXISTS lease_page_projection_writes ( + projection text NOT NULL, document_id text NOT NULL, xact text NOT NULL + )`) + await db.execute( + sql.raw(`CREATE OR REPLACE FUNCTION log_lease_page_projection_write() RETURNS trigger + LANGUAGE plpgsql AS $$ BEGIN + INSERT INTO lease_page_projection_writes VALUES (TG_TABLE_NAME, NEW.document_id, pg_current_xact_id()::text); + RETURN NEW; + END $$`) + ) + for (const projection of ['embedding_search', 'embedding_keyword_tin']) { + await db.execute( + sql.raw(`DROP TRIGGER IF EXISTS log_lease_page_projection_write ON ${projection}`) + ) + await db.execute( + sql.raw(`CREATE TRIGGER log_lease_page_projection_write AFTER UPDATE OF acl ON ${projection} + FOR EACH ROW EXECUTE FUNCTION log_lease_page_projection_write()`) + ) + } + }) + + beforeEach(async () => { + ids = createKnowledgeAclFixtureIds() + await seedKnowledgeAclFixture(ids, { connectorType: 'google_drive' }) + members = await seedKnowledgeMemberFixture(ids) + }) + + afterEach(async () => { + await db.execute(sql`DROP TRIGGER IF EXISTS fail_after_acl_writes ON document`) + await db.execute(sql`DELETE FROM lease_page_acl_writes`) + await db.execute(sql`DELETE FROM lease_page_projection_writes`) + await db.delete(workspace).where(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])) + }) + + afterAll(async () => { + await db.execute(sql`DROP TRIGGER IF EXISTS log_lease_page_acl_write ON document`) + await db.execute(sql`DROP FUNCTION IF EXISTS log_lease_page_acl_write()`) + await db.execute(sql`DROP FUNCTION IF EXISTS fail_after_acl_writes()`) + await db.execute(sql`DROP TABLE IF EXISTS lease_page_acl_writes`) + for (const projection of ['embedding_search', 'embedding_keyword_tin']) + await db.execute( + sql.raw(`DROP TRIGGER IF EXISTS log_lease_page_projection_write ON ${projection}`) + ) + await db.execute(sql`DROP FUNCTION IF EXISTS log_lease_page_projection_write()`) + await db.execute(sql`DROP TABLE IF EXISTS lease_page_projection_writes`) + vi.restoreAllMocks() + vi.unstubAllGlobals() + await db.$client.end() + }) + + const seedDocuments = async (connectorId: string, acl: string[], count = DOCUMENTS) => { + const rows = Array.from({ length: count }, (_unused, index) => ({ + id: generateId(), + knowledgeBaseId: ids.knowledgeBaseId, + connectorId, + externalId: `file-${String(index).padStart(3, '0')}`, + filename: `file-${index}`, + fileUrl: '', + fileSize: 0, + mimeType: 'text/plain', + processingStatus: 'completed', + contentHash: 'fixture-content', + /** Ten chunks each, so a page of projection rows holds 25 documents. */ + chunkCount: 10, + acl, + })) + await db.insert(document).values(rows) + return rows + } + + const storedAcls = async (connectorId: string) => + ( + await db + .select({ externalId: document.externalId, acl: document.acl }) + .from(document) + .where(eq(document.connectorId, connectorId)) + ).map((row) => row.acl) + + /** ACL assignments per transaction, for one connector. */ + const writesPerTransaction = async (connectorId: string) => + ( + await db.execute<{ writes: number }>(sql` + SELECT count(*)::int AS writes FROM lease_page_acl_writes + WHERE connector_id = ${connectorId} GROUP BY xact ORDER BY writes DESC`) + ).map((row) => row.writes) + + /** Every ACL assignment ran under the bounds of a connector-lease transaction. */ + const expectBounded = async (connectorId: string) => { + const bounds = await db.execute<{ lock: string; statement: string }>(sql` + SELECT DISTINCT lock_timeout AS lock, statement_timeout AS statement + FROM lease_page_acl_writes WHERE connector_id = ${connectorId}`) + expect([...bounds]).toEqual([{ lock: '15s', statement: '30s' }]) + } + + /** Takes the lease away once `held` pages have committed, as a reclaim between pages would. */ + const losingAfter = (held: number, inner: LeaseTransaction, lose: () => Promise) => { + let opened = 0 + const lossy: LeaseTransaction = async (write) => { + opened += 1 + if (opened === held + 1) await lose() + return inner(write) + } + return lossy + } + + const adminLease = () => ({ stillHeld: () => stillHoldsSyncLock(ids.connectorId, ids.lockId) }) + const reclaimAdmin = () => + db + .update(knowledgeConnector) + .set({ syncLockToken: generateId() }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + + describe('persistDocumentAcls', () => { + const changed = () => + new Map( + Array.from({ length: DOCUMENTS }, (_u, i) => [ + `file-${String(i).padStart(3, '0')}`, + [bob()], + ]) + ) + + it('assigns at most one page of ACLs per lease transaction, with unchanged results', async () => { + await seedDocuments(ids.connectorId, [alice()]) + + await expect( + persistDocumentAcls( + ids.connectorId, + changed(), + leaseTransaction(ids.connectorId, adminLease()) + ) + ).resolves.toEqual({ updated: DOCUMENTS, rejected: 0 }) + + expect(await writesPerTransaction(ids.connectorId)).toEqual([ + PAGE, + PAGE, + DOCUMENTS - 2 * PAGE, + ]) + await expectBounded(ids.connectorId) + expect((await storedAcls(ids.connectorId)).every((acl) => acl.join() === bob())).toBe(true) + }) + + it('writes nothing further once the lease is lost between pages, keeping the pages that landed', async () => { + await seedDocuments(ids.connectorId, [alice()]) + /** Page one is the evidence refresh (nothing unchanged), page two the first ACL batch. */ + const transaction = losingAfter( + 2, + leaseTransaction(ids.connectorId, adminLease()), + reclaimAdmin + ) + + await expect( + persistDocumentAcls(ids.connectorId, changed(), transaction) + ).rejects.toBeInstanceOf(SyncLockLostException) + + const acls = await storedAcls(ids.connectorId) + expect(acls.filter((acl) => acl.join() === bob())).toHaveLength(PAGE) + expect(acls.filter((acl) => acl.join() === alice())).toHaveLength(DOCUMENTS - PAGE) + }) + }) + + describe('search projection fan-out', () => { + const CHUNKS = 20 + + /** The production document trigger under test, whatever an earlier suite left installed. */ + beforeAll(async () => { + await installProjectionSourceAcl(db.$client) + }) + + /** Real chunks: the installed triggers create each chunk's search and keyword projection rows. */ + const seedChunks = async (documents: { id: string }[]) => { + const rows = documents.flatMap((entry) => + Array.from({ length: CHUNKS }, (_unused, chunkIndex) => ({ + id: generateId(), + knowledgeBaseId: ids.knowledgeBaseId, + documentId: entry.id, + chunkIndex, + chunkHash: `${entry.id}-${chunkIndex}`, + content: `lease page chunk ${chunkIndex}`, + contentLength: 20, + tokenCount: 4, + embedding: [1, ...Array(1535).fill(0)], + startOffset: 0, + endOffset: 20, + })) + ) + for (let offset = 0; offset < rows.length; offset += 200) + await db.insert(embedding).values(rows.slice(offset, offset + 200)) + /** + * Each chunk's projection rows, filled from its document as the backfill leaves them. The + * embedding insert writes the vector projection's row; the keyword projection's own sync + * trigger ships with the Tin migration, which a database without Tin skips. The fixture sets + * the filled state itself, whatever triggers an earlier suite left behind; what is under + * test is the document trigger that rewrites these rows. + */ + const chunkIds = sql.join( + documents.map((entry) => sql`${entry.id}`), + sql`, ` + ) + await db.execute(sql` + UPDATE embedding_search p SET enabled = true, connector_id = d.connector_id, acl = d.acl + FROM document d WHERE d.id = p.document_id AND d.id IN (${chunkIds})`) + await db.execute(sql` + INSERT INTO embedding_keyword_tin (id, knowledge_base_id, document_id, enabled, content, connector_id, acl) + SELECT e.id, e.knowledge_base_id, e.document_id, true, e.content, d.connector_id, d.acl + FROM embedding e JOIN document d ON d.id = e.document_id WHERE e.document_id IN (${chunkIds}) + ON CONFLICT (id) DO UPDATE SET enabled = true, connector_id = EXCLUDED.connector_id, acl = EXCLUDED.acl`) + /** Only writes made by the code under test are counted. */ + await db.execute(sql`DELETE FROM lease_page_projection_writes`) + await db + .update(document) + .set({ chunkCount: CHUNKS }) + .where( + inArray( + document.id, + documents.map((entry) => entry.id) + ) + ) + } + + /** Every projection row of the connector's documents, with its ACL and the document's, as text. */ + const projectionAcls = async (connectorId: string) => + db.execute<{ projection: string; acl: string | null; expected: string }>(sql` + SELECT 'embedding_search' AS projection, array_to_string(p.acl, ',') AS acl, + array_to_string(d.acl, ',') AS expected + FROM embedding_search p JOIN document d ON d.id = p.document_id WHERE d.connector_id = ${connectorId} + UNION ALL + SELECT 'embedding_keyword_tin', array_to_string(p.acl, ','), array_to_string(d.acl, ',') + FROM embedding_keyword_tin p JOIN document d ON d.id = p.document_id WHERE d.connector_id = ${connectorId}`) + + /** Projection rows each transaction rewrote, per table. */ + const projectionRowsPerTransaction = async () => + ( + await db.execute<{ rows: number }>(sql` + SELECT count(*)::int AS rows FROM lease_page_projection_writes + GROUP BY projection, xact ORDER BY rows DESC`) + ).map((row) => row.rows) + + it('mirrors a paged ACL write onto every projection row, one page of rows per transaction', async () => { + const seeded = await seedDocuments(ids.connectorId, [alice()], 30) + await seedChunks(seeded) + const before = [...(await projectionAcls(ids.connectorId))] + expect(before).toHaveLength(2 * 30 * CHUNKS) + + await expect( + persistDocumentAcls( + ids.connectorId, + new Map(seeded.map((row) => [row.externalId, [bob()]])), + leaseTransaction(ids.connectorId, adminLease()) + ) + ).resolves.toEqual({ updated: 30, rejected: 0 }) + + const after = [...(await projectionAcls(ids.connectorId))] + expect(after.filter((row) => row.acl !== bob() || row.expected !== bob())).toEqual([]) + const perTransaction = await projectionRowsPerTransaction() + expect(perTransaction.reduce((total, rows) => total + rows, 0)).toBe(2 * 30 * CHUNKS) + expect(Math.max(...perTransaction)).toBeLessThanOrEqual(PROJECTION_ROW_BATCH_SIZE) + }) + + it('hides a members connector across its projection rows, one page of rows per transaction', async () => { + const seeded = await seedDocuments(members.connectorId, [alice()], 30) + await seedChunks(seeded) + + await expect( + rewriteConnectorAcls(members.connectorId, [], { + lease: { + stillHeld: () => stillHoldsMemberSyncLock(members.connectorId, members.runId), + }, + }) + ).resolves.toBe(true) + + const after = [...(await projectionAcls(members.connectorId))] + expect(after).toHaveLength(2 * 30 * CHUNKS) + expect(after.filter((row) => row.acl !== '' || row.expected !== '')).toEqual([]) + expect(Math.max(...(await projectionRowsPerTransaction()))).toBeLessThanOrEqual( + PROJECTION_ROW_BATCH_SIZE + ) + }) + }) + + describe('fence-last pages', () => { + /** + * A processing commit holds a document row for its whole write. A page waiting on it must not + * hold the connector row meanwhile, or every heartbeat, edit and reclaim queues behind it. + */ + it('waits on a locked document row without holding any lock on the connector table', async () => { + const [locked] = await seedDocuments(ids.connectorId, [alice()], 1) + let release!: () => void + let held!: () => void + const holding = new Promise((resolve) => { + held = resolve + }) + const holder = db.transaction(async (tx) => { + await tx + .select({ id: document.id }) + .from(document) + .where(eq(document.id, locked.id)) + .for('update') + held() + await new Promise((resolve) => { + release = resolve + }) + }) + await holding + const write = persistDocumentAcls( + ids.connectorId, + new Map([[locked.externalId, [bob()]]]), + leaseTransaction(ids.connectorId, adminLease()) + ) + try { + let waiter: number | undefined + for (let attempt = 0; attempt < 100 && waiter === undefined; attempt++) { + const [row] = await db.execute<{ pid: number }>(sql` + SELECT pid FROM pg_stat_activity + WHERE wait_event_type = 'Lock' AND datname = current_database() + AND (query ILIKE 'update "document" set "acl"%' + OR query ILIKE 'select "id", "chunk_count" from "document"%for update')`) + waiter = row?.pid + if (waiter === undefined) await new Promise((resolve) => setImmediate(resolve)) + } + expect(waiter).toBeDefined() + const connectorLocks = await db.execute<{ mode: string }>(sql` + SELECT mode FROM pg_locks + WHERE pid = ${waiter} AND relation = 'knowledge_connector'::regclass`) + expect([...connectorLocks]).toEqual([]) + } finally { + release() + await holder + } + await expect(write).resolves.toEqual({ updated: 1, rejected: 0 }) + expect((await storedAcls(ids.connectorId)).map((acl) => acl.join())).toEqual([bob()]) + }, 30_000) + + /** + * Pages are sized from a read taken without a lock; a reprocess can change a document's chunks + * before the page writes. The page locks its documents and rereads their counts, so what it + * writes still fits one page of projection rows. + */ + it('resizes a page whose documents gained chunks after it was planned', async () => { + const seeded = await seedDocuments(ids.connectorId, [alice()], 3) + const reprocess = losingAfter(1, leaseTransaction(ids.connectorId, adminLease()), () => + db + .update(document) + .set({ chunkCount: 1_000 }) + .where(eq(document.connectorId, ids.connectorId)) + ) + + await expect( + persistDocumentAcls( + ids.connectorId, + new Map(seeded.map((row) => [row.externalId, [bob()]])), + reprocess + ) + ).resolves.toEqual({ updated: 3, rejected: 0 }) + + expect(await writesPerTransaction(ids.connectorId)).toEqual([1, 1, 1]) + }) + + /** Observation changes that must commit with their ACLs are paged by projection rows too. */ + it('rewrites a membership page by projection rows, a large document alone', async () => { + const seeded = await seedDocuments(members.connectorId, [], 3) + const [member] = members.members + await recordMemberObservations( + db, + member.id, + seeded.map((row) => row.id), + members.runId + ) + const huge = [...seeded].sort((a, b) => (a.id < b.id ? -1 : 1))[1] + await db.update(document).set({ chunkCount: 1_000 }).where(eq(document.id, huge.id)) + await db + .update(knowledgeConnectorMember) + .set({ listingCheckpoint: { kind: 'membership', cursor: null, removeMember: false } }) + .where(eq(knowledgeConnectorMember.id, member.id)) + + await expect( + resumeMembershipRewrites({ + connectorId: members.connectorId, + runId: members.runId, + deadlineAt: Date.now() + 60_000, + lease: createMemberSyncLease(members.connectorId, members.runId), + }) + ).resolves.toBe(true) + + expect(await writesPerTransaction(members.connectorId)).toEqual([1, 1, 1]) + expect( + (await storedAcls(members.connectorId)).every((acl) => acl.join() === member.subjectToken) + ).toBe(true) + }) + + /** The member engine's ACL pages prove the lease last as well. */ + it('holds no connector lock while a member ACL page waits on a locked document row', async () => { + const seeded = await seedDocuments(members.connectorId, [], 1) + const [member] = members.members + await recordMemberObservations(db, member.id, [seeded[0].id], members.runId) + await db + .update(knowledgeConnectorMember) + .set({ listingCheckpoint: { kind: 'membership', cursor: null, removeMember: false } }) + .where(eq(knowledgeConnectorMember.id, member.id)) + let release!: () => void + let held!: () => void + const holding = new Promise((resolve) => { + held = resolve + }) + const holder = db.transaction(async (tx) => { + await tx + .select({ id: document.id }) + .from(document) + .where(eq(document.id, seeded[0].id)) + .for('update') + held() + await new Promise((resolve) => { + release = resolve + }) + }) + await holding + const rewrite = resumeMembershipRewrites({ + connectorId: members.connectorId, + runId: members.runId, + deadlineAt: Date.now() + 60_000, + lease: createMemberSyncLease(members.connectorId, members.runId), + }) + try { + let waiter: number | undefined + for (let attempt = 0; attempt < 200 && waiter === undefined; attempt++) { + const [row] = await db.execute<{ pid: number }>(sql` + SELECT pid FROM pg_stat_activity + WHERE wait_event_type = 'Lock' AND datname = current_database() + AND (query ILIKE 'update "document" set "acl"%' + OR query ILIKE 'select "id", "chunk_count" from "document"%for update')`) + waiter = row?.pid + if (waiter === undefined) await new Promise((resolve) => setImmediate(resolve)) + } + expect(waiter).toBeDefined() + const connectorLocks = await db.execute<{ mode: string }>(sql` + SELECT mode FROM pg_locks + WHERE pid = ${waiter} AND relation = 'knowledge_connector'::regclass`) + expect([...connectorLocks]).toEqual([]) + } finally { + release() + await holder + } + await expect(rewrite).resolves.toBe(true) + expect((await storedAcls(members.connectorId)).map((acl) => acl.join())).toEqual([ + member.subjectToken, + ]) + }, 30_000) + + /** One page is bounded by projection rows; a document larger than the cap still lands, alone. */ + it('gives a document larger than one page of projection rows a page alone', async () => { + const seeded = await seedDocuments(ids.connectorId, [alice()], 3) + await db.update(document).set({ chunkCount: 1_000 }).where(eq(document.id, seeded[1].id)) + await db + .update(document) + .set({ chunkCount: 200 }) + .where(inArray(document.id, [seeded[0].id, seeded[2].id])) + + await expect( + persistDocumentAcls( + ids.connectorId, + new Map(seeded.map((row) => [row.externalId, [bob()]])), + leaseTransaction(ids.connectorId, adminLease()) + ) + ).resolves.toEqual({ updated: 3, rejected: 0 }) + + expect(await writesPerTransaction(ids.connectorId)).toEqual([1, 1, 1]) + }) + }) + + describe('restoreWorkspaceDocumentAcls', () => { + beforeEach(async () => { + await db + .update(knowledgeConnector) + .set({ accessMode: 'workspace' }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + }) + + it('restores a hidden connector one page per transaction and reports every document', async () => { + await seedDocuments(ids.connectorId, []) + + await expect( + restoreWorkspaceDocumentAcls( + ids.connectorId, + leaseTransaction(ids.connectorId, adminLease()) + ) + ).resolves.toEqual({ restored: DOCUMENTS, finished: true }) + + expect(await writesPerTransaction(ids.connectorId)).toEqual([ + PAGE, + PAGE, + DOCUMENTS - 2 * PAGE, + ]) + await expectBounded(ids.connectorId) + expect((await storedAcls(ids.connectorId)).every((acl) => acl.join() === 'ws')).toBe(true) + await expect( + restoreWorkspaceDocumentAcls( + ids.connectorId, + leaseTransaction(ids.connectorId, adminLease()) + ) + ).resolves.toEqual({ restored: 0, finished: true }) + }) + + it('finishes a pending rewrite before a sync completes, outside the completion transaction', async () => { + await db + .update(knowledgeConnector) + .set({ status: 'active', syncLockToken: null, accessRewritePending: true }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + const seeded = await seedDocuments(ids.connectorId, []) + await db + .update(document) + .set({ sourceSeenAt: sql`now() + interval '1 day'`, storageKey: sql`'kb/fixture/' || id` }) + .where(eq(document.connectorId, ids.connectorId)) + provider.list.mockResolvedValue({ + documents: seeded.map((row) => ({ + externalId: row.externalId, + title: row.filename, + content: '', + contentDeferred: true, + contentHash: 'fixture-content', + mimeType: 'text/plain', + })), + hasMore: false, + }) + const token = vi + .spyOn(connectorTokens, 'resolveConnectorAccessToken') + .mockResolvedValue({ accessToken: 'fixture-token' } as never) + try { + const result = await executeSync(ids.connectorId, { + billingAttribution: await resolveBillingAttribution({ + actorUserId: ids.aliceId, + workspaceId: ids.workspaceId, + }), + }) + expect(result.error).toBeUndefined() + } finally { + token.mockRestore() + } + + expect((await storedAcls(ids.connectorId)).every((acl) => acl.join() === 'ws')).toBe(true) + expect(await writesPerTransaction(ids.connectorId)).toEqual([ + PAGE, + PAGE, + DOCUMENTS - 2 * PAGE, + ]) + await expectBounded(ids.connectorId) + const [connector] = await db + .select({ + status: knowledgeConnector.status, + syncLockToken: knowledgeConnector.syncLockToken, + accessRewritePending: knowledgeConnector.accessRewritePending, + }) + .from(knowledgeConnector) + .where(eq(knowledgeConnector.id, ids.connectorId)) + expect(connector).toEqual({ + status: 'active', + syncLockToken: null, + accessRewritePending: false, + }) + }) + + /** A restore stops between pages at its deadline; a later walk resumes from what is still off. */ + it('stops a restore between pages at its deadline and resumes it on the next walk', async () => { + await seedDocuments(ids.connectorId, []) + let clock = Date.now() + const now = vi.spyOn(Date, 'now').mockImplementation(() => clock) + let pages = 0 + try { + const partial = await restoreWorkspaceDocumentAcls( + ids.connectorId, + leaseTransaction(ids.connectorId, adminLease()), + { + deadlineAt: clock + 1_000, + beforePage: async () => { + pages += 1 + /** The window read and the first page run; the budget passes before the next window. */ + if (pages === 3) clock += 2_000 + }, + } + ) + expect(partial).toEqual({ restored: PAGE, finished: false }) + } finally { + now.mockRestore() + } + await expect( + restoreWorkspaceDocumentAcls( + ids.connectorId, + leaseTransaction(ids.connectorId, adminLease()) + ) + ).resolves.toEqual({ restored: DOCUMENTS - PAGE, finished: true }) + expect((await storedAcls(ids.connectorId)).every((acl) => acl.join() === 'ws')).toBe(true) + }) + + /** A sync whose restore ran out of budget keeps the flag and comes back at once to finish it. */ + it('keeps the pending rewrite of a sync whose restore did not finish', async () => { + await db + .update(knowledgeConnector) + .set({ status: 'active', syncLockToken: null, accessRewritePending: true }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + const seeded = await seedDocuments(ids.connectorId, []) + await db + .update(document) + .set({ sourceSeenAt: sql`now() + interval '1 day'`, storageKey: sql`'kb/fixture/' || id` }) + .where(eq(document.connectorId, ids.connectorId)) + provider.list.mockResolvedValue({ + documents: seeded.map((row) => ({ + externalId: row.externalId, + title: row.filename, + content: '', + contentDeferred: true, + contentHash: 'fixture-content', + mimeType: 'text/plain', + })), + hasMore: false, + }) + const token = vi + .spyOn(connectorTokens, 'resolveConnectorAccessToken') + .mockResolvedValue({ accessToken: 'fixture-token' } as never) + const original = syncPersistence.restoreWorkspaceDocumentAcls + const restore = vi + .spyOn(syncPersistence, 'restoreWorkspaceDocumentAcls') + .mockImplementationOnce((connectorId, transaction, options) => + original(connectorId, transaction, { ...options, deadlineAt: Date.now() - 1 }) + ) + const billing = await resolveBillingAttribution({ + actorUserId: ids.aliceId, + workspaceId: ids.workspaceId, + }) + const state = async () => { + const [row] = await db + .select({ + accessRewritePending: knowledgeConnector.accessRewritePending, + nextSyncAt: knowledgeConnector.nextSyncAt, + }) + .from(knowledgeConnector) + .where(eq(knowledgeConnector.id, ids.connectorId)) + return row + } + try { + expect((await executeSync(ids.connectorId, { billingAttribution: billing })).error).toBe( + undefined + ) + /** The sync hands its own run budget to the restore. */ + expect(restore).toHaveBeenCalledWith( + ids.connectorId, + expect.any(Function), + expect.objectContaining({ deadlineAt: expect.any(Number) }) + ) + const unfinished = await state() + expect(unfinished?.accessRewritePending).toBe(true) + expect(unfinished?.nextSyncAt?.getTime()).toBeLessThanOrEqual(Date.now()) + expect((await storedAcls(ids.connectorId)).every((acl) => acl.length === 0)).toBe(true) + + expect((await executeSync(ids.connectorId, { billingAttribution: billing })).error).toBe( + undefined + ) + expect((await state())?.accessRewritePending).toBe(false) + expect((await storedAcls(ids.connectorId)).every((acl) => acl.join() === 'ws')).toBe(true) + } finally { + token.mockRestore() + restore.mockRestore() + } + }) + + /** An admin connector still hiding its documents lists nothing until the walk is done. */ + it('lists nothing and keeps the pending rewrite while an admin hide is unfinished', async () => { + provider.list.mockClear() + await db + .update(knowledgeConnector) + .set({ + accessMode: 'admin', + status: 'active', + syncLockToken: null, + accessRewritePending: true, + }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + await seedDocuments(ids.connectorId, ['ws']) + const token = vi + .spyOn(connectorTokens, 'resolveConnectorAccessToken') + .mockResolvedValue({ accessToken: 'fixture-token' } as never) + const hide = vi + .spyOn(memberObservations, 'rewriteConnectorAcls') + .mockImplementationOnce(async (_connectorId, _target, options) => { + expect(options?.deadlineAt).toEqual(expect.any(Number)) + return false + }) + try { + const result = await executeSync(ids.connectorId, { + billingAttribution: await resolveBillingAttribution({ + actorUserId: ids.aliceId, + workspaceId: ids.workspaceId, + }), + }) + expect(result.error).toBeUndefined() + expect(hide).toHaveBeenCalledOnce() + expect(provider.list).not.toHaveBeenCalled() + const [row] = await db + .select({ + accessRewritePending: knowledgeConnector.accessRewritePending, + syncLockToken: knowledgeConnector.syncLockToken, + }) + .from(knowledgeConnector) + .where(eq(knowledgeConnector.id, ids.connectorId)) + expect(row).toEqual({ accessRewritePending: true, syncLockToken: null }) + } finally { + token.mockRestore() + hide.mockRestore() + } + }) + + /** Only a pending switch leaves workspace documents off the workspace ACL; a healthy sync never walks them. */ + it('does not walk a workspace connector without a pending rewrite', async () => { + await db + .update(knowledgeConnector) + .set({ status: 'active', syncLockToken: null, accessRewritePending: false }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + const seeded = await seedDocuments(ids.connectorId, ['ws']) + await db + .update(document) + .set({ sourceSeenAt: sql`now() + interval '1 day'`, storageKey: sql`'kb/fixture/' || id` }) + .where(eq(document.connectorId, ids.connectorId)) + provider.list.mockResolvedValue({ + documents: seeded.map((row) => ({ + externalId: row.externalId, + title: row.filename, + content: '', + contentDeferred: true, + contentHash: 'fixture-content', + mimeType: 'text/plain', + })), + hasMore: false, + }) + const token = vi + .spyOn(connectorTokens, 'resolveConnectorAccessToken') + .mockResolvedValue({ accessToken: 'fixture-token' } as never) + const restore = vi.spyOn(syncPersistence, 'restoreWorkspaceDocumentAcls') + try { + const result = await executeSync(ids.connectorId, { + billingAttribution: await resolveBillingAttribution({ + actorUserId: ids.aliceId, + workspaceId: ids.workspaceId, + }), + }) + expect(result.error).toBeUndefined() + expect(restore).not.toHaveBeenCalled() + } finally { + token.mockRestore() + restore.mockRestore() + } + }) + + it('restores nothing once the connector has left workspace mode', async () => { + await seedDocuments(ids.connectorId, []) + await db + .update(knowledgeConnector) + .set({ accessMode: 'admin' }) + .where(eq(knowledgeConnector.id, ids.connectorId)) + + await expect( + restoreWorkspaceDocumentAcls( + ids.connectorId, + leaseTransaction(ids.connectorId, adminLease()) + ) + ).resolves.toEqual({ restored: 0, finished: true }) + expect((await storedAcls(ids.connectorId)).every((acl) => acl.length === 0)).toBe(true) + }) + + it('stops at the first page after the lease is lost', async () => { + await seedDocuments(ids.connectorId, []) + + await expect( + restoreWorkspaceDocumentAcls( + ids.connectorId, + losingAfter(1, leaseTransaction(ids.connectorId, adminLease()), reclaimAdmin) + ) + ).rejects.toBeInstanceOf(SyncLockLostException) + + expect((await storedAcls(ids.connectorId)).filter((acl) => acl.length > 0)).toHaveLength(PAGE) + }) + }) + + describe('rewriteConnectorAcls', () => { + const memberLease = () => ({ + stillHeld: () => stillHoldsMemberSyncLock(members.connectorId, members.runId), + }) + + it('hides a members connector one page per transaction', async () => { + await seedDocuments(members.connectorId, [alice()]) + + await expect( + rewriteConnectorAcls(members.connectorId, [], { lease: memberLease() }) + ).resolves.toBe(true) + + expect(await writesPerTransaction(members.connectorId)).toEqual([ + PAGE, + PAGE, + DOCUMENTS - 2 * PAGE, + ]) + await expectBounded(members.connectorId) + expect((await storedAcls(members.connectorId)).every((acl) => acl.length === 0)).toBe(true) + }) + + it('stops writing once the lease is lost between pages', async () => { + await seedDocuments(members.connectorId, [alice()]) + let pages = 0 + + await expect( + rewriteConnectorAcls(members.connectorId, [], { + lease: memberLease(), + beforeBatch: async () => { + pages += 1 + /** The first beat precedes the window read, the second the first page. */ + if (pages === 3) + await db + .update(knowledgeConnector) + .set({ memberSyncLockToken: generateId() }) + .where(eq(knowledgeConnector.id, members.connectorId)) + }, + }) + ).rejects.toBeInstanceOf(SyncLockLostException) + + const acls = await storedAcls(members.connectorId) + expect(acls.filter((acl) => acl.length === 0)).toHaveLength(PAGE) + expect(acls.filter((acl) => acl.length > 0)).toHaveLength(DOCUMENTS - PAGE) + }) + }) + + describe('stale member sweep', () => { + it('defers a connector whose row a member run holds and still sweeps the others', async () => { + const busy = members + const idle = await seedKnowledgeMemberFixture(ids) + for (const fixture of [busy, idle]) { + await db + .update(knowledgeConnector) + .set({ + status: 'active', + memberSyncStatus: 'idle', + memberSyncLockToken: null, + syncIntervalMinutes: 60, + lastMemberSyncAt: new Date(), + }) + .where(eq(knowledgeConnector.id, fixture.connectorId)) + /** Enrolled long enough ago that never having listed makes them stale. */ + await db + .update(knowledgeConnectorMember) + .set({ createdAt: new Date(Date.now() - 3 * 24 * 60 * 60 * 1000) }) + .where(eq(knowledgeConnectorMember.connectorId, fixture.connectorId)) + const seeded = await seedDocuments( + fixture.connectorId, + fixture.members.map((member) => member.subjectToken).sort(), + 3 + ) + for (const member of fixture.members) + await recordMemberObservations( + db, + member.id, + seeded.map((row) => row.id), + fixture.runId + ) + } + const observed = async (connectorId: string) => + ( + await db + .select({ id: knowledgeDocumentObservation.memberId }) + .from(knowledgeDocumentObservation) + .innerJoin( + knowledgeConnectorMember, + eq(knowledgeConnectorMember.id, knowledgeDocumentObservation.memberId) + ) + .where(eq(knowledgeConnectorMember.connectorId, connectorId)) + ).length + + /** A member page of a running sync holds the busy connector's row for longer than a sweep page waits. */ + let release!: () => void + let locked!: () => void + const held = new Promise((resolve) => { + locked = resolve + }) + const holder = db.transaction(async (tx) => { + await tx + .select({ id: knowledgeConnector.id }) + .from(knowledgeConnector) + .where(eq(knowledgeConnector.id, busy.connectorId)) + .for('update') + locked() + await new Promise((resolve) => { + release = resolve + }) + }) + await held + try { + const result = await sweepStaleMemberObservations(new Date(Date.now() + 1_000)) + expect(result.members).toBe(idle.members.length) + } finally { + release() + await holder + } + + expect(await observed(busy.connectorId)).toBe(busy.members.length * 3) + expect(await observed(idle.connectorId)).toBe(0) + expect((await storedAcls(idle.connectorId)).every((acl) => acl.length === 0)).toBe(true) + }, 30_000) + }) + + describe('resumeMembershipRewrites', () => { + it('rematerialises a changed member one page per lease transaction', async () => { + const seeded = await seedDocuments(members.connectorId, []) + const [member] = members.members + await recordMemberObservations( + db, + member.id, + seeded.map((row) => row.id), + members.runId + ) + await db + .update(knowledgeConnectorMember) + .set({ listingCheckpoint: { kind: 'membership', cursor: null, removeMember: false } }) + .where(eq(knowledgeConnectorMember.id, member.id)) + + await expect( + resumeMembershipRewrites({ + connectorId: members.connectorId, + runId: members.runId, + deadlineAt: Date.now() + 60_000, + lease: createMemberSyncLease(members.connectorId, members.runId), + }) + ).resolves.toBe(true) + + expect( + (await storedAcls(members.connectorId)).every((acl) => acl.join() === member.subjectToken) + ).toBe(true) + expect(await writesPerTransaction(members.connectorId)).toEqual([ + PAGE, + PAGE, + DOCUMENTS - 2 * PAGE, + ]) + await expectBounded(members.connectorId) + }) + }) + + describe('member listing materialisation', () => { + beforeEach(async () => { + provider.list.mockReset() + provider.get.mockReset() + provider.changes.mockReset() + await db + .insert(resourcePolicy) + .values({ + id: generateId(), + workspaceId: ids.workspaceId, + resourceType: 'credential_group', + resourceId: members.groupId, + document: compileCredentialGroupWorkflowAccessPolicy({ + credentialGroupId: members.groupId, + allowedWorkflowIds: [], + }), + createdBy: ids.aliceId, + updatedBy: ids.aliceId, + }) + .onConflictDoNothing() + await memberAccess.grantKnowledgeConnectorCredentialAccess( + { + workspaceId: ids.workspaceId, + credentialGroupId: members.groupId, + credentialGroupOptionId: members.optionId, + connectorId: members.connectorId, + }, + ids.aliceId + ) + await db + .update(knowledgeConnector) + .set({ status: 'active', memberSyncStatus: 'idle', memberSyncLockToken: null }) + .where(eq(knowledgeConnector.id, members.connectorId)) + }) + + it('observes and materialises a large listing one page per lease transaction', async () => { + const seeded = await seedDocuments(members.connectorId, []) + /** Content this run already read stays unchanged, so only visibility is written. */ + await db + .update(document) + .set({ sourceSeenAt: sql`now() + interval '1 day'` }) + .where(eq(document.connectorId, members.connectorId)) + provider.list.mockResolvedValue({ + documents: seeded.map((row) => ({ + externalId: row.externalId, + title: row.filename, + content: '', + contentDeferred: true, + contentHash: 'fixture-content', + mimeType: 'text/plain', + })), + hasMore: false, + }) + + const result = await executeMemberSync(members.connectorId, { + billingAttribution: await resolveBillingAttribution({ + actorUserId: ids.aliceId, + workspaceId: ids.workspaceId, + }), + }) + + expect(result.error).toBeUndefined() + expect(result.membersCompleted).toBe(2) + expect(provider.get).not.toHaveBeenCalled() + const tokens = members.members.map((member) => member.subjectToken).sort() + expect( + (await storedAcls(members.connectorId)).every((acl) => acl.join() === tokens.join()) + ).toBe(true) + const perTransaction = await writesPerTransaction(members.connectorId) + expect(perTransaction.reduce((total, writes) => total + writes, 0)).toBe(2 * DOCUMENTS) + expect(Math.max(...perTransaction)).toBeLessThanOrEqual(PAGE) + await expectBounded(members.connectorId) + }) + + it('rematerialises what a change feed withdrew one page per lease transaction', async () => { + const tokens = members.members.map((member) => member.subjectToken).sort() + const seeded = await seedDocuments(members.connectorId, tokens) + for (const member of members.members) + await recordMemberObservations( + db, + member.id, + seeded.map((row) => row.id), + members.runId + ) + await db + .update(knowledgeConnectorMember) + .set({ + changeCursor: 'fixture-start', + lastCompleteListingAt: new Date(), + memberSyncedThrough: new Date(), + }) + .where(eq(knowledgeConnectorMember.connectorId, members.connectorId)) + provider.changes.mockResolvedValue({ + changes: seeded.map((row) => ({ kind: 'removed', externalId: row.externalId })), + hasMore: false, + nextCursor: 'fixture-drained', + }) + + const result = await executeMemberSync(members.connectorId, { + billingAttribution: await resolveBillingAttribution({ + actorUserId: ids.aliceId, + workspaceId: ids.workspaceId, + }), + }) + + expect(result.error).toBeUndefined() + expect(result.observationsRemoved).toBe(2 * DOCUMENTS) + expect((await storedAcls(members.connectorId)).every((acl) => acl.length === 0)).toBe(true) + const perTransaction = await writesPerTransaction(members.connectorId) + expect(perTransaction.reduce((total, writes) => total + writes, 0)).toBe(2 * DOCUMENTS) + expect(Math.max(...perTransaction)).toBeLessThanOrEqual(PAGE) + await expectBounded(members.connectorId) + }) + }) + + describe('member sync disable', () => { + const billing = () => + resolveBillingAttribution({ actorUserId: ids.aliceId, workspaceId: ids.workspaceId }) + const connectorState = async () => { + const [row] = await db + .select({ + memberSyncStatus: knowledgeConnector.memberSyncStatus, + memberSyncLockToken: knowledgeConnector.memberSyncLockToken, + }) + .from(knowledgeConnector) + .where(eq(knowledgeConnector.id, members.connectorId)) + return row + } + + beforeEach(async () => { + /** The option the connector synced through is gone, and no run holds the lease. */ + await db + .update(knowledgeConnector) + .set({ + credentialGroupOptionId: null, + memberSyncStatus: 'idle', + memberSyncLockToken: null, + }) + .where(eq(knowledgeConnector.id, members.connectorId)) + }) + + it('resumes an interrupted disable and ends disabled with every ACL revoked', async () => { + const tokens = members.members.map((member) => member.subjectToken).sort() + const seeded = await seedDocuments(members.connectorId, tokens) + await recordMemberObservations( + db, + members.members[0].id, + seeded.map((row) => row.id), + members.runId + ) + /** The third page fails, as a statement timeout would: two pages have committed. */ + await db.execute( + sql.raw(`CREATE OR REPLACE FUNCTION fail_after_acl_writes() RETURNS trigger + LANGUAGE plpgsql AS $$ BEGIN + IF (SELECT count(*) FROM lease_page_acl_writes WHERE connector_id = NEW.connector_id) >= 30 + THEN RAISE EXCEPTION 'fixture statement failure'; END IF; + RETURN NEW; + END $$`) + ) + await db.execute( + sql`CREATE TRIGGER fail_after_acl_writes BEFORE UPDATE OF acl ON document FOR EACH ROW + WHEN (NEW.connector_id = ${sql.raw(`'${members.connectorId}'`)}) + EXECUTE FUNCTION fail_after_acl_writes()` + ) + + const interrupted = await executeMemberSync(members.connectorId, { + billingAttribution: await billing(), + }) + + expect(interrupted.error).toBeTruthy() + expect((await connectorState())?.memberSyncStatus).not.toBe('disabled') + const midway = await storedAcls(members.connectorId) + expect(midway.filter((acl) => acl.length === 0)).toHaveLength(2 * PAGE) + /** A document not yet revoked keeps only the grant it had; nothing is broadened. */ + expect( + midway.filter((acl) => acl.length > 0).every((acl) => acl.join() === tokens.join()) + ).toBe(true) + /** Members are suspended first, so a rematerialisation in the meantime grants nobody. */ + const suspended = await db + .select({ status: knowledgeConnectorMember.status }) + .from(knowledgeConnectorMember) + .where(eq(knowledgeConnectorMember.connectorId, members.connectorId)) + expect(suspended.every((row) => row.status === 'suspended')).toBe(true) + + await db.execute(sql`DROP TRIGGER fail_after_acl_writes ON document`) + const unrevoked = seeded.at(-1)!.id + expect( + await leaseTransaction(members.connectorId)((tx) => + materializeDocumentAcls(members.connectorId, [unrevoked], tx) + ) + ).toBe(1) + const [rematerialized] = await db + .select({ acl: document.acl }) + .from(document) + .where(eq(document.id, unrevoked)) + expect(rematerialized.acl).toEqual([]) + const resumed = await executeMemberSync(members.connectorId, { + billingAttribution: await billing(), + }) + + expect(resumed.skipReason).toBe('connector_not_syncable') + expect(await connectorState()).toEqual({ + memberSyncStatus: 'disabled', + memberSyncLockToken: null, + }) + expect((await storedAcls(members.connectorId)).every((acl) => acl.length === 0)).toBe(true) + expect(Math.max(...(await writesPerTransaction(members.connectorId)))).toBeLessThanOrEqual( + PAGE + ) + await expectBounded(members.connectorId) + const [{ granted }] = await db + .select({ granted: sql`count(*)::int` }) + .from(document) + .where( + and(eq(document.connectorId, members.connectorId), sql`cardinality(${document.acl}) > 0`) + ) + expect(granted).toBe(0) + }) + + it('closes a run whose disable of a removed option fails as an ordinary failure', async () => { + /** The option id is set but no longer exists, which membership reconciliation reports. */ + await db + .update(knowledgeConnector) + .set({ credentialGroupOptionId: generateId() }) + .where(eq(knowledgeConnector.id, members.connectorId)) + await seedDocuments( + members.connectorId, + members.members.map((member) => member.subjectToken).sort() + ) + await db.execute( + sql.raw(`CREATE OR REPLACE FUNCTION fail_after_acl_writes() RETURNS trigger + LANGUAGE plpgsql AS $$ BEGIN RAISE EXCEPTION 'fixture statement failure'; END $$`) + ) + await db.execute( + sql`CREATE TRIGGER fail_after_acl_writes BEFORE UPDATE OF acl ON document FOR EACH ROW + WHEN (NEW.connector_id = ${sql.raw(`'${members.connectorId}'`)}) + EXECUTE FUNCTION fail_after_acl_writes()` + ) + + const failed = await executeMemberSync(members.connectorId, { + billingAttribution: await billing(), + }) + + expect(failed.error).toBeTruthy() + const state = await connectorState() + expect(state?.memberSyncStatus).toBe('error') + expect(state?.memberSyncLockToken).toBeNull() + const [log] = await db + .select({ status: knowledgeConnectorMemberSyncLog.status }) + .from(knowledgeConnectorMemberSyncLog) + .where(eq(knowledgeConnectorMemberSyncLog.connectorId, members.connectorId)) + expect(log?.status).toBe('failed') + + await db.execute(sql`DROP TRIGGER fail_after_acl_writes ON document`) + await executeMemberSync(members.connectorId, { billingAttribution: await billing() }) + expect((await connectorState())?.memberSyncStatus).toBe('disabled') + expect((await storedAcls(members.connectorId)).every((acl) => acl.length === 0)).toBe(true) + }) + }) +}) diff --git a/apps/sim/lib/knowledge/__integration__/migration-fixture.ts b/apps/sim/lib/knowledge/__integration__/migration-fixture.ts index 5d59ab207e4..bd6db49e409 100644 --- a/apps/sim/lib/knowledge/__integration__/migration-fixture.ts +++ b/apps/sim/lib/knowledge/__integration__/migration-fixture.ts @@ -57,7 +57,8 @@ export async function createEnterpriseSearchMigrationFixture(databaseUrl: string id text PRIMARY KEY, external_id text, connector_id text, knowledge_base_id text, tag1 text, tag2 text, tag3 text, tag4 text, tag5 text, tag6 text, tag7 text, acl text[] NOT NULL DEFAULT '{ws}', storage_key text, - user_excluded boolean NOT NULL DEFAULT false, archived_at timestamp + user_excluded boolean NOT NULL DEFAULT false, archived_at timestamp, + chunk_count integer NOT NULL DEFAULT 0 ); CREATE INDEX doc_connector_id_idx ON document(connector_id); CREATE TABLE knowledge_connector ( diff --git a/apps/sim/lib/knowledge/access/predicate.postgres.test.ts b/apps/sim/lib/knowledge/access/predicate.postgres.test.ts index 06ef356fc36..3e105c6eba4 100644 --- a/apps/sim/lib/knowledge/access/predicate.postgres.test.ts +++ b/apps/sim/lib/knowledge/access/predicate.postgres.test.ts @@ -22,6 +22,7 @@ vi.mock('@/connectors/registry.server', () => ({ CONNECTOR_REGISTRY: {} })) const { drizzle } = await import('drizzle-orm/postgres-js') const schema = await import('@sim/db/schema') const { persistDocumentAcls } = await import('@/lib/knowledge/connectors/sync-persistence') +const { leaseTransaction } = await import('@/lib/knowledge/connectors/sync-lock') const { mergeMirroredAcls, hideUnlistedDocuments } = await import( '@/lib/knowledge/connectors/mirrored-acls' ) @@ -750,7 +751,9 @@ describe.runIf(Boolean(databaseUrl))('knowledge ACLs in PostgreSQL', () => { const space = 'g:confluence:tenant:space' const page = 'g:confluence:tenant:page' const input = new Map([['page', { acl: [space], requirements: [[page]] }]]) - expect(await persistDocumentAcls('admin', input, executor)).toEqual({ updated: 1, rejected: 0 }) + expect( + await persistDocumentAcls('admin', input, leaseTransaction('admin', undefined, executor)) + ).toEqual({ updated: 1, rejected: 0 }) expect(await readable([space, page], 'persisted')).toBe(true) expect(await readable([page], 'persisted')).toBe(false) await connection.unsafe("UPDATE document SET acl = string_to_array($1, E'\\n')", [page]) @@ -760,7 +763,7 @@ describe.runIf(Boolean(databaseUrl))('knowledge ACLs in PostgreSQL', () => { "UPDATE document SET acl_verified_at = statement_timestamp() - interval '25 hours'" ) expect(await readable([space, page], 'persisted')).toBe(false) - await persistDocumentAcls('admin', input, executor) + await persistDocumentAcls('admin', input, leaseTransaction('admin', undefined, executor)) expect(await readable([space, page], 'persisted')).toBe(true) const [stored] = await connection.unsafe( "SELECT jsonb_typeof(acl_requirements) AS shape, acl_requirements FROM document WHERE id = 'persisted'" @@ -844,10 +847,15 @@ describe.runIf(Boolean(databaseUrl))('knowledge ACLs in PostgreSQL', () => { {} ) if (step === 'unlisted') hideUnlistedDocuments(merged.acls, ['shared-file']) - await persistDocumentAcls('admin', merged.acls, executor, { - unresolvedExternalIds: merged.unresolvedExternalIds, - generationStartedAt, - }) + await persistDocumentAcls( + 'admin', + merged.acls, + leaseTransaction('admin', undefined, executor), + { + unresolvedExternalIds: merged.unresolvedExternalIds, + generationStartedAt, + } + ) const [stored] = await connection.unsafe( "SELECT to_jsonb(acl) AS acl, acl_verified_at FROM document WHERE id = 'shared'" ) @@ -882,10 +890,15 @@ describe.runIf(Boolean(databaseUrl))('knowledge ACLs in PostgreSQL', () => { [verifiedAt] ) const executor = drizzle(connection, { schema }) - const result = await persistDocumentAcls('admin', new Map([['file', []]]), executor, { - unresolvedExternalIds: new Set(['file']), - generationStartedAt, - }) + const result = await persistDocumentAcls( + 'admin', + new Map([['file', []]]), + leaseTransaction('admin', undefined, executor), + { + unresolvedExternalIds: new Set(['file']), + generationStartedAt, + } + ) expect(result.updated).toBe(preserved ? 0 : 1) const [stored] = await connection.unsafe( "SELECT to_jsonb(acl) AS acl, acl_verified_at FROM document WHERE id = 'boundary'" diff --git a/apps/sim/lib/knowledge/connectors/detachment.ts b/apps/sim/lib/knowledge/connectors/detachment.ts index fee3c8e4d04..65add4459d9 100644 --- a/apps/sim/lib/knowledge/connectors/detachment.ts +++ b/apps/sim/lib/knowledge/connectors/detachment.ts @@ -24,10 +24,10 @@ import { import type { DbOrTx } from '@/lib/db/types' import { removeDrainedConnector } from '@/lib/knowledge/connectors/deletion' import { revokeKnowledgeConnectorCredentialAccess } from '@/lib/knowledge/connectors/member-access' +import { PROJECTION_ROW_BATCH_SIZE } from '@/lib/knowledge/connectors/sync-limits' export const KNOWLEDGE_CONNECTOR_DETACH_EVENT = 'knowledge.connector.detach' const DOCUMENT_BATCH_SIZE = 100 -const PROJECTION_ROW_BATCH_SIZE = 250 const MAX_BATCHES_PER_RUN = 4 const RUN_BUDGET_MS = 30_000 /** How often a detachment paused on a deleted knowledge base checks for its restore or purge. */ diff --git a/apps/sim/lib/knowledge/connectors/member-observations.test.ts b/apps/sim/lib/knowledge/connectors/member-observations.test.ts index 6abdc39e2e8..3f161739ff8 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.test.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.test.ts @@ -20,17 +20,21 @@ import { db } from '@sim/db' import { inArray } from 'drizzle-orm' import { applyMemberDocumentLifecycle, + materializeDocumentAcls, + pagesByProjectionRows, + rematerializeDocumentAcls, removeUnseenMemberObservations, renewMemberObservationsInScopes, rewriteConnectorAcls, staleMemberWindowMs, sweepStaleMemberObservations, + writeProjectionPages, } from '@/lib/knowledge/connectors/member-observations' import { MEMBER_OBSERVATION_STALE_AFTER_HOURS, MEMBER_TOMBSTONE_RECONCILE_PAGES_PER_RUN, } from '@/lib/knowledge/connectors/sync-limits' -import { SyncLockLostException } from '@/lib/knowledge/connectors/sync-lock' +import { type LeaseTransaction, SyncLockLostException } from '@/lib/knowledge/connectors/sync-lock' import { ConnectorSyncDeletionGuardError, hardDeleteDocuments, @@ -45,23 +49,46 @@ describe('removeUnseenMemberObservations', () => { resetDbChainMock() }) it('rematerializes bounded batches before reading more absent observations', async () => { + const candidates = (count: number, prefix: string) => + Array.from({ length: count }, (_, i) => ({ documentId: `${prefix}-${i}` })) + queueTableRows(schemaMock.knowledgeDocumentObservation, candidates(25, 'd')) + queueTableRows(schemaMock.knowledgeDocumentObservation, candidates(1, 'last')) dbChainMockFns.returning - .mockResolvedValueOnce(Array.from({ length: 500 }, (_, i) => ({ documentId: `d-${i}` }))) - .mockResolvedValueOnce([{ documentId: 'last' }]) - .mockResolvedValueOnce([]) + .mockResolvedValueOnce(candidates(25, 'd')) + .mockResolvedValueOnce(candidates(1, 'last')) const onRemoved = vi.fn(async (_ids: string[]) => undefined) await expect( removeUnseenMemberObservations(db, 'member', 'generation', onRemoved) - ).resolves.toEqual({ removed: 500, finished: false }) + ).resolves.toEqual({ removed: 25, finished: false }) await expect( removeUnseenMemberObservations(db, 'member', 'generation', onRemoved) ).resolves.toEqual({ removed: 1, finished: true }) - expect(onRemoved.mock.calls.map(([ids]) => (ids as string[]).length)).toEqual([500, 1]) - expect(dbChainMockFns.limit).toHaveBeenCalledWith(500) + expect(onRemoved.mock.calls.map(([ids]) => (ids as string[]).length)).toEqual([25, 1]) + /** A page rematerialises in the caller's lease transaction, so it holds one ACL batch. */ + expect(dbChainMockFns.limit).toHaveBeenCalledWith(25) expect(onRemoved.mock.invocationCallOrder[0]).toBeLessThan( dbChainMockFns.delete.mock.invocationCallOrder[1] ) }) + + /** One document above the row cap is a page alone; the rest of the candidates wait. */ + it('removes only the leading page of projection rows and reports more to come', async () => { + queueTableRows(schemaMock.knowledgeDocumentObservation, [ + { documentId: 'huge' }, + { documentId: 'small' }, + ]) + queueTableRows(schemaMock.document, [ + { id: 'huge', chunkCount: 1_000 }, + { id: 'small', chunkCount: 1 }, + ]) + dbChainMockFns.returning.mockResolvedValueOnce([{ documentId: 'huge' }]) + const onRemoved = vi.fn(async (_ids: string[]) => undefined) + + await expect( + removeUnseenMemberObservations(db, 'member', 'generation', onRemoved) + ).resolves.toEqual({ removed: 1, finished: false }) + expect(onRemoved).toHaveBeenCalledWith(['huge']) + }) }) describe('renewMemberObservationsInScopes', () => { @@ -159,6 +186,10 @@ describe('sweepStaleMemberObservations', () => { queueTableRows(schemaMock.knowledgeConnectorMember, [STALE_MEMBER]) queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) queueTableRows(schemaMock.knowledgeConnectorMember, [{ id: 'm-1' }]) + queueTableRows(schemaMock.knowledgeDocumentObservation, [ + { documentId: 'd-1' }, + { documentId: 'd-2' }, + ]) dbChainMockFns.returning .mockResolvedValueOnce([{ documentId: 'd-1' }, { documentId: 'd-2' }]) .mockResolvedValueOnce([{ id: 'd-1' }, { id: 'd-2' }]) @@ -172,12 +203,157 @@ describe('sweepStaleMemberObservations', () => { }) expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() - expect(dbChainMockFns.for).toHaveBeenNthCalledWith(1, 'share') - expect(dbChainMockFns.for).toHaveBeenNthCalledWith(2, 'update') + /** The member row first, then the page's documents, the connector row last. */ + expect(dbChainMockFns.for.mock.calls.map(([mode]) => mode)).toEqual([ + 'update', + 'update', + 'share', + ]) + expect(dbChainMockFns.for.mock.invocationCallOrder.at(-1)).toBeGreaterThan( + dbChainMockFns.set.mock.invocationCallOrder.at(-1)! + ) expect(dbChainMockFns.delete).toHaveBeenCalledWith(schemaMock.knowledgeDocumentObservation) expect(dbChainMockFns.set).toHaveBeenLastCalledWith({ deletedAt: NOW }) }) + /** Each page rematerialises at most one ACL batch under the shared connector row. */ + it('sweeps a large member in pages of 25, one bounded transaction each', async () => { + const ids = (count: number, prefix: string) => + Array.from({ length: count }, (_unused, index) => `${prefix}-${index}`) + queueTableRows(schemaMock.knowledgeConnectorMember, [STALE_MEMBER]) + for (const page of [ids(25, 'a'), ids(3, 'b')]) { + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) + queueTableRows(schemaMock.knowledgeConnectorMember, [{ id: 'm-1' }]) + queueTableRows( + schemaMock.knowledgeDocumentObservation, + page.map((documentId) => ({ documentId })) + ) + dbChainMockFns.returning + .mockResolvedValueOnce(page.map((documentId) => ({ documentId }))) + .mockResolvedValueOnce(page.map((id) => ({ id }))) + .mockResolvedValueOnce([]) + } + + await expect(sweepStaleMemberObservations(NOW)).resolves.toEqual({ + members: 1, + observationsRemoved: 28, + documentsRematerialized: 28, + docsTombstoned: 0, + }) + + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(2) + expect(dbChainMockFns.limit).toHaveBeenCalledWith(25) + const bounds = dbChainMockFns.execute.mock.calls.filter((call: unknown[]) => + JSON.stringify(call).includes('lock_timeout') + ) + expect(bounds).toHaveLength(2) + }) + + /** A document above the row cap is swept alone; the rest of the member waits for the next page. */ + it('sweeps a document larger than one page of projection rows in a page alone', async () => { + queueTableRows(schemaMock.knowledgeConnectorMember, [STALE_MEMBER]) + for (const page of [['huge', 'small'], ['small']]) { + queueTableRows(schemaMock.knowledgeConnectorMember, [{ id: 'm-1' }]) + queueTableRows( + schemaMock.knowledgeDocumentObservation, + page.map((documentId) => ({ documentId })) + ) + queueTableRows( + schemaMock.document, + page.map((id) => ({ id, chunkCount: id === 'huge' ? 1_000 : 1 })) + ) + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) + dbChainMockFns.returning + .mockResolvedValueOnce([{ documentId: page[0] }]) + .mockResolvedValueOnce([{ id: page[0] }]) + .mockResolvedValueOnce([]) + } + + await expect(sweepStaleMemberObservations(NOW)).resolves.toMatchObject({ + members: 1, + observationsRemoved: 2, + }) + const deletedPages = dbChainMockFns.where.mock.calls + .map(([condition]) => + flattenMockConditions(condition).find( + (node) => + node.type === 'inArray' && + node.column === schemaMock.knowledgeDocumentObservation.documentId + ) + ) + .filter((node) => node !== undefined) + .map((node) => node!.values) + expect(deletedPages).toEqual([['huge'], ['small']]) + }) + + /** The budget is checked before every page, so one large member cannot hold a tick past it. */ + it('leaves the rest of a large member for the next tick once the budget passes mid-member', async () => { + let clock = Date.now() + const now = vi.spyOn(Date, 'now').mockImplementation(() => clock) + try { + const page = Array.from({ length: 25 }, (_unused, index) => `a-${index}`) + queueTableRows(schemaMock.knowledgeConnectorMember, [STALE_MEMBER]) + queueTableRows(schemaMock.knowledgeConnectorMember, [{ id: 'm-1' }]) + queueTableRows( + schemaMock.knowledgeDocumentObservation, + page.map((documentId) => ({ documentId })) + ) + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) + dbChainMockFns.returning + .mockResolvedValueOnce(page.map((documentId) => ({ documentId }))) + .mockResolvedValueOnce(page.map((id) => ({ id }))) + .mockImplementationOnce(async () => { + clock += 2_000 + return [] + }) + + await expect(sweepStaleMemberObservations(NOW, clock + 1_000)).resolves.toEqual({ + members: 1, + observationsRemoved: 25, + documentsRematerialized: 25, + docsTombstoned: 0, + }) + expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() + } finally { + now.mockRestore() + } + }) + + /** A connector whose running member page holds its row must not stall the other connectors. */ + it('defers a member whose page hits a lock timeout and still sweeps the next one', async () => { + queueTableRows(schemaMock.knowledgeConnectorMember, [ + STALE_MEMBER, + /** Another member of the busy connector waits a tick rather than wait out the same lock. */ + { ...STALE_MEMBER, id: 'm-3' }, + { ...STALE_MEMBER, id: 'm-2', connectorId: 'c-2' }, + ]) + dbChainMockFns.transaction.mockRejectedValueOnce( + Object.assign(new Error('canceling statement due to lock timeout'), { code: '55P03' }) + ) + queueTableRows(schemaMock.knowledgeConnectorMember, [{ id: 'm-2' }]) + queueTableRows(schemaMock.knowledgeDocumentObservation, [{ documentId: 'd-1' }]) + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-2' }]) + dbChainMockFns.returning + .mockResolvedValueOnce([{ documentId: 'd-1' }]) + .mockResolvedValueOnce([{ id: 'd-1' }]) + .mockResolvedValueOnce([]) + + await expect(sweepStaleMemberObservations(NOW)).resolves.toEqual({ + members: 1, + observationsRemoved: 1, + documentsRematerialized: 1, + docsTombstoned: 0, + }) + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(2) + }) + + it('fails the sweep on an error that is not a lock or capacity failure', async () => { + queueTableRows(schemaMock.knowledgeConnectorMember, [STALE_MEMBER]) + dbChainMockFns.transaction.mockRejectedValueOnce(new TypeError('broken')) + + await expect(sweepStaleMemberObservations(NOW)).rejects.toThrow('broken') + }) + /** * A run that claimed the member between the selection and the lock moved * `lastStartedAt` forward, so the re-check under `FOR UPDATE` finds nothing @@ -199,16 +375,194 @@ describe('sweepStaleMemberObservations', () => { expect(dbChainMockFns.update).not.toHaveBeenCalled() }) - /** A connector that left members mode after the selection no longer matches the shared lock's re-check. */ + /** A connector that left members mode after the selection fails the last check and rolls the page back. */ it('leaves a connector that left members mode after it was selected', async () => { queueTableRows(schemaMock.knowledgeConnectorMember, [STALE_MEMBER]) + queueTableRows(schemaMock.knowledgeConnectorMember, [{ id: 'm-1' }]) queueTableRows(schemaMock.knowledgeConnector, []) + dbChainMockFns.returning.mockResolvedValueOnce([{ documentId: 'd-1' }]) - await expect(sweepStaleMemberObservations(NOW)).resolves.toMatchObject({ members: 0 }) + await expect(sweepStaleMemberObservations(NOW)).resolves.toMatchObject({ + members: 0, + observationsRemoved: 0, + }) - expect(dbChainMockFns.for).toHaveBeenCalledWith('share') - expect(dbChainMockFns.for).not.toHaveBeenCalledWith('update') - expect(dbChainMockFns.delete).not.toHaveBeenCalled() + expect(dbChainMockFns.for).toHaveBeenLastCalledWith('share') + expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() + }) + + it('stops a tick at its wall-clock budget, leaving the rest for the next', async () => { + queueTableRows(schemaMock.knowledgeConnectorMember, [STALE_MEMBER]) + + await expect(sweepStaleMemberObservations(NOW, Date.now() - 1)).resolves.toMatchObject({ + members: 0, + }) + expect(dbChainMockFns.transaction).not.toHaveBeenCalled() + }) +}) + +describe('materializeDocumentAcls', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + }) + + /** Each rematerialised document rewrites every projection row of its chunks. */ + it('rematerialises 25 documents per statement', async () => { + const ids = Array.from({ length: 60 }, (_unused, index) => `d-${index}`) + dbChainMockFns.returning + .mockResolvedValueOnce(ids.slice(0, 25).map((id) => ({ id }))) + .mockResolvedValueOnce(ids.slice(25, 50).map((id) => ({ id }))) + .mockResolvedValueOnce(ids.slice(50).map((id) => ({ id }))) + + await expect(materializeDocumentAcls('c-1', ids, db)).resolves.toBe(60) + + const sizes = dbChainMockFns.where.mock.calls.map( + ([condition]) => + flattenMockConditions(condition).find( + (node) => node.type === 'inArray' && node.column === schemaMock.document.id + )?.values as string[] + ) + expect(sizes.map((values) => values.length)).toEqual([25, 25, 10]) + }) +}) + +describe('pagesByProjectionRows', () => { + const doc = (id: string, chunkCount: number) => ({ id, chunkCount }) + + it('fills a page up to the projection row cap and keeps order', () => { + expect( + pagesByProjectionRows([doc('a', 100), doc('b', 150), doc('c', 1), doc('d', 249)]) + ).toEqual([ + ['a', 'b'], + ['c', 'd'], + ]) + }) + + it('gives a document above the cap a page alone so it still makes progress', () => { + expect(pagesByProjectionRows([doc('a', 1), doc('huge', 5_000), doc('b', 1)])).toEqual([ + ['a'], + ['huge'], + ['b'], + ]) + }) + + it('starts with a page for a leading document above the cap, never an empty one', () => { + expect(pagesByProjectionRows([doc('huge', 5_000), doc('a', 1)])).toEqual([['huge'], ['a']]) + }) + + it('returns no pages for no documents', () => { + expect(pagesByProjectionRows([])).toEqual([]) + }) +}) + +describe('rematerializeDocumentAcls', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + }) + const page = vi.fn() + const transaction: LeaseTransaction = (write) => { + page() + return write(db) + } + + it('opens no transaction for documents whose ACL already matches their observers', async () => { + queueTableRows(schemaMock.document, []) + + await expect(rematerializeDocumentAcls('c-1', ['d-1', 'd-2'], transaction)).resolves.toBe(0) + + expect(page).not.toHaveBeenCalled() + expect(dbChainMockFns.update).not.toHaveBeenCalled() + }) + + it('rematerialises only what differs, in pages bounded by projection rows', async () => { + queueTableRows( + schemaMock.document, + Array.from({ length: 30 }, (_unused, index) => ({ id: `d-${index}`, chunkCount: 10 })) + ) + dbChainMockFns.returning + .mockResolvedValueOnce(Array.from({ length: 25 }, (_unused, index) => ({ id: `d-${index}` }))) + .mockResolvedValueOnce(Array.from({ length: 5 }, (_unused, index) => ({ id: `e-${index}` }))) + + await expect( + rematerializeDocumentAcls( + 'c-1', + Array.from({ length: 40 }, (_unused, index) => `d-${index}`), + transaction + ) + ).resolves.toBe(30) + + expect(page).toHaveBeenCalledTimes(2) + }) +}) + +describe('writeProjectionPages', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + }) + + /** A planned page that splits under the locked reread stops at the budget, not after it. */ + it('stops before the next page once the deadline passes and reports it unfinished', async () => { + let clock = Date.now() + const now = vi.spyOn(Date, 'now').mockImplementation(() => clock) + try { + queueTableRows(schemaMock.document, [ + { id: 'a', chunkCount: 200 }, + { id: 'b', chunkCount: 200 }, + { id: 'c', chunkCount: 200 }, + ]) + const write = vi.fn(async (_tx: unknown, page: string[]) => { + clock += 2_000 + return page.length + }) + + await expect( + writeProjectionPages(['a', 'b', 'c'], (fn) => fn(db), write, { + deadlineAt: clock + 1_000, + }) + ).resolves.toEqual({ written: 1, finished: false }) + expect(write).toHaveBeenCalledOnce() + } finally { + now.mockRestore() + } + }) + + /** The lease heartbeat runs before each page and can itself use up what is left of the budget. */ + it('opens no transaction when the heartbeat before a page outlasts the deadline', async () => { + let clock = Date.now() + const now = vi.spyOn(Date, 'now').mockImplementation(() => clock) + try { + const transaction = vi.fn() + const write = vi.fn() + const beforePage = vi.fn(async () => { + clock += 2_000 + }) + + await expect( + writeProjectionPages(['a'], transaction, write, { beforePage, deadlineAt: clock + 1_000 }) + ).resolves.toEqual({ written: 0, finished: false }) + expect(beforePage).toHaveBeenCalledOnce() + expect(transaction).not.toHaveBeenCalled() + } finally { + now.mockRestore() + } + }) + + it('writes every page when no deadline is given', async () => { + queueTableRows(schemaMock.document, [ + { id: 'a', chunkCount: 200 }, + { id: 'b', chunkCount: 200 }, + ]) + queueTableRows(schemaMock.document, [{ id: 'b', chunkCount: 200 }]) + const write = vi.fn(async (_tx: unknown, page: string[]) => page.length) + + await expect(writeProjectionPages(['a', 'b'], (fn) => fn(db), write)).resolves.toEqual({ + written: 2, + finished: true, + }) + expect(write.mock.calls.map(([, page]) => page)).toEqual([['a'], ['b']]) }) }) @@ -218,25 +572,167 @@ describe('rewriteConnectorAcls', () => { resetDbChainMock() }) - it('proves the lease inside each batch transaction before rewriting', async () => { + const stale = (id: string, aclDiffers = true, evidencePresent = false, chunkCount = 10) => ({ + id, + externalId: `ext-${id}`, + chunkCount, + aclDiffers, + evidencePresent, + }) + const held = { stillHeld: () => 'held' as never } + /** The ids each `acl` assignment targets, in call order. */ + const assignedPages = () => + dbChainMockFns.set.mock.calls + .map(([values], index) => ({ values, where: dbChainMockFns.where.mock.calls[index] })) + .filter(({ values }) => 'acl' in values) + const pageSizes = () => + dbChainMockFns.set.mock.calls + .map(([values], index) => ({ + values, + order: dbChainMockFns.set.mock.invocationCallOrder[index], + })) + .filter(({ values }) => 'acl' in values || 'aclRequirements' in values) + .map(({ order }) => { + const whereIndex = dbChainMockFns.where.mock.invocationCallOrder.findIndex( + (whereOrder) => whereOrder > order + ) + return ( + flattenMockConditions(dbChainMockFns.where.mock.calls[whereIndex]?.[0]).find( + (node) => node.type === 'inArray' && node.column === schemaMock.document.id + )?.values as string[] + ).length + }) + + it('proves the lease inside each page transaction before rewriting', async () => { + queueTableRows(schemaMock.document, [stale('d-1')]) queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'd-1' }]) + queueTableRows(schemaMock.document, []) - await expect( - rewriteConnectorAcls('c-1', [], { lease: { stillHeld: () => 'held' as never } }) - ).resolves.toBe(true) + await expect(rewriteConnectorAcls('c-1', [], { lease: held })).resolves.toBe(true) expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() expect(dbChainMockFns.for).toHaveBeenCalledWith('share') expect(dbChainMockFns.update).toHaveBeenCalledWith(schemaMock.document) }) - it('stops without writing once the lease is gone', async () => { + /** The lease is the page's last statement: a lost lease rolls the page back and ends the rewrite. */ + it('proves the lease after the page writes and stops once it is gone', async () => { + queueTableRows( + schemaMock.document, + Array.from({ length: 60 }, (_unused, index) => stale(`d-${index}`)) + ) queueTableRows(schemaMock.knowledgeConnector, []) + await expect(rewriteConnectorAcls('c-1', [], { lease: held })).rejects.toBeInstanceOf( + SyncLockLostException + ) + + expect(dbChainMockFns.update).toHaveBeenCalledOnce() + const leaseCheck = dbChainMockFns.for.mock.calls.findIndex(([mode]) => mode === 'share') + expect(dbChainMockFns.for.mock.invocationCallOrder[leaseCheck]).toBeGreaterThan( + dbChainMockFns.update.mock.invocationCallOrder[0] + ) + }) + + it('gives a document larger than one page of projection rows a page alone', async () => { + queueTableRows(schemaMock.document, [ + stale('d-small'), + stale('d-huge', true, false, 1_000), + stale('d-next'), + ]) + for (let page = 0; page < 3; page++) { + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) + dbChainMockFns.returning.mockResolvedValueOnce([{ id: `written-${page}` }]) + } + + await expect(rewriteConnectorAcls('c-1', [], { lease: held })).resolves.toBe(true) + + expect(pageSizes()).toEqual([1, 1, 1]) + }) + + /** Each assignment rewrites every projection row of its documents; one page per transaction. */ + it('assigns acl in pages of 25, each in a bounded transaction of its own', async () => { + const window = Array.from({ length: 60 }, (_unused, index) => stale(`d-${index}`)) + queueTableRows(schemaMock.document, window) + for (let page = 0; page < 3; page++) { + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) + dbChainMockFns.returning.mockResolvedValueOnce([{ id: `written-${page}` }]) + } + queueTableRows(schemaMock.document, []) + + await expect(rewriteConnectorAcls('c-1', [], { lease: held })).resolves.toBe(true) + + expect(assignedPages()).toHaveLength(3) + expect(pageSizes()).toEqual([25, 25, 10]) + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(3) + const bounds = dbChainMockFns.execute.mock.calls.filter((call: unknown[]) => + JSON.stringify(call).includes('lock_timeout') + ) + expect(bounds).toHaveLength(3) + }) + + it('clears evidence without assigning acl where only the evidence is stale', async () => { + queueTableRows(schemaMock.document, [stale('d-1', false, true), stale('d-2', false, false)]) + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) + dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'd-1' }]) + queueTableRows(schemaMock.document, []) + + await expect(rewriteConnectorAcls('c-1', [], { lease: held })).resolves.toBe(true) + + expect(assignedPages()).toHaveLength(0) + expect(dbChainMockFns.set).toHaveBeenCalledWith({ aclRequirements: [], aclVerifiedAt: null }) + expect(pageSizes()).toEqual([1]) + }) + + it('writes nothing and opens no transaction on a connector with nothing stale', async () => { + queueTableRows( + schemaMock.document, + Array.from({ length: 500 }, (_unused, index) => stale(`d-${index}`, false)) + ) + queueTableRows(schemaMock.document, [stale('d-last', false)]) + + await expect(rewriteConnectorAcls('c-1', [], { lease: held })).resolves.toBe(true) + + expect(dbChainMockFns.transaction).not.toHaveBeenCalled() + /** A full window is followed by the next one, from the last key read. */ + expect(dbChainMockFns.limit).toHaveBeenCalledTimes(2) + }) + + /** A planned page that splits under the locked reread still stops at the deadline. */ + it('stops unfinished when a split page reaches the deadline', async () => { + let clock = Date.now() + const now = vi.spyOn(Date, 'now').mockImplementation(() => clock) + try { + queueTableRows(schemaMock.document, [ + stale('d-1', true, false, 1), + stale('d-2', true, false, 1), + ]) + queueTableRows(schemaMock.document, [ + { id: 'd-1', chunkCount: 200 }, + { id: 'd-2', chunkCount: 200 }, + ]) + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) + dbChainMockFns.returning.mockImplementationOnce(async () => { + clock += 2_000 + return [{ id: 'd-1' }] + }) + + await expect( + rewriteConnectorAcls('c-1', [], { lease: held, deadlineAt: clock + 1_000 }) + ).resolves.toBe(false) + expect(dbChainMockFns.update).toHaveBeenCalledOnce() + } finally { + now.mockRestore() + } + }) + + it('stops unfinished at the deadline before the next page', async () => { + queueTableRows(schemaMock.document, [stale('d-1')]) + await expect( - rewriteConnectorAcls('c-1', [], { lease: { stillHeld: () => 'lost' as never } }) - ).rejects.toBeInstanceOf(SyncLockLostException) + rewriteConnectorAcls('c-1', [], { lease: held, deadlineAt: Date.now() - 1 }) + ).resolves.toBe(false) expect(dbChainMockFns.update).not.toHaveBeenCalled() }) diff --git a/apps/sim/lib/knowledge/connectors/member-observations.ts b/apps/sim/lib/knowledge/connectors/member-observations.ts index 9f77dc7d4e6..eac00e3567f 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.ts @@ -5,6 +5,9 @@ import { knowledgeConnectorMember, knowledgeDocumentObservation, } from '@sim/db/schema' +import { createLogger } from '@sim/logger' +import { getPostgresErrorCode, getTransientDatabaseFailure } from '@sim/utils/errors' +import { chunkArray } from '@sim/utils/helpers' import { and, asc, @@ -16,21 +19,28 @@ import { isNull, lt, ne, + not, notExists, or, + type SQL, sql, } from 'drizzle-orm' import type { DbOrTx } from '@/lib/db/types' import { textArrayLiteral } from '@/lib/knowledge/access/predicate' import { + ACL_CHANGE_BATCH_SIZE, + ACL_WRITE_BATCH_SIZE, MEMBER_OBSERVATION_STALE_AFTER_HOURS, MEMBER_PURGE_MAX_PER_RUN, MEMBER_TOMBSTONE_PURGE_DAYS, MEMBER_TOMBSTONE_RECONCILE_PAGES_PER_RUN, + PROJECTION_ROW_BATCH_SIZE, } from '@/lib/knowledge/connectors/sync-limits' import { - assertSyncLeaseHeldInTx, + boundLeaseTransaction, connectorIsLive, + type LeaseTransaction, + leaseTransaction, MEMBER_LOCKABLE_CONNECTOR_STATUSES, SyncLockLostException, type SyncRunLease, @@ -42,14 +52,21 @@ import { hardDeleteDocuments, } from '@/lib/knowledge/documents/service' -/** Documents rematerialised per `UPDATE`; keeps each statement's bind list and lock footprint small. */ -const MATERIALIZE_BATCH_SIZE = 500 +const logger = createLogger('MemberObservations') + +/** + * Documents per tombstone or resurrection page. Those writes set `deleted_at` alone and fire no + * projection fan-out, so they page wider than an ACL write. + */ +const LIFECYCLE_PAGE_SIZE = 500 /** Observation rows written per `INSERT`. */ const OBSERVATION_BATCH_SIZE = 500 /** Documents hard-deleted per call, so the lease heartbeat runs between chunks. */ const PURGE_CHUNK_SIZE = 25 /** Members one scheduler tick will sweep; the rest wait for the next tick. */ const STALE_MEMBER_SWEEP_LIMIT = 200 +/** Wall clock one sweep tick spends; the members it did not reach are still stale next tick. */ +const STALE_MEMBER_SWEEP_BUDGET_MS = 60_000 /** * The subject-token aggregate that is a members-mode document's ACL. Ordered @@ -219,7 +236,10 @@ export async function renewMemberObservationsInScopes(input: { /** * Removes every observation of one member that this run did not re-assert. * Only called after a full, complete, non-suspect listing: absence from any - * other kind of listing says nothing about access. + * other kind of listing says nothing about access. One page per call, of at + * most {@link ACL_CHANGE_BATCH_SIZE} documents and one page of projection rows + * ({@link lockProjectionPage}), because `onRemoved` rematerialises the page's + * ACLs in the caller's lease transaction. */ export async function removeUnseenMemberObservations( executor: DbOrTx, @@ -231,17 +251,27 @@ export async function removeUnseenMemberObservations( eq(knowledgeDocumentObservation.memberId, memberId), ne(knowledgeDocumentObservation.runId, runId) ) - const candidates = executor + const candidates = await executor .select({ documentId: knowledgeDocumentObservation.documentId }) .from(knowledgeDocumentObservation) .where(unseen) - .limit(OBSERVATION_BATCH_SIZE) - const removed = await executor - .delete(knowledgeDocumentObservation) - .where(and(unseen, inArray(knowledgeDocumentObservation.documentId, candidates))) - .returning({ documentId: knowledgeDocumentObservation.documentId }) + .limit(ACL_CHANGE_BATCH_SIZE) + const { page, rest } = await lockProjectionPage( + executor, + candidates.map((row) => row.documentId) + ) + const removed = + page.length === 0 + ? [] + : await executor + .delete(knowledgeDocumentObservation) + .where(and(unseen, inArray(knowledgeDocumentObservation.documentId, page))) + .returning({ documentId: knowledgeDocumentObservation.documentId }) if (removed.length > 0) await onRemoved(removed.map((row) => row.documentId)) - return { removed: removed.length, finished: removed.length < OBSERVATION_BATCH_SIZE } + return { + removed: removed.length, + finished: candidates.length < ACL_CHANGE_BATCH_SIZE && rest.length === 0, + } } /** @@ -319,19 +349,210 @@ export async function tombstoneDocumentsObservedOnlyBy( return tombstoned } +/** The connector's documents one keyset window reads while looking for ACLs to rewrite. */ +interface ConnectorDocumentCursor { + externalId: string + id: string +} + /** - * Rewrites `document.acl` from the observation graph: the sorted subject - * tokens of every active observer, or nobody. Scoped to the connector so a - * document id that was detached or re-owned since it was collected is left - * alone. + * Splits documents into ACL pages whose chunks stay within {@link PROJECTION_ROW_BATCH_SIZE} + * search projection rows per table, the rows an ACL assignment sends through the projection + * trigger. Every page holds at least one document, so one larger than the cap still makes progress + * alone. Documents keep their order. */ -/** Documents rewritten per statement while a mode switch rewrites a connector's ACLs. */ -const ACCESS_REWRITE_BATCH_SIZE = 1000 +export function pagesByProjectionRows( + documents: readonly { id: string; chunkCount: number }[] +): string[][] { + const pages: string[][] = [] + let page: string[] = [] + let rows = 0 + for (const entry of documents) { + const cost = Math.max(0, entry.chunkCount) + if (page.length > 0 && rows + cost > PROJECTION_ROW_BATCH_SIZE) { + pages.push(page) + page = [] + rows = 0 + } + page.push(entry.id) + rows += cost + } + if (page.length > 0) pages.push(page) + return pages +} /** - * Rewrites every document ACL of the connector to `target`, in bounded - * batches, until done or `deadlineAt` passes. Returns whether every row was - * rewritten. `beforeBatch` runs ahead of each statement, for a lease heartbeat. + * Locks `documentIds` (in id order, so pages never deadlock one another) and splits off the leading + * page, in the order given, whose chunks as committed now fit {@link PROJECTION_ROW_BATCH_SIZE}; + * the rest wait for a later page. A chunk count read + * earlier without a lock can be stale: a processing commit holds its document's row while it + * replaces the chunks and sets `chunk_count`, so locking first either waits for that commit and + * reads its count, or makes it wait until this page commits, when the rows it inserts copy the new + * ACL. The rows are the ones the page's write locks anyway. An id with no document row costs + * nothing. + */ +export async function lockProjectionPage( + tx: DbOrTx, + documentIds: readonly string[] +): Promise<{ page: string[]; rest: string[] }> { + if (documentIds.length === 0) return { page: [], rest: [] } + const locked = await tx + .select({ id: document.id, chunkCount: document.chunkCount }) + .from(document) + .where(inArray(document.id, [...documentIds])) + .orderBy(asc(document.id)) + .for('update') + const chunks = new Map(locked.map((row) => [row.id, row.chunkCount])) + const ordered = [...new Set(documentIds)].map((id) => ({ id, chunkCount: chunks.get(id) ?? 0 })) + const [page = []] = pagesByProjectionRows(ordered) + const taken = new Set(page) + return { page, rest: ordered.map(({ id }) => id).filter((id) => !taken.has(id)) } +} + +/** + * Writes `documentIds` one page per `transaction`, each page sized by the chunk counts it reads + * under {@link lockProjectionPage}; whatever no longer fits waits for the next page. Stops before + * the next transaction once `deadlineAt` passes, so one planned page that splits into many cannot + * run past a caller's budget; `finished` is then false and the caller must not record the + * documents as written. Returns the rows `write` reports. + */ +export async function writeProjectionPages( + documentIds: readonly string[], + transaction: LeaseTransaction, + write: (tx: DbOrTx, page: string[]) => Promise, + options: { beforePage?: () => Promise; deadlineAt?: number } = {} +): Promise<{ written: number; finished: boolean }> { + let pending = [...documentIds] + let written = 0 + while (pending.length > 0) { + await options.beforePage?.() + if (options.deadlineAt !== undefined && Date.now() >= options.deadlineAt) + return { written, finished: false } + const { rows, rest } = await transaction(async (tx) => { + const { page, rest } = await lockProjectionPage(tx, pending) + return { rows: page.length > 0 ? await write(tx, page) : 0, rest } + }) + written += rows + pending = rest + } + return { written, finished: true } +} + +/** + * Rewrites the ACL of every document of the connector to `target`, clearing its permission + * evidence, one short transaction per page. The documents are walked once in keyset windows of + * {@link ACL_WRITE_BATCH_SIZE} through `doc_connector_source_lookup_idx`, so no read revisits the + * rows earlier pages fixed; every connector document carries an external id, since the sync's + * inserts are the only writers of `connector_id`. Within a window, the documents whose `acl` + * differs are assigned it in pages bounded by {@link pagesByProjectionRows}, each its own + * transaction, because each assignment rewrites the document's search projection rows; those whose + * `acl` already matches only have their evidence cleared, which fires no fan-out. Every write + * re-checks its row and `guard`, so a page is idempotent and a crash resumes by rewriting what is + * still stale. Callers hold a lease that keeps every other ACL writer of the connector off it, or + * rewrite only toward what such a writer would also write. Returns the documents written and + * whether it finished before `deadlineAt`. + */ +export async function rewriteConnectorDocumentAcls(input: { + connectorId: string + target: readonly string[] + transaction: LeaseTransaction + /** Must hold for the connector, in every read and write, or the rewrite stops. */ + guard?: SQL + beforePage?: () => Promise + deadlineAt?: number +}): Promise<{ rewritten: number; finished: boolean }> { + const { connectorId, target, transaction, guard } = input + const aclDiffers = + target.length === 0 + ? sql`cardinality(${document.acl}) > 0` + : sql`${document.acl} <> ${textArrayLiteral(target)}` + const evidencePresent = sql`(${document.aclRequirements} <> '[]'::jsonb OR ${document.aclVerifiedAt} IS NOT NULL)` + const expired = () => input.deadlineAt !== undefined && Date.now() >= input.deadlineAt + let rewritten = 0 + let after: ConnectorDocumentCursor | undefined + for (;;) { + await input.beforePage?.() + if (expired()) return { rewritten, finished: false } + const window = await db + .select({ + id: document.id, + externalId: document.externalId, + chunkCount: document.chunkCount, + aclDiffers: sql`${aclDiffers}`, + evidencePresent: sql`${evidencePresent}`, + }) + .from(document) + .where( + and( + eq(document.connectorId, connectorId), + isNotNull(document.externalId), + after + ? sql`${document.externalId} >= ${after.externalId} AND (${document.externalId} > ${after.externalId} OR ${document.id} > ${after.id})` + : undefined, + guard + ) + ) + .orderBy(asc(document.externalId), asc(document.id)) + .limit(ACL_WRITE_BATCH_SIZE) + const evidenceOnly = window + .filter((row) => !row.aclDiffers && row.evidencePresent) + .map((row) => row.id) + if (evidenceOnly.length > 0) { + if (expired()) return { rewritten, finished: false } + const rows = await transaction((tx) => + tx + .update(document) + .set({ aclRequirements: [], aclVerifiedAt: null }) + .where( + and( + eq(document.connectorId, connectorId), + inArray(document.id, evidenceOnly), + not(aclDiffers), + evidencePresent, + guard + ) + ) + .returning({ id: document.id }) + ) + rewritten += rows.length + } + for (const page of pagesByProjectionRows(window.filter((row) => row.aclDiffers))) { + if (expired()) return { rewritten, finished: false } + const written = await writeProjectionPages( + page, + transaction, + async (tx, locked) => + ( + await tx + .update(document) + .set({ acl: [...target], aclRequirements: [], aclVerifiedAt: null }) + .where( + and( + eq(document.connectorId, connectorId), + inArray(document.id, locked), + aclDiffers, + guard + ) + ) + .returning({ id: document.id }) + ).length, + { beforePage: input.beforePage, deadlineAt: input.deadlineAt } + ) + rewritten += written.written + if (!written.finished) return { rewritten, finished: false } + } + const last = window.at(-1) + if (window.length < ACL_WRITE_BATCH_SIZE || !last?.externalId) + return { rewritten, finished: true } + after = { externalId: last.externalId, id: last.id } + } +} + +/** + * Rewrites every document ACL of the connector to `target`, one short + * transaction per page, until done or `deadlineAt` passes. Returns whether + * every row was rewritten. `beforeBatch` runs ahead of each page, for a lease + * heartbeat. */ export async function rewriteConnectorAcls( connectorId: string, @@ -340,43 +561,31 @@ export async function rewriteConnectorAcls( deadlineAt?: number beforeBatch?: () => Promise /** - * The lease the caller holds on the connector, proved inside each batch's - * transaction: a heartbeat before the batch only says the lease was held + * The lease the caller holds on the connector, proved inside each page's + * transaction: a heartbeat before the page only says the lease was held * then, and a run reclaimed mid-rewrite must not land an empty ACL over * what its replacement has since materialised. */ lease?: SyncWriteLease } = {} ): Promise { - const aclMismatch = - target.length === 0 - ? sql`cardinality(${document.acl}) > 0` - : sql`${document.acl} <> ${textArrayLiteral(target)}` - const mismatch = sql`(${aclMismatch} OR ${document.aclRequirements} <> '[]'::jsonb OR ${document.aclVerifiedAt} IS NOT NULL)` - for (;;) { - await options.beforeBatch?.() - const rewritten = await db.transaction(async (tx) => { - if (options.lease) await assertSyncLeaseHeldInTx(tx, connectorId, options.lease) - return tx - .update(document) - .set({ acl: [...target], aclRequirements: [], aclVerifiedAt: null }) - .where( - eq( - document.id, - sql`ANY(ARRAY( - SELECT ${document.id} FROM ${document} - WHERE ${document.connectorId} = ${connectorId} AND ${mismatch} - LIMIT ${ACCESS_REWRITE_BATCH_SIZE} - ))` - ) - ) - .returning({ id: document.id }) - }) - if (rewritten.length < ACCESS_REWRITE_BATCH_SIZE) return true - if (options.deadlineAt !== undefined && Date.now() >= options.deadlineAt) return false - } + const { finished } = await rewriteConnectorDocumentAcls({ + connectorId, + target, + transaction: leaseTransaction(connectorId, options.lease), + beforePage: options.beforeBatch, + deadlineAt: options.deadlineAt, + }) + return finished } +/** + * Rewrites `document.acl` from the observation graph: the sorted subject + * tokens of every active observer, or nobody. Scoped to the connector so a + * document id that was detached or re-owned since it was collected is left + * alone. {@link ACL_CHANGE_BATCH_SIZE} documents per statement; callers keep + * the documents one transaction materialises to the same bound. + */ export async function materializeDocumentAcls( connectorId: string, documentIds: Iterable, @@ -384,8 +593,7 @@ export async function materializeDocumentAcls( ): Promise { const ids = [...new Set(documentIds)] let updated = 0 - for (let offset = 0; offset < ids.length; offset += MATERIALIZE_BATCH_SIZE) { - const batch = ids.slice(offset, offset + MATERIALIZE_BATCH_SIZE) + for (const batch of chunkArray(ids, ACL_CHANGE_BATCH_SIZE)) { const rows = await executor .update(document) .set({ acl: observedAcl(), aclRequirements: [], aclVerifiedAt: null }) @@ -402,6 +610,44 @@ export async function materializeDocumentAcls( return updated } +/** + * Rematerialises the ACLs of `documentIds` that differ from the observation graph, in pages bounded + * by {@link pagesByProjectionRows}, each its own `transaction`. The documents that differ are read + * first without a lock, {@link ACL_WRITE_BATCH_SIZE} at a time, so documents whose ACL already + * matches cost one read and no transaction; each page's write re-checks the difference. For pages + * whose observation writes have already committed. + */ +export async function rematerializeDocumentAcls( + connectorId: string, + documentIds: Iterable, + transaction: LeaseTransaction, + beforePage?: () => Promise +): Promise { + let updated = 0 + for (const window of chunkArray([...new Set(documentIds)], ACL_WRITE_BATCH_SIZE)) { + const stale = await db + .select({ id: document.id, chunkCount: document.chunkCount }) + .from(document) + .where( + and( + inArray(document.id, window), + eq(document.connectorId, connectorId), + sql`(${document.acl} IS DISTINCT FROM ${observedAcl()} OR ${document.aclRequirements} <> '[]'::jsonb OR ${document.aclVerifiedAt} IS NOT NULL)` + ) + ) + for (const page of pagesByProjectionRows(stale)) { + const { written } = await writeProjectionPages( + page, + transaction, + (tx, locked) => materializeDocumentAcls(connectorId, locked, tx), + { beforePage } + ) + updated += written + } + } + return updated +} + export interface MemberDocumentLifecycleResult { tombstoned: number resurrected: number @@ -501,10 +747,10 @@ async function tombstoneUnobserved( now: Date, result: MemberDocumentLifecycleResult ): Promise { - for (let offset = 0; offset < documentIds.length; offset += MATERIALIZE_BATCH_SIZE) { + for (let offset = 0; offset < documentIds.length; offset += LIFECYCLE_PAGE_SIZE) { if (Date.now() >= input.deadlineAt) return false await input.lease.beatIfDue() - const batch = documentIds.slice(offset, offset + MATERIALIZE_BATCH_SIZE) + const batch = documentIds.slice(offset, offset + LIFECYCLE_PAGE_SIZE) const changed = await input.withLease((tx) => tx .update(document) @@ -576,7 +822,7 @@ async function reconcileUnobservedPages( ) ) .orderBy(asc(document.externalId)) - .limit(MATERIALIZE_BATCH_SIZE) + .limit(LIFECYCLE_PAGE_SIZE) if (Date.now() >= input.deadlineAt) { finished = false break @@ -595,9 +841,7 @@ async function reconcileUnobservedPages( advanced = true const lastExternalId = rows.at(-1)?.externalId after = - rows.length < MATERIALIZE_BATCH_SIZE || !lastExternalId - ? null - : { externalId: lastExternalId } + rows.length < LIFECYCLE_PAGE_SIZE || !lastExternalId ? null : { externalId: lastExternalId } if (!after) break } /** One write per run, not per page: the connector row is hot, and a lost write only repeats pages. */ @@ -664,7 +908,7 @@ export async function applyMemberDocumentLifecycle( ) ) .orderBy(seenOrder, asc(document.id)) - .limit(MATERIALIZE_BATCH_SIZE) + .limit(LIFECYCLE_PAGE_SIZE) if (candidates.length === 0) break if (Date.now() >= input.deadlineAt) return result const changed = await input.withLease(async (tx) => { @@ -684,7 +928,7 @@ export async function applyMemberDocumentLifecycle( }) result.resurrected += changed.length after = candidates.at(-1) - if (candidates.length < MATERIALIZE_BATCH_SIZE) break + if (candidates.length < LIFECYCLE_PAGE_SIZE) break } const purgeCutoff = new Date(now.getTime() - MEMBER_TOMBSTONE_PURGE_DAYS * 24 * 60 * 60 * 1000) @@ -767,6 +1011,112 @@ function memberStillStale(memberId: string, cutoff: Date) { ) } +/** + * Stale-member observations one sweep tick removes per member: the same + * {@link OBSERVATION_BATCH_SIZE} as before, now in pages of + * {@link ACL_CHANGE_BATCH_SIZE}, one transaction each, so no page holds the + * connector row across more ACL fan-out than one statement's. + */ +const STALE_MEMBER_PAGES_PER_TICK = OBSERVATION_BATCH_SIZE / ACL_CHANGE_BATCH_SIZE + +/** + * One page of a stale member's sweep in one short, bounded transaction: the + * connector row shared and the member row locked, both re-checked, then up to + * {@link ACL_CHANGE_BATCH_SIZE} observations removed with the ACLs and + * tombstones they decide. Null when the member or connector no longer qualifies. + */ +async function sweepStaleMemberPage( + member: { id: string; connectorId: string }, + memberCutoff: Date, + now: Date +): Promise<{ + observationsRemoved: number + documentsRematerialized: number + docsTombstoned: number + /** Whether the member still had observations past this page. */ + more: boolean +} | null> { + try { + return await db.transaction(async (tx) => { + await boundLeaseTransaction(tx) + const [stale] = await tx + .select({ id: knowledgeConnectorMember.id }) + .from(knowledgeConnectorMember) + .where(memberStillStale(member.id, memberCutoff)) + .for('update') + if (!stale) return null + + const candidates = await tx + .select({ documentId: knowledgeDocumentObservation.documentId }) + .from(knowledgeDocumentObservation) + .where(eq(knowledgeDocumentObservation.memberId, member.id)) + .limit(ACL_CHANGE_BATCH_SIZE) + const { page, rest } = await lockProjectionPage( + tx, + candidates.map((row) => row.documentId) + ) + const removed = + page.length === 0 + ? [] + : await tx + .delete(knowledgeDocumentObservation) + .where( + and( + eq(knowledgeDocumentObservation.memberId, member.id), + inArray(knowledgeDocumentObservation.documentId, page) + ) + ) + .returning({ documentId: knowledgeDocumentObservation.documentId }) + const documentIds = removed.map((row) => row.documentId) + const rematerialized = await materializeDocumentAcls(member.connectorId, documentIds, tx) + const tombstoned = + documentIds.length === 0 + ? [] + : await tx + .update(document) + .set({ deletedAt: now }) + .where( + and( + inArray(document.id, documentIds), + eq(document.connectorId, member.connectorId), + eq(document.userExcluded, false), + isNull(document.archivedAt), + isNull(document.deletedAt), + hasNoObservation() + ) + ) + .returning({ id: document.id }) + /** Last, so the connector row is never held while the page waits on document rows. */ + const [connector] = await tx + .select({ id: knowledgeConnector.id }) + .from(knowledgeConnector) + .where( + and( + eq(knowledgeConnector.id, member.connectorId), + eq(knowledgeConnector.accessMode, 'members'), + inArray(knowledgeConnector.status, MEMBER_LOCKABLE_CONNECTOR_STATUSES), + ne(knowledgeConnector.memberSyncStatus, 'disabled'), + connectorIsLive() + ) + ) + .for('share') + if (!connector) throw new StaleSweepConnectorIneligible() + return { + observationsRemoved: documentIds.length, + documentsRematerialized: rematerialized, + docsTombstoned: tombstoned.length, + more: candidates.length === ACL_CHANGE_BATCH_SIZE || rest.length > 0, + } + }) + } catch (error) { + if (error instanceof StaleSweepConnectorIneligible) return null + throw error + } +} + +/** Rolls a sweep page back when its connector no longer qualifies by the time the page commits. */ +class StaleSweepConnectorIneligible extends Error {} + /** * Removes the observations of members whose crawls have stopped, so the * documents only they observed go dark instead of staying readable forever. @@ -783,15 +1133,21 @@ function memberStillStale(memberId: string, cutoff: Date) { * lists for them rebuilds their observations. Purging is left to a run holding * the lease. * - * Each member is swept in one transaction that first shares the connector row - * — which a member run holds `FOR UPDATE` while it writes and a mode switch - * updates when it flips — and then locks the member row, which `claimNextMember` - * skips while locked. Both are re-checked under those locks, so a run that - * claimed the member after the selection, or a switch that left members mode, - * makes the sweep skip rather than delete observations a run just wrote or - * rewrite ACLs the switch just set. + * Each page of a member's sweep is one bounded transaction that first locks the + * member row, which `claimNextMember` skips while locked, and last shares the + * connector row, which a member run's page locks `FOR UPDATE` before it commits + * and a mode switch updates when it flips. Both are re-checked under those + * locks, so a run that claimed the member after the selection, or a switch that + * left members mode, rolls the page back rather than delete observations a run + * just wrote or rewrite ACLs the switch just set; the connector row is never + * held while the page waits on document rows. A page that hits a lock or + * statement bound is left for the next tick, and a tick stops after + * `STALE_MEMBER_SWEEP_BUDGET_MS`. */ -export async function sweepStaleMemberObservations(now: Date): Promise { +export async function sweepStaleMemberObservations( + now: Date, + deadlineAt: number = Date.now() + STALE_MEMBER_SWEEP_BUDGET_MS +): Promise { const staleWindow = sql`GREATEST( ${MEMBER_OBSERVATION_STALE_AFTER_HOURS} * INTERVAL '1 hour', 2 * ${knowledgeConnector.syncIntervalMinutes} * INTERVAL '1 minute' @@ -851,74 +1207,41 @@ export async function sweepStaleMemberObservations(now: Date): Promise() for (const member of staleMembers) { + if (Date.now() >= deadlineAt) break + if (deferredConnectors.has(member.connectorId)) continue const memberCutoff = new Date(now.getTime() - staleMemberWindowMs(member.syncIntervalMinutes)) - const swept = await db.transaction(async (tx) => { - const [connector] = await tx - .select({ id: knowledgeConnector.id }) - .from(knowledgeConnector) - .where( - and( - eq(knowledgeConnector.id, member.connectorId), - eq(knowledgeConnector.accessMode, 'members'), - inArray(knowledgeConnector.status, MEMBER_LOCKABLE_CONNECTOR_STATUSES), - ne(knowledgeConnector.memberSyncStatus, 'disabled'), - connectorIsLive() - ) - ) - .for('share') - if (!connector) return null - const [stale] = await tx - .select({ id: knowledgeConnectorMember.id }) - .from(knowledgeConnectorMember) - .where(memberStillStale(member.id, memberCutoff)) - .for('update') - if (!stale) return null - - const candidates = tx - .select({ documentId: knowledgeDocumentObservation.documentId }) - .from(knowledgeDocumentObservation) - .where(eq(knowledgeDocumentObservation.memberId, member.id)) - .limit(OBSERVATION_BATCH_SIZE) - const removed = await tx - .delete(knowledgeDocumentObservation) - .where( - and( - eq(knowledgeDocumentObservation.memberId, member.id), - inArray(knowledgeDocumentObservation.documentId, candidates) - ) - ) - .returning({ documentId: knowledgeDocumentObservation.documentId }) - const documentIds = removed.map((row) => row.documentId) - const rematerialized = await materializeDocumentAcls(member.connectorId, documentIds, tx) - const tombstoned = - documentIds.length === 0 - ? [] - : await tx - .update(document) - .set({ deletedAt: now }) - .where( - and( - inArray(document.id, documentIds), - eq(document.connectorId, member.connectorId), - eq(document.userExcluded, false), - isNull(document.archivedAt), - isNull(document.deletedAt), - hasNoObservation() - ) - ) - .returning({ id: document.id }) - return { - observationsRemoved: documentIds.length, - documentsRematerialized: rematerialized, - docsTombstoned: tombstoned.length, + let sweptAny = false + try { + for (let page = 0; page < STALE_MEMBER_PAGES_PER_TICK; page++) { + /** Checked before every page, so one large member cannot run a tick past its budget. */ + if (Date.now() >= deadlineAt) break + const swept = await sweepStaleMemberPage(member, memberCutoff, now) + if (!swept) break + sweptAny = true + result.observationsRemoved += swept.observationsRemoved + result.documentsRematerialized += swept.documentsRematerialized + result.docsTombstoned += swept.docsTombstoned + if (!swept.more) break } - }) - if (!swept) continue - result.members += 1 - result.observationsRemoved += swept.observationsRemoved - result.documentsRematerialized += swept.documentsRematerialized - result.docsTombstoned += swept.docsTombstoned + } catch (error) { + /** + * A connector whose member run holds its row, or whose page outruns the bounds, is left + * for the next tick: its committed pages stand, and the other members are still swept. + */ + const failure = getTransientDatabaseFailure(error) + if (!failure) throw error + deferredConnectors.add(member.connectorId) + logger.warn('Deferred a stale member sweep to the next tick', { + connectorId: member.connectorId, + memberId: member.id, + failure, + code: getPostgresErrorCode(error), + }) + } + if (sweptAny) result.members += 1 } return result } diff --git a/apps/sim/lib/knowledge/connectors/member-sync-engine.integration.test.ts b/apps/sim/lib/knowledge/connectors/member-sync-engine.integration.test.ts index 5ac9820c142..6fea46424e3 100644 --- a/apps/sim/lib/knowledge/connectors/member-sync-engine.integration.test.ts +++ b/apps/sim/lib/knowledge/connectors/member-sync-engine.integration.test.ts @@ -18,6 +18,7 @@ const mocks = vi.hoisted(() => ({ removeUnseen: vi.fn(), removeForDocuments: vi.fn(), materialize: vi.fn(), + rematerialize: vi.fn(async () => 0), lifecycle: vi.fn(), credentials: vi.fn(), getChangeCursor: vi.fn(), @@ -60,6 +61,7 @@ vi.mock('@/lib/knowledge/connectors/member-access', () => ({ vi.mock('@/lib/knowledge/connectors/member-observations', () => ({ applyMemberDocumentLifecycle: mocks.lifecycle, materializeDocumentAcls: mocks.materialize, + rematerializeDocumentAcls: mocks.rematerialize, recordMemberObservations: mocks.observe, removeMemberObservationsForDocuments: mocks.removeForDocuments, removeUnseenMemberObservations: mocks.removeUnseen, @@ -487,6 +489,43 @@ describe('member engine with a dedicated content credential', () => { expect(mocks.observe).not.toHaveBeenCalled() }) + /** + * The content completion counts the whole connector. That scan runs before the lease + * transaction, which stays under the role's own timeouts, so a large connector that finished + * every content page cannot then fail its completion on a page bound. + */ + it('counts the connector before its content completion takes the lease', async () => { + const run = arrange({ members: true, noDueMembers: true }) + const result = await run() + expect(result.error).toBeUndefined() + const insertIndex = dbChainMockFns.insert.mock.calls.findIndex( + ([table]) => table === schemaMock.knowledgeConnectorSyncLog + ) + expect(insertIndex).toBeGreaterThanOrEqual(0) + const inserted = dbChainMockFns.insert.mock.invocationCallOrder[insertIndex] + const opened = Math.max( + ...dbChainMockFns.transaction.mock.invocationCallOrder.filter((order) => order < inserted) + ) + const counted = dbChainMockFns.select.mock.calls + .map(([fields], index) => ({ + fields, + order: dbChainMockFns.select.mock.invocationCallOrder[index], + })) + .filter(({ fields, order }) => fields && 'count' in fields && order < inserted) + .map(({ order }) => order) + expect(Math.max(...counted)).toBeLessThan(opened) + const bounded = dbChainMockFns.execute.mock.calls + .map((call: unknown[], index) => ({ + call, + order: dbChainMockFns.execute.mock.invocationCallOrder[index], + })) + .filter( + ({ call, order }) => + order > opened && order < inserted && JSON.stringify(call).includes('lock_timeout') + ) + expect(bounded).toEqual([]) + }) + it('reserves time for member permissions when a slow dedicated content page has more batches', async () => { const run = arrange({ members: true, contentIncomplete: true }) const now = Date.now() diff --git a/apps/sim/lib/knowledge/connectors/member-sync-engine.ts b/apps/sim/lib/knowledge/connectors/member-sync-engine.ts index f031c25cfc8..64e4de04f02 100644 --- a/apps/sim/lib/knowledge/connectors/member-sync-engine.ts +++ b/apps/sim/lib/knowledge/connectors/member-sync-engine.ts @@ -12,6 +12,7 @@ import { } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { getErrorMessage, getTransientDatabaseFailure, toError } from '@sim/utils/errors' +import { chunkArray } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' import { randomInt } from '@sim/utils/random' import { and, asc, eq, gt, inArray, isNull, lte, notExists, sql } from 'drizzle-orm' @@ -54,8 +55,10 @@ import { } from '@/lib/knowledge/connectors/member-access' import { applyMemberDocumentLifecycle, + lockProjectionPage, materializeDocumentAcls, recordMemberObservations, + rematerializeDocumentAcls, removeMemberObservationsForDocuments, removeUnseenMemberObservations, renewMemberObservationsInScopes, @@ -71,6 +74,8 @@ import { getConnectorSyncDeferral, } from '@/lib/knowledge/connectors/sync-deferral' import { + ACL_CHANGE_BATCH_SIZE, + ACL_WRITE_BATCH_SIZE, CONNECTOR_AUTO_DISABLED_ERROR, CONNECTOR_FAILURE_BACKOFF_CAP_MINUTES, connectorFailureBackoffMinutes, @@ -87,8 +92,10 @@ import { } from '@/lib/knowledge/connectors/sync-limits' import { assertSyncLeaseHeldInTx, + boundLeaseTransaction, createMemberSyncLease, holdsMemberSyncLockToken, + type LeaseTransaction, MEMBER_LOCKABLE_CONNECTOR_STATUSES, SyncLockLostException, stillHoldsMemberSyncLock, @@ -513,24 +520,32 @@ function createMemberTokenCache(input: { } /** - * Runs `fn` in a transaction that first proves this run still holds the - * connector's member lease, taking the connector row's lock so the scheduler - * cannot reclaim the lease mid-transaction. A run that stalled past the lease + * Runs `fn` in a transaction whose last statement proves this run still holds + * the connector's member lease, taking the connector row's lock so the + * scheduler cannot reclaim the lease before the transaction commits; a lost + * lease rolls everything `fn` wrote back. Proving it last keeps the connector + * row unlocked while `fn` waits on document rows. A run that stalled past the lease * TTL and resumed after a replacement took over therefore never lands its - * observations or ACLs over the replacement's; it ends as superseded. + * observations or ACLs over the replacement's; it ends as superseded. A page + * that assigns ACLs passes `aclPage`: it writes at most one page, which fires + * the projection fan-out, so it takes the lock and statement bounds of every + * connector-lease ACL page. Other bodies keep the role's own timeouts. */ async function withMemberLease( run: Pick, - fn: (tx: DbOrTx) => Promise + fn: (tx: DbOrTx) => Promise, + options: { aclPage?: boolean } = {} ): Promise { return db.transaction(async (tx) => { + if (options.aclPage) await boundLeaseTransaction(tx) + const written = await fn(tx) const [held] = await tx .select({ id: knowledgeConnector.id }) .from(knowledgeConnector) .where(stillHoldsMemberSyncLock(run.connectorId, run.runId)) .for('update') if (!held) throw new SyncLockLostException(run.connectorId) - return fn(tx) + return written }) } @@ -652,56 +667,71 @@ export async function resumeMembershipRewrites( if (!member) return true const checkpoint = membershipRewrite(member.checkpoint) if (!checkpoint) throw new Error('Invalid membership ACL checkpoint') - await withMemberLease(run, async (tx) => { - const documents = await tx - .select({ documentId: knowledgeDocumentObservation.documentId }) - .from(knowledgeDocumentObservation) - .where( - and( - eq(knowledgeDocumentObservation.memberId, member.id), - checkpoint.cursor - ? gt(knowledgeDocumentObservation.documentId, checkpoint.cursor) - : undefined + await withMemberLease( + run, + async (tx) => { + const documents = await tx + .select({ documentId: knowledgeDocumentObservation.documentId }) + .from(knowledgeDocumentObservation) + .where( + and( + eq(knowledgeDocumentObservation.memberId, member.id), + checkpoint.cursor + ? gt(knowledgeDocumentObservation.documentId, checkpoint.cursor) + : undefined + ) ) - ) - .orderBy(asc(knowledgeDocumentObservation.documentId)) - .limit(500) - await materializeDocumentAcls( - run.connectorId, - documents.map((row) => row.documentId), - tx - ) - if (checkpoint.removeMember && run.tombstonesUnobserved && documents.length > 0) { - const tombstoned = await tombstoneDocumentsObservedOnlyBy( + .orderBy(asc(knowledgeDocumentObservation.documentId)) + .limit(ACL_CHANGE_BATCH_SIZE) + /** + * The page is the leading documents whose chunks fit one page of projection rows, so the + * cursor only ever passes documents whose ACLs this transaction rewrote. + */ + const { page } = await lockProjectionPage( tx, - run.connectorId, - member.id, documents.map((row) => row.documentId) ) - if (run.result) run.result.docsTombstoned += tombstoned - } - if (!checkpoint.removeMember && run.tombstonesUnobserved && documents.length > 0) { - const resurrected = await resurrectObservedDocuments( - tx, + const pageDocuments = documents.slice(0, page.length) + await materializeDocumentAcls( run.connectorId, - documents.map((row) => row.documentId) + pageDocuments.map((row) => row.documentId), + tx ) - if (run.result) run.result.docsResurrected += resurrected - } - if (documents.length === 0 && checkpoint.removeMember) { - await tx.delete(knowledgeConnectorMember).where(eq(knowledgeConnectorMember.id, member.id)) - } else { - await tx - .update(knowledgeConnectorMember) - .set({ - listingCheckpoint: - documents.length === 0 - ? null - : { ...checkpoint, cursor: documents.at(-1)!.documentId }, - }) - .where(eq(knowledgeConnectorMember.id, member.id)) - } - }) + if (checkpoint.removeMember && run.tombstonesUnobserved && pageDocuments.length > 0) { + const tombstoned = await tombstoneDocumentsObservedOnlyBy( + tx, + run.connectorId, + member.id, + pageDocuments.map((row) => row.documentId) + ) + if (run.result) run.result.docsTombstoned += tombstoned + } + if (!checkpoint.removeMember && run.tombstonesUnobserved && pageDocuments.length > 0) { + const resurrected = await resurrectObservedDocuments( + tx, + run.connectorId, + pageDocuments.map((row) => row.documentId) + ) + if (run.result) run.result.docsResurrected += resurrected + } + if (documents.length === 0 && checkpoint.removeMember) { + await tx + .delete(knowledgeConnectorMember) + .where(eq(knowledgeConnectorMember.id, member.id)) + } else { + await tx + .update(knowledgeConnectorMember) + .set({ + listingCheckpoint: + documents.length === 0 + ? null + : { ...checkpoint, cursor: pageDocuments.at(-1)!.documentId }, + }) + .where(eq(knowledgeConnectorMember.id, member.id)) + } + }, + { aclPage: true } + ) } } @@ -1481,16 +1511,19 @@ async function applyMemberListing( break } await run.lease.beatIfDue() - const batch = await withMemberLease(run, (tx) => - removeUnseenMemberObservations( - tx, - outcome.member.id, - outcome.observationRunId ?? run.runId, - async (removed) => { - await materializeDocumentAcls(run.connectorId, removed, tx) - for (const documentId of removed) run.unobservedDocumentIds.add(documentId) - } - ) + const batch = await withMemberLease( + run, + (tx) => + removeUnseenMemberObservations( + tx, + outcome.member.id, + outcome.observationRunId ?? run.runId, + async (removed) => { + await materializeDocumentAcls(run.connectorId, removed, tx) + for (const documentId of removed) run.unobservedDocumentIds.add(documentId) + } + ), + { aclPage: true } ) run.result.observationsRemoved += batch.removed if (batch.finished) break @@ -1649,18 +1682,19 @@ async function syncDedicatedMemberContent(input: { !pass.complete || pass.checkpoint.unsafe || pass.checkpoint.contentFailures if (pass.checkpoint.contentFailures) run.result.listingIncomplete = true const contentNotice = pass.holdNotice - await withMemberLease(run, async (tx) => { - const [{ count }] = await tx - .select({ count: sql`count(*)::int` }) - .from(document) - .where( - and( - eq(document.connectorId, run.connectorId), - eq(document.userExcluded, false), - isNull(document.archivedAt), - isNull(document.deletedAt) - ) + /** Counted before the lease transaction: a scan of the whole connector never holds its row. */ + const [{ count }] = await db + .select({ count: sql`count(*)::int` }) + .from(document) + .where( + and( + eq(document.connectorId, run.connectorId), + eq(document.userExcluded, false), + isNull(document.archivedAt), + isNull(document.deletedAt) ) + ) + await withMemberLease(run, async (tx) => { const now = new Date() await tx.insert(knowledgeConnectorSyncLog).values({ id: run.runId, @@ -1866,12 +1900,21 @@ async function deferMemberSync(run: MemberSyncRun, syncIntervalMinutes: number): * group binding is gone, and suspends every member so their tokens leave * every ACL. Nothing is purged: re-enabling restores access from the retained * observations. + * + * Suspension lands first: a reader needs an active member's observation as + * well as an overlapping ACL, so suspending the members revokes their reads at + * once, and anything that rematerialises an ACL meanwhile computes nobody. + * Every ACL is then revoked in short lease-proving pages, and only after the + * last page does the connector flip to disabled: one statement over the whole + * connector outlasted the statement timeout, rolled back, and retried forever. + * No reader gains access between pages, and an interrupted run leaves the + * binding still gone, which sends the next run back here to finish. Returns false when the run's budget + * ended first; the caller re-dispatches. */ -async function disableMemberSync(run: MemberSyncRun, reason: string): Promise { +async function disableMemberSync(run: MemberSyncRun, reason: string): Promise { const now = new Date() - /** Suspension, the ACLs it changes, and the disable itself land together, and only under the lease. */ - await withMemberLease(run, async (tx) => { - await tx + await withMemberLease(run, (tx) => + tx .update(knowledgeConnectorMember) .set({ status: 'suspended', suspendedAt: now, updatedAt: now }) .where( @@ -1880,24 +1923,49 @@ async function disableMemberSync(run: MemberSyncRun, reason: string): Promise { + run.result.membersRemaining = true + const landed = await completeMemberSync(run, syncIntervalMinutes) + return landed ? run.result : skipped(run.result, 'sync_superseded') } /** @@ -2041,7 +2109,13 @@ export async function executeMemberSync( } } if (!connector.credentialGroupId || !connector.credentialGroupOptionId) { - await disableMemberSync(run, 'Connector is no longer attached to a Credential Group option') + if ( + !(await disableMemberSync( + run, + 'Connector is no longer attached to a Credential Group option' + )) + ) + return finishDisableLater(run, connector.syncIntervalMinutes) return { ...skipped(result, 'connector_not_syncable'), error: 'Connector is no longer attached to a Credential Group option', @@ -2139,6 +2213,8 @@ export async function executeMemberSync( } const credentialIdByMemberId = new Map() + /** One bounded ACL page under this run's lease. */ + const aclPage: LeaseTransaction = (fn) => withMemberLease(run, fn, { aclPage: true }) const tokens = createMemberTokenCache({ run, connectorConfig, @@ -2174,34 +2250,39 @@ export async function executeMemberSync( ) ).values(), ] - await withMemberLease(run, async (tx) => { - result.observationsAdded += await recordMemberObservations( - tx, - member.id, - documentIds, - checkpoint.generationId - ) - await materializeDocumentAcls(connectorId, documentIds, tx) - if (durableCheckpoint && checkpoint.contentFailures) { - await tx - .update(knowledgeConnectorMember) - .set({ listingCheckpoint: checkpoint }) - .where(eq(knowledgeConnectorMember.id, member.id)) - } - if (!serviceContent) { - for (let offset = 0; offset < documentIds.length; offset += 500) { + /** + * Observations and the read watermark land first, one lease transaction per + * {@link ACL_WRITE_BATCH_SIZE} documents (the failure checkpoint with the + * first, since forcing a later relist is the conservative side); only the + * documents whose ACL now differs are then rematerialised, in pages bounded + * by their projection rows. A crash in between leaves them hidden until the + * next run rematerialises every seen document, never shown too widely. + */ + const pages = + documentIds.length > 0 ? chunkArray(documentIds, ACL_WRITE_BATCH_SIZE) : [[]] + for (const [index, page] of pages.entries()) { + await withMemberLease(run, async (tx) => { + result.observationsAdded += await recordMemberObservations( + tx, + member.id, + page, + checkpoint.generationId + ) + if (index === 0 && durableCheckpoint && checkpoint.contentFailures) { + await tx + .update(knowledgeConnectorMember) + .set({ listingCheckpoint: checkpoint }) + .where(eq(knowledgeConnectorMember.id, member.id)) + } + if (!serviceContent && page.length > 0) { await tx .update(document) .set({ sourceSeenAt: run.runStartedAt }) - .where( - and( - eq(document.connectorId, connectorId), - inArray(document.id, documentIds.slice(offset, offset + 500)) - ) - ) + .where(and(eq(document.connectorId, connectorId), inArray(document.id, page))) } - } - }) + }) + } + await rematerializeDocumentAcls(connectorId, documentIds, aclPage, run.lease.beatIfDue) result.docsListed += attempted.length } if (!serviceContent) { @@ -2374,7 +2455,7 @@ export async function executeMemberSync( await loadDocumentIdsByExternalId(connectorId, relevantIds), connector.syncIntervalMinutes ) - await withMemberLease(run, (tx) => materializeDocumentAcls(connectorId, affected, tx)) + await rematerializeDocumentAcls(connectorId, affected, aclPage, run.lease.beatIfDue) } /** A service-owned corpus outlives its last observer; only the content pass removes it. */ @@ -2440,7 +2521,9 @@ export async function executeMemberSync( } logger.info('Member sync completed', { connectorId, runId, ...result }) return result - } catch (error) { + } catch (caught) { + /** A failed disable of a gone binding replaces the error the failure path records. */ + let error: unknown = caught if (error instanceof SyncLockLostException) { logger.warn('Member sync abandoned — lock was reclaimed while this run was executing', { connectorId, @@ -2459,17 +2542,27 @@ export async function executeMemberSync( return skipped(result, 'connector_deleted_during_sync') } if (error instanceof MemberBindingGoneError) { + const bindingError = error try { - await disableMemberSync(run, error.message) + if (!(await disableMemberSync(run, bindingError.message))) + return await finishDisableLater(run, connector.syncIntervalMinutes) + return { ...skipped(result, 'connector_not_syncable'), error: bindingError.message } } catch (disableError) { - if (!(disableError instanceof SyncLockLostException)) throw disableError - logger.warn('Member sync abandoned — lock was reclaimed before it could be disabled', { - connectorId, - runId, - }) - return skipped(result, 'sync_superseded') + if (disableError instanceof SyncLockLostException) { + logger.warn('Member sync abandoned — lock was reclaimed before it could be disabled', { + connectorId, + runId, + }) + return skipped(result, 'sync_superseded') + } + /** + * A disable that failed part-way is an ordinary run failure: the log closes as failed + * and the lease is released under its own guard, instead of the run escaping with the + * connector left running until the stale-lease reclaim. The binding is still gone, so + * the next run resumes the revocation. + */ + error = disableError } - return { ...skipped(result, 'connector_not_syncable'), error: error.message } } if (getConnectorSyncDeferral(error)) { diff --git a/apps/sim/lib/knowledge/connectors/sync-content-pass.postgres.test.ts b/apps/sim/lib/knowledge/connectors/sync-content-pass.postgres.test.ts index 49b7fdfd92a..98b9ff1ed41 100644 --- a/apps/sim/lib/knowledge/connectors/sync-content-pass.postgres.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-content-pass.postgres.test.ts @@ -32,6 +32,7 @@ const schema = await import('@sim/db/schema') const { beginListingCheckpoint } = await import('@/lib/knowledge/connectors/listing-checkpoint') const { runConnectorContentPass } = await import('@/lib/knowledge/connectors/sync-content-pass') const { revokeDocumentAcls } = await import('@/lib/knowledge/connectors/sync-persistence') +const { leaseTransaction } = await import('@/lib/knowledge/connectors/sync-lock') const { confluenceConnector } = await import('@/connectors/confluence/confluence') const databaseUrl = process.env.KNOWLEDGE_ACL_TEST_DATABASE_URL @@ -156,7 +157,7 @@ describe.runIf(Boolean(databaseUrl))('completed listing reconciliation in Postgr holder.db = drizzle(sql, { schema }) await sql`CREATE TABLE knowledge_connector (id text PRIMARY KEY)` await sql`CREATE TABLE document ( - id text PRIMARY KEY, external_id text, connector_id text, + id text PRIMARY KEY, external_id text, connector_id text, chunk_count integer NOT NULL DEFAULT 1, user_excluded boolean NOT NULL DEFAULT false, archived_at timestamp, deleted_at timestamp, source_seen_at timestamp, acl text[] NOT NULL DEFAULT '{ws}', acl_requirements jsonb NOT NULL DEFAULT '[]', @@ -274,8 +275,10 @@ describe.runIf(Boolean(databaseUrl))('completed listing reconciliation in Postgr await sql`UPDATE document SET acl_requirements = '[[], ["g:confluence:tenant:space"]]' WHERE id = 'stale-evidence'` - await revokeDocumentAcls(holder.db as never, ['stale-evidence', 'granted'], (batch) => - inArray(schema.document.id, batch) + await revokeDocumentAcls( + leaseTransaction(CONNECTOR, undefined, holder.db as never), + ['stale-evidence', 'granted'], + (batch) => inArray(schema.document.id, batch) ) expect(await fannedOut()).toEqual(['granted']) diff --git a/apps/sim/lib/knowledge/connectors/sync-content-pass.test.ts b/apps/sim/lib/knowledge/connectors/sync-content-pass.test.ts index 967dbe31b71..90e377ddad8 100644 --- a/apps/sim/lib/knowledge/connectors/sync-content-pass.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-content-pass.test.ts @@ -208,7 +208,16 @@ describe('completed listing removal counts', () => { ]) queueTableRows(schemaMock.document, options.revoked ?? []) if (options.revoked?.length) { - queueTableRows(schemaMock.knowledgeConnector, [{ id: 'connector' }]) + /** Each window reads what still grants someone; pages are bounded by their chunks' rows. */ + const granting = options.revoked.map(({ id }) => ({ id, chunkCount: 10 })) + queueTableRows(schemaMock.document, granting) + /** Each revocation page locks its documents and rereads their chunks before writing. */ + for (let offset = 0; offset < granting.length; offset += 25) + queueTableRows(schemaMock.document, granting.slice(offset, offset + 25)) + /** Every revocation page is its own lease-proving transaction: one evidence clear, then acl pages. */ + const batches = 1 + Math.ceil(options.revoked.length / 25) + for (let batch = 0; batch < batches; batch++) + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'connector' }]) queueTableRows(schemaMock.document, []) } if (!options.fullSync) { @@ -317,11 +326,16 @@ describe('completed listing removal counts', () => { expect( writes.map( ({ conditions }) => - (conditions.find((node) => node.type === 'inArray')?.values as string[]).length + (conditions.filter((node) => node.type === 'inArray').at(-1)?.values as string[]).length ) ).toEqual([25, 5]) expect(writes.every(({ conditions }) => conditions.some(grantsSomeone))).toBe(true) expect(result.docsDeleted).toBe(0) + /** One lease-proving, bounded transaction per revocation batch, never one across the page. */ + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(3) + expect(dbChainMockFns.for.mock.calls.filter(([mode]) => mode === 'share')).toHaveLength(3) + /** Each acl page locks its own documents first, and only those. */ + expect(dbChainMockFns.for.mock.calls.filter(([mode]) => mode === 'update')).toHaveLength(2) }) it('does not report a full-sync removal when the guarded delete removed no live rows', async () => { @@ -364,6 +378,12 @@ async function runPass( } queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb' }]) queueTableRows(schemaMock.document, options.existing ? [options.existing] : []) + /** A permission-only page revokes a changed body first: its read of what still grants someone. */ + if (options.permissionsOnly && options.existing && options.existing.contentHash === 'old-body') { + queueTableRows(schemaMock.document, [{ id: options.existing.id, chunkCount: 1 }]) + /** The revocation page locks the document and rereads its chunks before writing. */ + queueTableRows(schemaMock.document, [{ id: options.existing.id, chunkCount: 1 }]) + } if (options.readCurrent) { queueTableRows(schemaMock.document, [{ fileUrl: options.existing?.fileUrl ?? '' }]) if ( @@ -828,6 +848,19 @@ describe('permission refresh through the shared content pass', () => { const revocations = aclAssignments().filter(({ values }) => values.acl.length === 0) expect(revocations).toHaveLength(1) expect(revocations[0].conditions.some(grantsSomeone)).toBe(true) + /** The evidence clear and the revocation are separate lease transactions, never one across both. */ + const setOrder = (matches: (values: Record) => boolean) => + dbChainMockFns.set.mock.invocationCallOrder[ + dbChainMockFns.set.mock.calls.findIndex(([values]) => matches(values)) + ] + const cleared = setOrder((values) => !('acl' in values) && 'aclRequirements' in values) + const revoked = setOrder((values) => Array.isArray(values.acl) && values.acl.length === 0) + expect(cleared).toBeLessThan(revoked) + expect( + dbChainMockFns.transaction.mock.invocationCallOrder.some( + (order) => order > cleared && order < revoked + ) + ).toBe(true) }) it('never renews a changed body after its hydration fails', async () => { diff --git a/apps/sim/lib/knowledge/connectors/sync-content-pass.ts b/apps/sim/lib/knowledge/connectors/sync-content-pass.ts index 86f515bd805..b06e47c8d83 100644 --- a/apps/sim/lib/knowledge/connectors/sync-content-pass.ts +++ b/apps/sim/lib/knowledge/connectors/sync-content-pass.ts @@ -26,7 +26,12 @@ import { SOURCE_CONTENT_ERROR, SOURCE_PERMISSION_ERROR, } from '@/lib/knowledge/connectors/sync-limits' -import { assertSyncLeaseHeldInTx, type SyncRunLease } from '@/lib/knowledge/connectors/sync-lock' +import { + assertSyncLeaseHeldInTx, + type LeaseTransaction, + leaseTransaction, + type SyncRunLease, +} from '@/lib/knowledge/connectors/sync-lock' import { type KnowledgeBaseOwner, persistSourceDocumentFailures, @@ -87,11 +92,15 @@ interface ContentPassInput { /** One durable content cycle shared by content-owned and member-visibility connectors. */ export async function runConnectorContentPass(input: ContentPassInput) { const { matchContentHash } = input.connectorConfig - const withLease = (fn: (tx: DbOrTx) => Promise) => + /** The lease is proved last, so no write waits on document rows while holding the connector row. */ + const withLease: LeaseTransaction = (fn) => db.transaction(async (tx) => { + const written = await fn(tx) await assertSyncLeaseHeldInTx(tx, input.connectorId, input.lease) - return fn(tx) + return written }) + /** ACL revocations fire the projection fan-out, so each of their pages is also bounded. */ + const withAclPage = leaseTransaction(input.connectorId, input.lease) const readGenerationStartedAt = async (tx: DbOrTx): Promise => { const [clock] = await tx.execute<{ startedAt: string }>( sql`SELECT statement_timestamp()::text AS "startedAt"` @@ -180,17 +189,15 @@ export async function runConnectorContentPass(input: ContentPassInput) { }) /** Revoke grants without matching stored content, retaining the content crawl's observation for EOF reconciliation. */ if (changed.length) - await withLease((tx) => - revokeDocumentAcls( - tx, - changed.map((item) => item.externalId), - (batch) => - and( - eq(document.connectorId, input.connectorId), - inArray(document.externalId, batch), - isNull(document.archivedAt) - ) - ) + await revokeDocumentAcls( + withAclPage, + changed.map((item) => item.externalId), + (batch) => + and( + eq(document.connectorId, input.connectorId), + inArray(document.externalId, batch), + isNull(document.archivedAt) + ) ) } const state = createSyncRunState(input.result) @@ -326,7 +333,7 @@ export async function runConnectorContentPass(input: ContentPassInput) { }, }) const reconciliation = checkpoint.complete - ? await reconcileCompletedListing(input, checkpoint, withLease) + ? await reconcileCompletedListing(input, checkpoint, withLease, withAclPage) : { finished: false, notice: null } /** Unverified permissions, an incomplete listing and unrefreshed content are independent holds; an admin needs each, one per line. */ const holdNotice = @@ -349,7 +356,8 @@ export async function runConnectorContentPass(input: ContentPassInput) { async function reconcileCompletedListing( input: ContentPassInput, checkpoint: ListingCheckpoint, - withLease: (fn: (tx: DbOrTx) => Promise) => Promise + withLease: LeaseTransaction, + withAclPage: LeaseTransaction ): Promise<{ finished: boolean; notice: string | null }> { if (checkpoint.unsafe || (checkpoint.listingFailures?.count ?? 0) > 0) return { @@ -426,12 +434,10 @@ async function reconcileCompletedListing( await input.lease.beatIfDue() const rows = await loadBatch(and(absent, sql`cardinality(${document.acl}) > 0`), 500, after) if (rows.length === 0) break - await withLease((tx) => - revokeDocumentAcls( - tx, - rows.map((row) => row.id), - (batch) => and(absent, inArray(document.id, batch)) - ) + await revokeDocumentAcls( + withAclPage, + rows.map((row) => row.id), + (batch) => and(absent, inArray(document.id, batch)) ) after = rows.at(-1) } diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts index 20968907f68..0a933b61738 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts @@ -2384,7 +2384,6 @@ describe('completeSuccessfulSync', () => { queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) queueTableRows(schemaMock.document, [{ count: 4 }]) dbChainMockFns.returning - .mockResolvedValueOnce([]) .mockResolvedValueOnce([{ id: 'log-1' }]) .mockResolvedValueOnce([{ id: 'c-1' }]) const directoryNotice = @@ -2448,13 +2447,55 @@ describe('completeSuccessfulSync', () => { } ) + /** An unfinished pending rewrite keeps its flag and comes straight back, instead of clearing it. */ + it('keeps the pending access rewrite and re-runs at once when the walk did not finish', async () => { + const { completeSuccessfulSync } = await import('@/lib/knowledge/connectors/sync-engine') + queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-1' }]) + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) + queueTableRows(schemaMock.document, [{ count: 4 }]) + dbChainMockFns.returning + .mockResolvedValueOnce([{ id: 'log-1' }]) + .mockResolvedValueOnce([{ id: 'c-1' }]) + + await expect( + completeSuccessfulSync('c-1', 'kb-1', 'log-1', 60, RESULT, null, undefined, null, true) + ).resolves.toBe(true) + + const logUpdate = dbChainMockFns.set.mock.calls + .map(([value]) => value as Record) + .find((value) => 'completedAt' in value) + const connectorUpdate = dbChainMockFns.set.mock.calls + .map(([value]) => value as Record) + .find((value) => value.status === 'active') + expect(logUpdate?.status).toBe('partial') + expect(connectorUpdate).not.toHaveProperty('accessRewritePending') + expect((connectorUpdate?.nextSyncAt as Date).getTime()).toBeLessThanOrEqual(Date.now()) + }) + + it('clears the pending access rewrite once the walk finished', async () => { + const { completeSuccessfulSync } = await import('@/lib/knowledge/connectors/sync-engine') + queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-1' }]) + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) + queueTableRows(schemaMock.document, [{ count: 4 }]) + dbChainMockFns.returning + .mockResolvedValueOnce([{ id: 'log-1' }]) + .mockResolvedValueOnce([{ id: 'c-1' }]) + + await expect(completeSuccessfulSync('c-1', 'kb-1', 'log-1', 60, RESULT, null)).resolves.toBe( + true + ) + + expect(dbChainMockFns.set).toHaveBeenCalledWith( + expect.objectContaining({ status: 'active', accessRewritePending: false }) + ) + }) + it('counts the documents before taking the completion locks', async () => { const { completeSuccessfulSync } = await import('@/lib/knowledge/connectors/sync-engine') queueTableRows(schemaMock.knowledgeBase, [{ id: 'kb-1' }]) queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) queueTableRows(schemaMock.document, [{ count: 4 }]) dbChainMockFns.returning - .mockResolvedValueOnce([]) .mockResolvedValueOnce([{ id: 'log-1' }]) .mockResolvedValueOnce([{ id: 'c-1' }]) @@ -2487,7 +2528,6 @@ describe('completeSuccessfulSync', () => { ) ) dbChainMockFns.returning - .mockResolvedValueOnce([]) .mockResolvedValueOnce([{ id: 'log-1' }]) .mockResolvedValueOnce([{ id: 'c-1' }]) @@ -2520,8 +2560,6 @@ describe('completeSuccessfulSync', () => { queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) queueTableRows(schemaMock.document, [{ count: 4 }]) dbChainMockFns.returning - /** The workspace ACL restore finds nothing drifted. */ - .mockResolvedValueOnce([]) .mockResolvedValueOnce([{ id: 'log-1' }]) .mockResolvedValueOnce([{ id: 'c-1' }]) @@ -2561,7 +2599,6 @@ describe('completeSuccessfulSync', () => { queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) queueTableRows(schemaMock.document, [{ count: 4 }]) dbChainMockFns.returning - .mockResolvedValueOnce([]) .mockResolvedValueOnce([{ id: 'log-1' }]) .mockResolvedValueOnce([{ id: 'c-1' }]) @@ -2608,7 +2645,6 @@ describe('completeSuccessfulSync', () => { queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) queueTableRows(schemaMock.document, [{ count: 4 }]) dbChainMockFns.returning - .mockResolvedValueOnce([]) .mockResolvedValueOnce([{ id: 'log-1' }]) .mockResolvedValueOnce([{ id: 'c-1' }]) const holdNotice = 'Source listing is incomplete; unlisted documents were kept.' @@ -2676,7 +2712,6 @@ describe('completeSuccessfulSync', () => { queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) queueTableRows(schemaMock.document, [{ count: 4 }]) dbChainMockFns.returning - .mockResolvedValueOnce([]) .mockResolvedValueOnce([{ id: 'log-1' }]) .mockResolvedValueOnce([{ id: 'c-1' }]) @@ -3386,6 +3421,9 @@ describe('executeSync heartbeats during the listing phase', () => { primeSyncUpToListing() dbChainMockFns.returning.mockReset() dbChainMockFns.returning.mockResolvedValueOnce([{ ...CONNECTOR, accessMode: 'admin' }]) + /** No tombstone, then the stored document whose ACL the mirrored write changes. */ + queueTableRows(schemaMock.document, []) + queueTableRows(schemaMock.document, [{ id: 'doc-1', chunkCount: 1 }]) let permissionResult: { permissionsIncomplete: boolean } | undefined const pass = vi .spyOn(contentPass, 'runConnectorContentPass') @@ -3413,6 +3451,68 @@ describe('executeSync heartbeats during the listing phase', () => { } ) + it('proves the lease last inside each bounded transaction that writes mirrored permissions', async () => { + const contentPass = await import('@/lib/knowledge/connectors/sync-content-pass') + primeSyncUpToListing() + dbChainMockFns.returning.mockReset() + dbChainMockFns.returning.mockResolvedValueOnce([{ ...CONNECTOR, accessMode: 'admin' }]) + queueTableRows(schemaMock.document, []) + queueTableRows(schemaMock.document, [{ id: 'doc-1', chunkCount: 1 }]) + const pass = vi + .spyOn(contentPass, 'runConnectorContentPass') + .mockImplementation(async (input) => { + await input.onPage?.( + [ + { + externalId: 'page-1', + title: 'Page', + content: 'Body', + contentHash: 'hash-1', + mimeType: 'text/plain', + acl: ['u:reader@example.com'], + }, + ], + new Date() + ) + throw new Error('Stopped after permission persistence') + }) + try { + await executeSync('c-1', { billingAttribution: { workspaceId: 'ws-1' } as never }) + const writeIndex = dbChainMockFns.set.mock.calls.findIndex(([values]) => 'acl' in values) + expect(writeIndex).toBeGreaterThanOrEqual(0) + const written = dbChainMockFns.set.mock.invocationCallOrder[writeIndex] + const opened = Math.max( + ...dbChainMockFns.transaction.mock.invocationCallOrder.filter((order) => order < written) + ) + const between = (orders: number[]) => + orders.some((order) => order > opened && order < written) + const leaseChecks = dbChainMockFns.for.mock.calls + .map(([mode], index) => ({ + mode, + order: dbChainMockFns.for.mock.invocationCallOrder[index], + })) + .filter(({ mode }) => mode === 'share') + .map(({ order }) => order) + /** Proved after the write and before the next transaction opens: the page's last statement. */ + const closed = Math.min( + ...dbChainMockFns.transaction.mock.invocationCallOrder.filter((order) => order > written), + Number.POSITIVE_INFINITY + ) + expect(leaseChecks.some((order) => order > written && order < closed)).toBe(true) + expect(between(leaseChecks)).toBe(false) + const bounds = dbChainMockFns.execute.mock.calls + .map((query: unknown[], index) => ({ + query, + order: dbChainMockFns.execute.mock.invocationCallOrder[index], + })) + .filter(({ query }) => JSON.stringify(query).includes('lock_timeout')) + .map(({ order }) => order) + expect(between(bounds)).toBe(true) + } finally { + pass.mockRestore() + } + }) + it('beats between pages and abandons the run when the lock was reclaimed', async () => { const { executeSync } = await import('@/lib/knowledge/connectors/sync-engine') const { SYNC_LOCK_HEARTBEAT_INTERVAL_MS } = await import( diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.ts b/apps/sim/lib/knowledge/connectors/sync-engine.ts index d8545ff6b32..dca69a75090 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.ts @@ -69,6 +69,7 @@ import { createContentSyncLease, holdsSyncLockToken, LOCKABLE_CONNECTOR_STATUSES, + leaseTransaction, RUNNABLE_CONNECTOR_STATUSES, SyncLockLostException, type SyncRunLease, @@ -187,12 +188,13 @@ async function applySourceMirroredAcls(input: { */ const unlisted = hideUnlistedDocuments(acls, input.ownedExternalIds) - const written = input.lease - ? await db.transaction(async (tx) => { - await assertSyncLeaseHeldInTx(tx, connectorId, input.lease!) - return persistDocumentAcls(connectorId, acls, tx, evidence) - }) - : await persistDocumentAcls(connectorId, acls, db, evidence) + /** One short transaction per batch, each proving the lease, rather than one across all of them. */ + const written = await persistDocumentAcls( + connectorId, + acls, + leaseTransaction(connectorId, input.lease), + evidence + ) logger.info('Mirrored source permissions onto connector documents', { connectorId, listed, @@ -388,7 +390,9 @@ export async function completeSuccessfulSync( result: SyncResult, reconciliationHoldNotice: string | null, contentPass?: ContentPassOutcome, - directoryNotice: string | null = null + directoryNotice: string | null = null, + /** A pending ACL rewrite this run began but did not finish: the flag stays and the next run resumes it. */ + accessRewriteUnfinished = false ): Promise { const processingDispatchFailed = result.processingDispatch.failed > 0 const contentNotice = @@ -449,21 +453,6 @@ export async function completeSuccessfulSync( .for('update') if (!lockedConnector) throw new SyncCompletionOwnershipLost() - /** - * Self-healing invariant of workspace mode: a mode switch back from - * members that was interrupted, or any other drift, leaves no document - * of this connector hidden from the workspace once a sync completes. - * Inside the completion transaction, after the lock is proven held, so a - * reclaimed run cannot rewrite a connector that has since changed mode. - */ - const restoredAcls = await restoreWorkspaceDocumentAcls(tx, connectorId) - if (restoredAcls > 0) { - logger.warn('Restored workspace access on connector documents that had drifted', { - connectorId, - restoredAcls, - }) - } - const now = new Date() const [closedLog] = await tx .update(knowledgeConnectorSyncLog) @@ -471,6 +460,7 @@ export async function completeSuccessfulSync( status: directoryNotice || processingDispatchFailed || + accessRewriteUnfinished || (contentPass && isContentPassIncomplete(contentPass)) ? 'partial' : 'completed', @@ -509,16 +499,23 @@ export async function completeSuccessfulSync( ...buildSyncSuccessUpdate( now, actualDocCount, - contentPass && !contentPass.complete - ? contentPass.checkpoint.resumeAt - ? new Date(contentPass.checkpoint.resumeAt) - : now - : calculateNextSyncTime(syncIntervalMinutes), + accessRewriteUnfinished + ? now + : contentPass && !contentPass.complete + ? contentPass.checkpoint.resumeAt + ? new Date(contentPass.checkpoint.resumeAt) + : now + : calculateNextSyncTime(syncIntervalMinutes), completionNotice, - result.docsFailed === 0 && (!contentPass || !isContentPassIncomplete(contentPass)) + result.docsFailed === 0 && + !accessRewriteUnfinished && + (!contentPass || !isContentPassIncomplete(contentPass)) ), - /** Restored above under this same lock, or hidden by the admin pass before the ACLs it wrote. */ - accessRewritePending: false, + /** + * Restored before completion under this run's lease, or hidden by the admin pass before + * the ACLs it wrote; cleared only once that walk reached the end of the connector. + */ + ...(accessRewriteUnfinished ? {} : { accessRewritePending: false }), ...(contentPass?.complete ? { listingCheckpoint: null } : {}), ...(contentPass && !isContentPassIncomplete(contentPass) && result.docsFailed === 0 ? { lastSyncAt: new Date(contentPass.checkpoint.startedAt) } @@ -1087,6 +1084,9 @@ export async function executeSync( const mirrored = mirrorsSourceAcls(connector.accessMode) const sourceConfig = connector.sourceConfig as Record const syncStartedAt = new Date() + /** One budget for every page walker of the run, ending before the worker's own limit. */ + const runDeadlineAt = + syncStartedAt.getTime() + (CONNECTOR_SYNC_MAX_DURATION_SECONDS - 300) * 1000 const lease = createContentSyncLease(connectorId, syncLogId) await db.insert(knowledgeConnectorSyncLog).values({ id: syncLogId, @@ -1246,10 +1246,29 @@ export async function executeSync( * is safe to do last; hiding is not. */ if (connector.accessRewritePending) { - await rewriteConnectorAcls(connectorId, EMPTY_ACL, { + const hidden = await rewriteConnectorAcls(connectorId, EMPTY_ACL, { beforeBatch: lease.beatIfDue, lease, + deadlineAt: runDeadlineAt, }) + if (!hidden) { + /** Nothing is listed while documents are still readable; the next run resumes the walk. */ + const landed = await completeSuccessfulSync( + connectorId, + connector.knowledgeBaseId, + syncLogId, + effectiveConnectorSyncIntervalMinutes( + connector.accessMode, + connector.syncIntervalMinutes + ), + result, + null, + undefined, + null, + true + ) + return landed ? result : markSyncSuperseded(result) + } } /** * Started before the listing and awaited before the ACLs are written: a @@ -1309,7 +1328,7 @@ export async function executeSync( accessMode: connector.accessMode, }), fullSync: options.fullSync, - deadlineAt: syncStartedAt.getTime() + (CONNECTOR_SYNC_MAX_DURATION_SECONDS - 300) * 1000, + deadlineAt: runDeadlineAt, onPage: mirrored ? async (externalDocs, generationStartedAt) => { await directoryRefreshed @@ -1361,6 +1380,34 @@ export async function executeSync( lease, }) + /** + * Finishes a switch into workspace mode that outgrew its request budget + * or was interrupted: every document of the connector becomes readable by + * the workspace before the completion write clears the pending flag. The + * flag is the only source of such drift: every other writer of a + * workspace-mode document's ACL writes the workspace ACL, and both the + * mode switch and an ACL-resetting edit set the flag before anything can + * hide a document. One short lease-proving transaction per page that also + * re-checks the mode, before the completion transaction, so a large + * restore never holds the connector row across its projection fan-out. + */ + let accessRewriteUnfinished = false + if (accessMode === 'workspace' && connector.accessRewritePending) { + const restore = await restoreWorkspaceDocumentAcls( + connectorId, + leaseTransaction(connectorId, lease), + { beforePage: lease.beatIfDue, deadlineAt: runDeadlineAt } + ) + accessRewriteUnfinished = !restore.finished + if (restore.restored > 0) { + logger.warn('Restored workspace access on connector documents that had drifted', { + connectorId, + restoredAcls: restore.restored, + finished: restore.finished, + }) + } + } + const completionLanded = await completeSuccessfulSync( connectorId, connector.knowledgeBaseId, @@ -1369,7 +1416,8 @@ export async function executeSync( result, reconciliationHoldNotice, contentPass, - directoryNotice + directoryNotice, + accessRewriteUnfinished ) if (!completionLanded) { diff --git a/apps/sim/lib/knowledge/connectors/sync-limits.ts b/apps/sim/lib/knowledge/connectors/sync-limits.ts index 798258e78b9..f4134be13c7 100644 --- a/apps/sim/lib/knowledge/connectors/sync-limits.ts +++ b/apps/sim/lib/knowledge/connectors/sync-limits.ts @@ -155,3 +155,40 @@ export const SOURCE_PERMISSION_ERROR = /** Source downloads are retried by connector listing, never by parsing the retained file again. */ export const SOURCE_CONTENT_ERROR = 'Source content could not be refreshed. The connector will retry at its next scheduled sync.' + +/** + * Documents whose permission evidence is refreshed per statement. Documents are + * grouped by identical ACL first — files under one folder overwhelmingly share + * theirs — so a crawl of thousands usually resolves to a handful of statements. + * A refresh never assigns `acl`, so it fires no projection fan-out. + */ +export const ACL_WRITE_BATCH_SIZE = 500 + +/** + * Search projection rows one statement rewrites per projection table: a page of + * detachment releases, and the most chunk rows a page of ACL assignments may + * send through the projection trigger. A page always holds at least one + * document, so a document larger than this still makes progress alone. + */ +export const PROJECTION_ROW_BATCH_SIZE = 250 + +/** + * How long a connector-lease ACL page waits on any lock before it fails. The + * connector row is locked last, so the wait is on document rows, which a + * processing commit may hold for its whole embedding write. + */ +export const LEASE_PAGE_LOCK_TIMEOUT_MS = 15_000 + +/** The longest one statement of a connector-lease ACL page may run. */ +export const LEASE_PAGE_STATEMENT_TIMEOUT_MS = 30_000 + +/** + * Documents whose ACL actually changes, per statement. Assigning `acl` fires the + * document trigger that copies it onto every chunk's search projection rows, and + * each of those rows is re-inserted into the vector index, so one statement costs + * the chunks of every document in it rather than the documents. Kept small so a + * page of changed documents cannot outrun the statement timeout. Also the page + * of the transactions that remove observations and rematerialise the ACLs they + * decide together, which must commit as one and so cannot be split by rows. + */ +export const ACL_CHANGE_BATCH_SIZE = 25 diff --git a/apps/sim/lib/knowledge/connectors/sync-lock.ts b/apps/sim/lib/knowledge/connectors/sync-lock.ts index 9f5a4e4b397..8d2ddc72dce 100644 --- a/apps/sim/lib/knowledge/connectors/sync-lock.ts +++ b/apps/sim/lib/knowledge/connectors/sync-lock.ts @@ -1,7 +1,12 @@ import { db } from '@sim/db' import { knowledgeConnector } from '@sim/db/schema' -import { and, eq, isNull } from 'drizzle-orm' -import { SYNC_LOCK_HEARTBEAT_INTERVAL_MS } from '@/lib/knowledge/connectors/sync-limits' +import { and, eq, isNull, sql } from 'drizzle-orm' +import type { DbOrTx } from '@/lib/db/types' +import { + LEASE_PAGE_LOCK_TIMEOUT_MS, + LEASE_PAGE_STATEMENT_TIMEOUT_MS, + SYNC_LOCK_HEARTBEAT_INTERVAL_MS, +} from '@/lib/knowledge/connectors/sync-limits' /** * Raised when a run discovers mid-flight that it no longer holds its sync lock. @@ -235,6 +240,42 @@ export async function assertSyncLeaseHeldInTx( if (!held) throw new SyncLockLostException(connectorId) } +/** + * The bounds of a connector-lease ACL page. The `document` ACL trigger rewrites every filled + * search projection row of a document whose ACL is assigned, so a page that waits on a lock or + * runs long fails within the bounds and rolls back only itself. + */ +export async function boundLeaseTransaction(tx: Pick): Promise { + await tx.execute( + sql`SELECT set_config('lock_timeout', ${`${LEASE_PAGE_LOCK_TIMEOUT_MS}ms`}, true), set_config('statement_timeout', ${`${LEASE_PAGE_STATEMENT_TIMEOUT_MS}ms`}, true)` + ) +} + +/** Runs one bounded page of writes in a short transaction of its own. */ +export type LeaseTransaction = (write: (tx: DbOrTx) => Promise) => Promise + +/** + * One short, bounded transaction per call that proves `lease` as its last statement, so a run + * that lost its lease writes nothing further: the proof fails and the page rolls back. Proving it + * last keeps the connector row unlocked while the page waits on document rows, which a processing + * commit may hold for its whole write, and the share lock it then takes keeps the reclaim from + * landing until the page commits. Without a lease the page is only bounded: callers outside a + * sync run have no lease to prove. + */ +export function leaseTransaction( + connectorId: string, + lease?: SyncWriteLease, + executor: Pick = db +): LeaseTransaction { + return (write) => + executor.transaction(async (tx) => { + await boundLeaseTransaction(tx) + const written = await write(tx) + if (lease) await assertSyncLeaseHeldInTx(tx, connectorId, lease) + return written + }) +} + /** * The lease of the content sync engine, held through `syncLockToken`. The * heartbeat clock is seeded at lock acquisition, which opened `syncLockLeaseAt`. diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.postgres.test.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.postgres.test.ts index 223a7a5507f..b8c29fc5f00 100644 --- a/apps/sim/lib/knowledge/connectors/sync-persistence.postgres.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.postgres.test.ts @@ -17,6 +17,7 @@ vi.mock('@/connectors/registry.server', () => ({ CONNECTOR_REGISTRY: {} })) const { drizzle } = await import('drizzle-orm/postgres-js') const schema = await import('@sim/db/schema') const { persistDocumentAcls } = await import('@/lib/knowledge/connectors/sync-persistence') +const { leaseTransaction } = await import('@/lib/knowledge/connectors/sync-lock') const databaseUrl = process.env.KNOWLEDGE_ACL_TEST_DATABASE_URL @@ -43,8 +44,10 @@ describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => SELECT id, acl FROM embedding_search UNION ALL SELECT id, acl FROM embedding_keyword_tin ORDER BY id` - const persist = (acls: Map) => - persistDocumentAcls('admin', acls, drizzle(sql, { schema })) + /** Bounded page transactions on this schema's connection; these fixtures hold no lease. */ + const pages = () => leaseTransaction('admin', undefined, drizzle(sql, { schema })) + + const persist = (acls: Map) => persistDocumentAcls('admin', acls, pages()) beforeAll(async () => { const url = new URL(databaseUrl!) @@ -62,7 +65,7 @@ describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => connection: { search_path: schemaName }, }) await sql`CREATE TABLE document ( - id text PRIMARY KEY, external_id text, connector_id text, + id text PRIMARY KEY, external_id text, connector_id text, chunk_count integer NOT NULL DEFAULT 1, acl text[] NOT NULL DEFAULT '{ws}', acl_requirements jsonb NOT NULL DEFAULT '[]', acl_verified_at timestamp )` @@ -147,7 +150,7 @@ describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => await persistDocumentAcls( 'admin', new Map([['file-same', { acl: [ALICE], requirements: [['g:confluence:tenant:space']] }]]), - drizzle(sql, { schema }) + pages() ) const [stored] = await sql<{ requirements: string[][] }[]>` @@ -227,15 +230,10 @@ describe.runIf(Boolean(databaseUrl))('persistDocumentAcls in PostgreSQL', () => await sql`UPDATE document SET acl_verified_at = ${verifiedAt}::timestamptz AT TIME ZONE 'UTC' WHERE id = 'doc-same'` - const result = await persistDocumentAcls( - 'admin', - new Map([['file-same', []]]), - drizzle(sql, { schema }), - { - unresolvedExternalIds: new Set(['file-same']), - generationStartedAt, - } - ) + const result = await persistDocumentAcls('admin', new Map([['file-same', []]]), pages(), { + unresolvedExternalIds: new Set(['file-same']), + generationStartedAt, + }) expect(result).toEqual({ updated: preserved ? 0 : 1, rejected: 0 }) const [stored] = await sql<{ acl: string[]; verified: boolean }[]>` diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts index 86931ddbed5..5ba27d1ff6c 100644 --- a/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts @@ -49,6 +49,7 @@ vi.mock('@/connectors/registry.server', () => ({ import { MAX_ACL_TOKENS } from '@/lib/knowledge/access/tokens' import { getConnectorFailureDiagnostic } from '@/lib/knowledge/connectors/connector-error' +import { type LeaseTransaction, SyncLockLostException } from '@/lib/knowledge/connectors/sync-lock' import { addDocument, persistDocumentAcls, @@ -59,13 +60,42 @@ import { const CONNECTOR = 'connector-1' -/** Each `update(...).where(...)` chain ends in `returning()`; one row per changed document. */ -function queueUpdatedCounts(...counts: number[]) { - for (const count of counts) { +/** Runs each page straight on the mocked client, counting the transactions a writer opens. */ +const pages = vi.fn() +const direct: LeaseTransaction = (write) => { + pages() + return write(db) +} + +/** A lease that holds for `held` pages and is lost from then on. */ +function losingLease(held: number): LeaseTransaction { + let opened = 0 + return (write) => { + opened += 1 + if (opened > held) return Promise.reject(new SyncLockLostException(CONNECTOR)) + return write(db) + } +} + +/** + * Queues one ACL group's writes: the evidence refresh reports the external ids it matched, the + * same transaction then reads the documents whose ACL changes (only when some were not + * refreshed), and each change page reports the rows it wrote. + */ +function queueGroup(refreshed: string[], changed = 0, unrefreshed = changed > 0) { + dbChainMockFns.returning.mockResolvedValueOnce(refreshed.map((externalId) => ({ externalId }))) + if (!unrefreshed) return + const rows = Array.from({ length: changed }, (_unused, index) => ({ + id: `doc-${index}`, + chunkCount: 1, + })) + queueTableRows(schemaMock.document, rows) + /** The change page locks its documents and rereads their chunk counts before writing. */ + if (changed > 0) queueTableRows(schemaMock.document, rows) + if (changed > 0) dbChainMockFns.returning.mockResolvedValueOnce( - Array.from({ length: count }, (_unused, index) => ({ id: `doc-${index}` })) + Array.from({ length: changed }, (_unused, index) => ({ id: `doc-${index}` })) ) - } } describe('persistDocumentAcls', () => { @@ -81,7 +111,7 @@ describe('persistDocumentAcls', () => { * corpus every time somebody joined a group. */ it('refreshes only access fields, so no document is re-embedded', async () => { - queueUpdatedCounts(0, 1) + queueGroup([], 1) await persistDocumentAcls(CONNECTOR, new Map([['file-1', ['u:alice@corp.com']]])) @@ -102,7 +132,7 @@ describe('persistDocumentAcls', () => { * differs, so a document whose ACL did not change must only have its evidence refreshed. */ it('refreshes the evidence of an unchanged ACL without assigning it', async () => { - queueUpdatedCounts(1, 0) + queueGroup(['file-1']) await expect( persistDocumentAcls(CONNECTOR, new Map([['file-1', ['u:alice@corp.com']]])) @@ -117,7 +147,7 @@ describe('persistDocumentAcls', () => { }) it('reports how many documents received current permission evidence', async () => { - queueUpdatedCounts(2) + queueGroup(['file-1', 'file-2']) await expect( persistDocumentAcls( @@ -135,7 +165,8 @@ describe('persistDocumentAcls', () => { * keeps a crawl of thousands to a handful of statements. */ it('writes one refresh and one change statement per distinct ACL, not per document', async () => { - queueUpdatedCounts(0, 2, 0, 1) + queueGroup([], 2) + queueGroup([], 1) await persistDocumentAcls( CONNECTOR, @@ -166,7 +197,7 @@ describe('persistDocumentAcls', () => { }) it('groups ACLs that differ only in order or duplication', async () => { - queueUpdatedCounts(2) + queueGroup([], 2) await persistDocumentAcls( CONNECTOR, @@ -189,7 +220,7 @@ describe('persistDocumentAcls', () => { describe('an ACL we cannot store', () => { it('rejects workspace escape tokens even inside a source restriction', async () => { - queueUpdatedCounts(1) + queueGroup([], 1) const result = await persistDocumentAcls( CONNECTOR, new Map([['file-1', { acl: ['u:alice@corp.com'], requirements: [['ws']] }]]) @@ -203,7 +234,8 @@ describe('persistDocumentAcls', () => { }) it('retains an empty restriction and separately persists different clauses', async () => { - queueUpdatedCounts(0, 1, 0, 1) + queueGroup([], 1) + queueGroup([], 1) await persistDocumentAcls( CONNECTOR, new Map([ @@ -231,7 +263,7 @@ describe('persistDocumentAcls', () => { }) it('hides a document whose ACL carries a malformed token', async () => { - queueUpdatedCounts(1) + queueGroup([], 1) await expect( persistDocumentAcls(CONNECTOR, new Map([['file-1', ['u:NOT-FOLDED@corp.com']]])) @@ -244,7 +276,7 @@ describe('persistDocumentAcls', () => { }) it('hides a document whose ACL exceeds the ceiling', async () => { - queueUpdatedCounts(1) + queueGroup([], 1) const huge = Array.from({ length: MAX_ACL_TOKENS + 1 }, (_u, i) => `u:p${i}@corp.com`) await expect(persistDocumentAcls(CONNECTOR, new Map([['file-1', huge]]))).resolves.toEqual({ @@ -259,7 +291,7 @@ describe('persistDocumentAcls', () => { }) it('stores an ACL exactly at the ceiling', async () => { - queueUpdatedCounts(1) + queueGroup(['file-1']) const atLimit = Array.from({ length: MAX_ACL_TOKENS }, (_u, i) => `u:p${i}@corp.com`) await expect(persistDocumentAcls(CONNECTOR, new Map([['file-1', atLimit]]))).resolves.toEqual( @@ -268,7 +300,8 @@ describe('persistDocumentAcls', () => { }) it('still writes the documents whose ACLs are fine', async () => { - queueUpdatedCounts(1, 1) + queueGroup(['file-1']) + queueGroup(['file-2']) await expect( persistDocumentAcls( @@ -291,6 +324,101 @@ describe('persistDocumentAcls', () => { }) }) +describe('persistDocumentAcls paging', () => { + beforeEach(() => { + vi.clearAllMocks() + resetDbChainMock() + }) + + const changedGroup = (count: number) => + new Map( + Array.from({ length: count }, (_unused, index) => [`file-${index}`, ['u:bob@corp.com']]) + ) + /** The documents the window's read finds changed, each carrying `chunkCount` chunks. */ + const queueChanged = (count: number, chunkCount: number) => + queueTableRows( + schemaMock.document, + Array.from({ length: count }, (_unused, index) => ({ id: `doc-${index}`, chunkCount })) + ) + const assignedPages = () => + dbChainMockFns.set.mock.calls + .map(([values], index) => ({ + values, + order: dbChainMockFns.set.mock.invocationCallOrder[index], + })) + .filter(({ values }) => 'acl' in values) + .map(({ order }) => { + const whereIndex = dbChainMockFns.where.mock.invocationCallOrder.findIndex( + (whereOrder) => whereOrder > order + ) + const ids = flattenMockConditions(dbChainMockFns.where.mock.calls[whereIndex]?.[0]).find( + (node) => node.type === 'inArray' && node.column === schemaMock.document.id + )?.values as string[] + return ids.length + }) + + /** One transaction per page: a lease lock held across pages outlasted the statement timeout. */ + it('writes every page in a transaction of its own, bounded by projection rows', async () => { + queueChanged(60, 10) + + await persistDocumentAcls(CONNECTOR, changedGroup(60), direct) + + expect(assignedPages()).toEqual([25, 25, 10]) + expect(pages).toHaveBeenCalledTimes(dbChainMockFns.set.mock.calls.length) + }) + + it('packs small documents into one page and gives a document above the cap a page alone', async () => { + queueTableRows(schemaMock.document, [ + { id: 'small-1', chunkCount: 1 }, + { id: 'huge', chunkCount: 1_000 }, + { id: 'small-2', chunkCount: 1 }, + { id: 'small-3', chunkCount: 0 }, + ]) + + await persistDocumentAcls(CONNECTOR, changedGroup(4), direct) + + expect(assignedPages()).toEqual([1, 1, 2]) + }) + + /** An unchanged crawl costs the refresh alone: no change read, no change transaction. */ + it('writes a window whose every ACL is unchanged in one transaction', async () => { + dbChainMockFns.returning.mockResolvedValueOnce( + Array.from({ length: 500 }, (_unused, index) => ({ externalId: `file-${index}` })) + ) + + await expect(persistDocumentAcls(CONNECTOR, changedGroup(500), direct)).resolves.toEqual({ + updated: 500, + rejected: 0, + }) + + expect(pages).toHaveBeenCalledOnce() + expect(dbChainMockFns.select).not.toHaveBeenCalled() + }) + + it('writes nothing further once the lease is lost between pages', async () => { + queueChanged(60, 10) + + await expect( + persistDocumentAcls(CONNECTOR, changedGroup(60), losingLease(2)) + ).rejects.toBeInstanceOf(SyncLockLostException) + + /** The evidence refresh and the first change page landed; the rest never ran. */ + expect(dbChainMockFns.set).toHaveBeenCalledTimes(2) + }) + + it('bounds each page in a transaction of its own by default', async () => { + queueChanged(30, 10) + + await persistDocumentAcls(CONNECTOR, changedGroup(30)) + + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(dbChainMockFns.set.mock.calls.length) + const bounds = dbChainMockFns.execute.mock.calls.filter((call: unknown[]) => + JSON.stringify(call).includes('lock_timeout') + ) + expect(bounds).toHaveLength(dbChainMockFns.transaction.mock.calls.length) + }) +}) + describe('revokeDocumentAcls', () => { beforeEach(() => { vi.clearAllMocks() @@ -307,20 +435,33 @@ describe('revokeDocumentAcls', () => { const grants = (node: Record) => sqlText(node)?.startsWith('cardinality(') && sqlText(node)?.endsWith(') > 0') - /** Every `where` condition, paired with the `set` of the same statement. */ + /** Every `set`, paired with the `where` of the same statement. */ function statements() { - return dbChainMockFns.set.mock.calls.map(([values], index) => ({ - values, - conditions: flattenMockConditions(dbChainMockFns.where.mock.calls[index]?.[0]), - })) + return dbChainMockFns.set.mock.calls.map(([values], index) => { + const order = dbChainMockFns.set.mock.invocationCallOrder[index] + const whereIndex = dbChainMockFns.where.mock.invocationCallOrder.findIndex( + (whereOrder) => whereOrder > order + ) + return { + values, + conditions: flattenMockConditions(dbChainMockFns.where.mock.calls[whereIndex]?.[0]), + } + }) } + /** The documents the window's read finds still granting someone. */ + const queueGranting = (count: number, chunkCount = 1) => + queueTableRows( + schemaMock.document, + Array.from({ length: count }, (_unused, index) => ({ id: `doc-${index}`, chunkCount })) + ) /** * Assigning `acl` fires the projection fan-out whether or not the value changes, so a * document that already grants nobody must never be in an `acl` assignment. */ it('assigns acl only to documents that still grant someone', async () => { - await revokeDocumentAcls(db, ['a', 'b'], scope) + queueGranting(1) + await revokeDocumentAcls(direct, ['a', 'b'], scope) const writes = statements().filter(({ values }) => 'acl' in values) expect(writes).toHaveLength(1) @@ -328,8 +469,16 @@ describe('revokeDocumentAcls', () => { expect(writes[0].conditions.some(grants)).toBe(true) }) + it('assigns nothing when no document still grants someone', async () => { + await revokeDocumentAcls(direct, ['a', 'b'], scope) + + expect(statements().filter(({ values }) => 'acl' in values)).toHaveLength(0) + expect(pages).toHaveBeenCalledOnce() + }) + it('clears leftover evidence on an already-empty ACL without assigning acl', async () => { - await revokeDocumentAcls(db, ['a', 'b'], scope) + queueGranting(1) + await revokeDocumentAcls(direct, ['a', 'b'], scope) const clears = statements().filter(({ values }) => !('acl' in values)) expect(clears).toHaveLength(1) @@ -342,20 +491,41 @@ describe('revokeDocumentAcls', () => { }) /** Each document in an `acl` assignment costs a rewrite of every one of its chunks' projection rows. */ - it('assigns acl in batches of 25 and clears evidence in batches of 500', async () => { + it('assigns acl in pages bounded by projection rows and clears evidence per window', async () => { + const ids = Array.from({ length: 60 }, (_unused, index) => `doc-${index}`) + queueGranting(60, 10) + + await revokeDocumentAcls(direct, ids, scope) + + const pageSizes = statements() + .filter(({ values }) => 'acl' in values) + .map(({ conditions }) => { + const pageIds = conditions.filter((node) => node.type === 'inArray').at(-1) + return (pageIds?.values as string[]).length + }) + expect(pageSizes).toEqual([25, 25, 10]) + expect(statements().filter(({ values }) => !('acl' in values))).toHaveLength(1) + }) + + it('runs every write in a transaction of its own', async () => { const ids = Array.from({ length: 60 }, (_unused, index) => `doc-${index}`) + queueGranting(60, 10) - await revokeDocumentAcls(db, ids, scope) - - const batchSizes = (writesAcl: boolean) => - statements() - .filter(({ values }) => 'acl' in values === writesAcl) - .map(({ conditions }) => { - const inArray = conditions.find((node) => node.type === 'inArray') - return (inArray?.values as string[]).length - }) - expect(batchSizes(true)).toEqual([25, 25, 10]) - expect(batchSizes(false)).toEqual([60]) + await revokeDocumentAcls(direct, ids, scope) + + expect(pages).toHaveBeenCalledTimes(dbChainMockFns.set.mock.calls.length) + expect(pages).toHaveBeenCalledTimes(4) + }) + + it('stops writing at the first page whose lease is gone', async () => { + const ids = Array.from({ length: 60 }, (_unused, index) => `doc-${index}`) + queueGranting(60, 10) + + await expect(revokeDocumentAcls(losingLease(2), ids, scope)).rejects.toBeInstanceOf( + SyncLockLostException + ) + + expect(dbChainMockFns.set).toHaveBeenCalledTimes(2) }) }) diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.ts index ec3799f1dd5..0cbe765e1db 100644 --- a/apps/sim/lib/knowledge/connectors/sync-persistence.ts +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.ts @@ -6,7 +6,6 @@ import { generateId } from '@sim/utils/id' import { truncateAtCodePoint } from '@sim/utils/string' import { and, eq, exists, inArray, isNull, lt, not, or, type SQL, sql } from 'drizzle-orm' import { getInternalApiBaseUrl } from '@/lib/core/utils/urls' -import type { DbOrTx } from '@/lib/db/types' import { textArrayLiteral } from '@/lib/knowledge/access/predicate' import { EMPTY_ACL, @@ -16,9 +15,19 @@ import { import type { MirroredDocumentAcl } from '@/lib/knowledge/access/types' import { aclIsDerived, type ConnectorAccessMode } from '@/lib/knowledge/connectors/access-modes' import type { ConnectorFailureDiagnostic } from '@/lib/knowledge/connectors/connector-error' +import { + pagesByProjectionRows, + rewriteConnectorDocumentAcls, + writeProjectionPages, +} from '@/lib/knowledge/connectors/member-observations' import { resolveSourceModifiedAt } from '@/lib/knowledge/connectors/source-modified-at' -import { SOURCE_CONTENT_ERROR } from '@/lib/knowledge/connectors/sync-limits' -import { assertSyncLeaseHeldInTx, type SyncWriteLease } from '@/lib/knowledge/connectors/sync-lock' +import { ACL_WRITE_BATCH_SIZE, SOURCE_CONTENT_ERROR } from '@/lib/knowledge/connectors/sync-limits' +import { + assertSyncLeaseHeldInTx, + type LeaseTransaction, + leaseTransaction, + type SyncWriteLease, +} from '@/lib/knowledge/connectors/sync-lock' import { MAX_DOCUMENT_INDEXED_TEXT_LENGTH } from '@/lib/knowledge/constants' import type { DocumentData } from '@/lib/knowledge/documents/service' import { enqueueKnowledgeStorageCleanup } from '@/lib/knowledge/documents/storage-cleanup' @@ -47,54 +56,38 @@ function updatedDocumentAcl(access: ConnectorAccessMode) { * The workspace-mode invariant, applied after every successful content sync: * a document a workspace-mode connector owns is readable by the workspace, * whatever a mode switch or an interrupted rewrite left behind. Idempotent and - * a no-op on a healthy connector. + * a no-op on a healthy connector. Paged in short transactions that each prove + * the connector is still in workspace mode, so it never holds the connector row + * across the projection fan-out of a large restore. Stops between pages once + * `deadlineAt` passes and reports it unfinished; a later walk resumes by + * rewriting whatever is still off the workspace ACL. */ export async function restoreWorkspaceDocumentAcls( - executor: DbOrTx, - connectorId: string -): Promise { - const workspaceAcl = textArrayLiteral(WORKSPACE_ACL) - const restored = await executor - .update(document) - .set({ acl: [...WORKSPACE_ACL], aclRequirements: [], aclVerifiedAt: null }) - .where( - and( - eq(document.connectorId, connectorId), - sql`(${document.acl} <> ${workspaceAcl} OR ${document.aclRequirements} <> '[]'::jsonb OR ${document.aclVerifiedAt} IS NOT NULL)`, - exists( - executor - .select({ one: sql`1` }) - .from(knowledgeConnector) - .where( - and( - eq(knowledgeConnector.id, connectorId), - eq(knowledgeConnector.accessMode, 'workspace') - ) - ) + connectorId: string, + transaction: LeaseTransaction, + options: { beforePage?: () => Promise; deadlineAt?: number } = {} +): Promise<{ restored: number; finished: boolean }> { + const { rewritten, finished } = await rewriteConnectorDocumentAcls({ + connectorId, + target: WORKSPACE_ACL, + transaction, + beforePage: options.beforePage, + deadlineAt: options.deadlineAt, + guard: exists( + db + .select({ one: sql`1` }) + .from(knowledgeConnector) + .where( + and( + eq(knowledgeConnector.id, connectorId), + eq(knowledgeConnector.accessMode, 'workspace') + ) ) - ) - ) - .returning({ id: document.id }) - return restored.length + ), + }) + return { restored: rewritten, finished } } -/** - * Documents whose permission evidence is refreshed per statement. Documents are - * grouped by identical ACL first — files under one folder overwhelmingly share - * theirs — so a crawl of thousands usually resolves to a handful of statements. - * A refresh never assigns `acl`, so it fires no projection fan-out. - */ -const ACL_WRITE_BATCH_SIZE = 500 - -/** - * Documents whose ACL actually changes, per statement. Assigning `acl` fires the - * document trigger that copies it onto every chunk's search projection rows, and - * each of those rows is re-inserted into the vector index, so one statement costs - * the chunks of every document in it rather than the documents. Kept small so a - * page of changed documents cannot outrun the statement timeout. - */ -const ACL_CHANGE_BATCH_SIZE = 25 - export interface DocumentAclWriteResult { /** Documents whose ACL or authoritative evidence timestamp was refreshed. */ updated: number @@ -111,11 +104,14 @@ export interface DocumentAclWriteResult { * survive failed verification. * An unresolved duplicate may retain evidence verified during this durable crawl, * without refreshing its timestamp; explicit empty ACLs always revoke access. + * Every batch is its own `transaction`, which proves the run's lease when it has one: a + * connector row lock held across every batch outlasted the statement timeout and rolled back + * the batches that had landed. Batches are idempotent, so a retry rewrites only what is stale. */ export async function persistDocumentAcls( connectorId: string, acls: ReadonlyMap, - executor: DbOrTx = db, + transaction: LeaseTransaction = leaseTransaction(connectorId), evidence?: { unresolvedExternalIds: ReadonlySet generationStartedAt: Date @@ -174,22 +170,46 @@ export async function persistDocumentAcls( unchanged ? stored : not(stored), evidenceGuard ) - /** Refreshed first, so a row the change write below has just rewritten is not counted twice. */ - for (const batch of chunkArray(externalIds, ACL_WRITE_BATCH_SIZE)) { - const rows = await executor - .update(document) - .set({ aclVerifiedAt }) - .where(target(batch, true)) - .returning({ id: document.id }) - updated += rows.length - } - for (const batch of chunkArray(externalIds, ACL_CHANGE_BATCH_SIZE)) { - const rows = await executor - .update(document) - .set({ acl, aclRequirements: requirements, aclVerifiedAt }) - .where(target(batch, false)) - .returning({ id: document.id }) - updated += rows.length + /** + * Each window's evidence refresh and its read of the documents whose ACL changes share one + * transaction; the refresh reports the documents it matched, so an unchanged crawl costs one + * transaction per window and the change pages below see only what actually changes. Refreshed + * first, so a row a change page then rewrites is never counted twice. + */ + for (const window of chunkArray(externalIds, ACL_WRITE_BATCH_SIZE)) { + const { refreshed, changed } = await transaction(async (tx) => { + const refreshed = await tx + .update(document) + .set({ aclVerifiedAt }) + .where(target(window, true)) + .returning({ externalId: document.externalId }) + const matched = new Set(refreshed.map((row) => row.externalId)) + const remaining = window.filter((externalId) => !matched.has(externalId)) + const changed = + remaining.length === 0 + ? [] + : await tx + .select({ id: document.id, chunkCount: document.chunkCount }) + .from(document) + .where(target(remaining, false)) + return { refreshed: refreshed.length, changed } + }) + updated += refreshed + for (const page of pagesByProjectionRows(changed)) { + const { written } = await writeProjectionPages( + page, + transaction, + async (tx, locked) => + ( + await tx + .update(document) + .set({ acl, aclRequirements: requirements, aclVerifiedAt }) + .where(and(inArray(document.id, locked), target(window, false))) + .returning({ id: document.id }) + ).length + ) + updated += written + } } } @@ -199,34 +219,44 @@ export async function persistDocumentAcls( /** * Revokes every grant on the documents `target` selects from `ids`, leaving each readable by * nobody with its permission evidence cleared. Only a document that still grants someone has - * `acl` assigned, {@link ACL_CHANGE_BATCH_SIZE} at a time: the projection trigger fires on every + * `acl` assigned, in pages bounded by their chunks' projection rows: the projection trigger fires on every * assignment of `acl`, changed or not, and each document costs a rewrite of its chunks' * projection rows. A document already readable by nobody only has leftover evidence cleared, - * which fires no fan-out. + * which fires no fan-out. Each batch is its own `transaction`, so a lease lost between batches + * stops the rest and every committed batch stays revoked. */ export async function revokeDocumentAcls( - executor: DbOrTx, + transaction: LeaseTransaction, ids: string[], target: (batch: string[]) => SQL | undefined ): Promise { const grants = sql`cardinality(${document.acl}) > 0` - for (const batch of chunkArray(ids, ACL_WRITE_BATCH_SIZE)) { - await executor - .update(document) - .set({ aclRequirements: [], aclVerifiedAt: null }) - .where( - and( - target(batch), - not(grants), - sql`(${document.aclRequirements} <> '[]'::jsonb OR ${document.aclVerifiedAt} IS NOT NULL)` + for (const window of chunkArray(ids, ACL_WRITE_BATCH_SIZE)) { + const granting = await transaction(async (tx) => { + await tx + .update(document) + .set({ aclRequirements: [], aclVerifiedAt: null }) + .where( + and( + target(window), + not(grants), + sql`(${document.aclRequirements} <> '[]'::jsonb OR ${document.aclVerifiedAt} IS NOT NULL)` + ) ) - ) - } - for (const batch of chunkArray(ids, ACL_CHANGE_BATCH_SIZE)) { - await executor - .update(document) - .set({ acl: [], aclRequirements: [], aclVerifiedAt: null }) - .where(and(target(batch), grants)) + return tx + .select({ id: document.id, chunkCount: document.chunkCount }) + .from(document) + .where(and(target(window), grants)) + }) + for (const page of pagesByProjectionRows(granting)) { + await writeProjectionPages(page, transaction, async (tx, locked) => { + await tx + .update(document) + .set({ acl: [], aclRequirements: [], aclVerifiedAt: null }) + .where(and(target(window), inArray(document.id, locked), grants)) + return 0 + }) + } } }