From bb4623a8b64933a87bf9eeec6dd16847869b34d5 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 22 Sep 2026 23:21:52 -0700 Subject: [PATCH 1/8] fix(knowledge): bound every connector-lease ACL transaction to one short page --- .github/workflows/test-build.yml | 1 + .../connector-lease-pages.integration.ts | 653 ++++++++++++++++++ .../connectors/member-observations.test.ts | 172 ++++- .../connectors/member-observations.ts | 377 ++++++---- .../connectors/member-sync-engine.ts | 130 ++-- .../sync-content-pass.postgres.test.ts | 7 +- .../connectors/sync-content-pass.test.ts | 21 +- .../knowledge/connectors/sync-content-pass.ts | 45 +- .../knowledge/connectors/sync-engine.test.ts | 61 +- .../lib/knowledge/connectors/sync-engine.ts | 56 +- .../lib/knowledge/connectors/sync-limits.ts | 19 + .../sim/lib/knowledge/connectors/sync-lock.ts | 37 +- .../sync-persistence.postgres.test.ts | 22 +- .../connectors/sync-persistence.test.ts | 83 ++- .../knowledge/connectors/sync-persistence.ts | 144 ++-- 15 files changed, 1513 insertions(+), 315 deletions(-) create mode 100644 apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts diff --git a/.github/workflows/test-build.yml b/.github/workflows/test-build.yml index 4736123fd93..8c49548f7c4 100644 --- a/.github/workflows/test-build.yml +++ b/.github/workflows/test-build.yml @@ -276,6 +276,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/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..4cef8073038 --- /dev/null +++ b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts @@ -0,0 +1,653 @@ +/** + * 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, + knowledgeConnector, + knowledgeConnectorMember, + organization, + resourcePolicy, + user, + workspace, +} from '@sim/db/schema' +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 { + materializeDocumentAcls, + recordMemberObservations, + rewriteConnectorAcls, +} 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 { + createMemberSyncLease, + type LeaseTransaction, + leaseTransaction, + SyncLockLostException, + stillHoldsMemberSyncLock, + stillHoldsSyncLock, +} from '@/lib/knowledge/connectors/sync-lock' +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()`) + }) + + 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.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`) + 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', + 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: '5s', 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('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.toBe(DOCUMENTS) + + 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.toBe(0) + }) + + it('restores drift before a sync completes, outside the completion transaction', async () => { + await db + .update(knowledgeConnector) + .set({ status: 'active', syncLockToken: null }) + .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, + }) + .from(knowledgeConnector) + .where(eq(knowledgeConnector.id, ids.connectorId)) + expect(connector).toEqual({ status: 'active', syncLockToken: null }) + }) + + 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.toBe(0) + 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('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) + }) + }) +}) diff --git a/apps/sim/lib/knowledge/connectors/member-observations.test.ts b/apps/sim/lib/knowledge/connectors/member-observations.test.ts index 6abdc39e2e8..ce91f892068 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.test.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.test.ts @@ -20,6 +20,7 @@ import { db } from '@sim/db' import { inArray } from 'drizzle-orm' import { applyMemberDocumentLifecycle, + materializeDocumentAcls, removeUnseenMemberObservations, renewMemberObservationsInScopes, rewriteConnectorAcls, @@ -46,18 +47,19 @@ describe('removeUnseenMemberObservations', () => { }) it('rematerializes bounded batches before reading more absent observations', async () => { dbChainMockFns.returning - .mockResolvedValueOnce(Array.from({ length: 500 }, (_, i) => ({ documentId: `d-${i}` }))) + .mockResolvedValueOnce(Array.from({ length: 25 }, (_, i) => ({ documentId: `d-${i}` }))) .mockResolvedValueOnce([{ documentId: 'last' }]) .mockResolvedValueOnce([]) 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] ) @@ -178,6 +180,35 @@ describe('sweepStaleMemberObservations', () => { 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' }]) + 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(([query]) => + JSON.stringify(query).includes('lock_timeout') + ) + expect(bounds).toHaveLength(2) + }) + /** * A run that claimed the member between the selection and the lock moved * `lastStartedAt` forward, so the re-check under `FOR UPDATE` finds nothing @@ -212,19 +243,67 @@ describe('sweepStaleMemberObservations', () => { }) }) +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('rewriteConnectorAcls', () => { beforeEach(() => { vi.clearAllMocks() resetDbChainMock() }) - it('proves the lease inside each batch transaction before rewriting', async () => { + const stale = (id: string, aclDiffers = true, evidencePresent = false) => ({ + id, + externalId: `ext-${id}`, + 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.where.mock.calls + .map(([condition]) => + flattenMockConditions(condition).find( + (node) => node.type === 'inArray' && node.column === schemaMock.document.id + ) + ) + .filter((node) => node !== undefined) + .map((node) => (node!.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') @@ -232,14 +311,87 @@ describe('rewriteConnectorAcls', () => { }) it('stops without writing once the lease is gone', async () => { + queueTableRows(schemaMock.document, [stale('d-1')]) queueTableRows(schemaMock.knowledgeConnector, []) + await expect(rewriteConnectorAcls('c-1', [], { lease: held })).rejects.toBeInstanceOf( + SyncLockLostException + ) + + expect(dbChainMockFns.update).not.toHaveBeenCalled() + }) + + /** 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(([query]) => + JSON.stringify(query).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) + }) + + 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() }) + + it('walks again after a pass that wrote, so a row changed behind the walk is reached', async () => { + queueTableRows(schemaMock.document, [stale('d-1')]) + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) + dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'd-1' }]) + queueTableRows(schemaMock.document, [stale('d-0')]) + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) + dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'd-0' }]) + queueTableRows(schemaMock.document, []) + + await expect(rewriteConnectorAcls('c-1', [], { lease: held })).resolves.toBe(true) + + expect(assignedPages()).toHaveLength(2) + }) }) describe('applyMemberDocumentLifecycle', () => { diff --git a/apps/sim/lib/knowledge/connectors/member-observations.ts b/apps/sim/lib/knowledge/connectors/member-observations.ts index 9f77dc7d4e6..0d8e0a8ef0c 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.ts @@ -5,6 +5,7 @@ import { knowledgeConnectorMember, knowledgeDocumentObservation, } from '@sim/db/schema' +import { chunkArray } from '@sim/utils/helpers' import { and, asc, @@ -16,21 +17,27 @@ 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, } from '@/lib/knowledge/connectors/sync-limits' import { - assertSyncLeaseHeldInTx, + boundLeaseTransaction, connectorIsLive, + type LeaseTransaction, + leaseTransaction, MEMBER_LOCKABLE_CONNECTOR_STATUSES, SyncLockLostException, type SyncRunLease, @@ -42,8 +49,11 @@ 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 +/** + * 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. */ @@ -219,7 +229,9 @@ 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 of + * {@link ACL_CHANGE_BATCH_SIZE} per call, because `onRemoved` rematerialises + * the page's ACLs in the caller's lease transaction. */ export async function removeUnseenMemberObservations( executor: DbOrTx, @@ -235,13 +247,13 @@ export async function removeUnseenMemberObservations( .select({ documentId: knowledgeDocumentObservation.documentId }) .from(knowledgeDocumentObservation) .where(unseen) - .limit(OBSERVATION_BATCH_SIZE) + .limit(ACL_CHANGE_BATCH_SIZE) const removed = await executor .delete(knowledgeDocumentObservation) .where(and(unseen, inArray(knowledgeDocumentObservation.documentId, candidates))) .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: removed.length < ACL_CHANGE_BATCH_SIZE } } /** @@ -319,19 +331,125 @@ 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. + * Rewrites the ACL of every document of the connector to `target`, clearing its permission + * evidence, one short transaction per page. The documents are walked 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, a document whose `acl` differs is + * assigned it {@link ACL_CHANGE_BATCH_SIZE} at a time, each batch its own transaction, because each + * assignment rewrites the document's search projection rows; one whose `acl` already matches only + * has its 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. A pass that wrote + * anything is followed by another, so a document a concurrent writer changed behind the walk is + * still reached. Returns the documents written and whether it finished before `deadlineAt`. */ -/** Documents rewritten per statement while a mode switch rewrites a connector's ACLs. */ -const ACCESS_REWRITE_BATCH_SIZE = 1000 +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 + for (;;) { + let passWrote = false + let after: ConnectorDocumentCursor | undefined + for (;;) { + if (expired()) return { rewritten, finished: false } + await input.beforePage?.() + const window = await db + .select({ + id: document.id, + externalId: document.externalId, + 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 + if (rows.length > 0) passWrote = true + } + const changed = window.filter((row) => row.aclDiffers).map((row) => row.id) + for (const batch of chunkArray(changed, ACL_CHANGE_BATCH_SIZE)) { + if (expired()) return { rewritten, finished: false } + await input.beforePage?.() + const rows = await transaction((tx) => + tx + .update(document) + .set({ acl: [...target], aclRequirements: [], aclVerifiedAt: null }) + .where( + and( + eq(document.connectorId, connectorId), + inArray(document.id, batch), + aclDiffers, + guard + ) + ) + .returning({ id: document.id }) + ) + rewritten += rows.length + if (rows.length > 0) passWrote = true + } + const last = window.at(-1) + if (window.length < ACL_WRITE_BATCH_SIZE || !last?.externalId) break + after = { externalId: last.externalId, id: last.id } + } + if (!passWrote) return { rewritten, finished: true } + } +} /** - * 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. + * 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 +458,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 +490,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 }) @@ -501,10 +606,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 +681,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 +700,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 +767,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 +787,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 +870,93 @@ 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 +} | null> { + return db.transaction(async (tx) => { + await boundLeaseTransaction(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(ACL_CHANGE_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, + } + }) +} + /** * Removes the observations of members whose crawls have stopped, so the * documents only they observed go dark instead of staying readable forever. @@ -783,7 +973,7 @@ 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 + * Each page of a member's sweep is 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 @@ -853,72 +1043,17 @@ export async function sweepStaleMemberObservations(now: Date): Promise { - 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, - } - }) - if (!swept) continue - result.members += 1 - result.observationsRemoved += swept.observationsRemoved - result.documentsRematerialized += swept.documentsRematerialized - result.docsTombstoned += swept.docsTombstoned + let sweptAny = false + for (let page = 0; page < STALE_MEMBER_PAGES_PER_TICK; page++) { + 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.observationsRemoved < ACL_CHANGE_BATCH_SIZE) break + } + if (sweptAny) result.members += 1 } return result } diff --git a/apps/sim/lib/knowledge/connectors/member-sync-engine.ts b/apps/sim/lib/knowledge/connectors/member-sync-engine.ts index 53c9de10a09..9d27962e6cc 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, 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' @@ -70,6 +71,7 @@ import { getConnectorSyncDeferral, } from '@/lib/knowledge/connectors/sync-deferral' import { + ACL_CHANGE_BATCH_SIZE, CONNECTOR_AUTO_DISABLED_ERROR, CONNECTOR_FAILURE_BACKOFF_CAP_MINUTES, connectorFailureBackoffMinutes, @@ -86,6 +88,7 @@ import { } from '@/lib/knowledge/connectors/sync-limits' import { assertSyncLeaseHeldInTx, + boundLeaseTransaction, createMemberSyncLease, holdsMemberSyncLockToken, MEMBER_LOCKABLE_CONNECTOR_STATUSES, @@ -435,13 +438,16 @@ function createMemberTokenCache(input: { * 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 * 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. Bounded + * like every connector-lease transaction, and callers keep the ACLs one call + * writes to a page, so the connector row is never held across unbounded fan-out. */ async function withMemberLease( run: Pick, fn: (tx: DbOrTx) => Promise ): Promise { return db.transaction(async (tx) => { + await boundLeaseTransaction(tx) const [held] = await tx .select({ id: knowledgeConnector.id }) .from(knowledgeConnector) @@ -583,7 +589,7 @@ export async function resumeMembershipRewrites( ) ) .orderBy(asc(knowledgeDocumentObservation.documentId)) - .limit(500) + .limit(ACL_CHANGE_BATCH_SIZE) await materializeDocumentAcls( run.connectorId, documents.map((row) => row.documentId), @@ -1784,12 +1790,20 @@ 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, so 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. Readers decide from the ACL alone, so 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( @@ -1798,11 +1812,20 @@ async function disableMemberSync(run: MemberSyncRun, reason: string): Promise + tx .update(knowledgeConnector) .set({ memberSyncStatus: 'disabled', @@ -1810,12 +1833,27 @@ 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') } /** @@ -1959,7 +1997,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', @@ -2092,34 +2136,36 @@ 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) { + /** + * One lease transaction per page of documents, each observing and + * materialising its page together; the failure checkpoint lands with + * the first, since forcing a later relist is the conservative side. + */ + const pages = + documentIds.length > 0 ? chunkArray(documentIds, ACL_CHANGE_BATCH_SIZE) : [[]] + for (const [index, page] of pages.entries()) { + await withMemberLease(run, async (tx) => { + result.observationsAdded += await recordMemberObservations( + tx, + member.id, + page, + checkpoint.generationId + ) + await materializeDocumentAcls(connectorId, page, tx) + 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))) } - } - }) + }) + } result.docsListed += attempted.length } if (!serviceContent) { @@ -2292,7 +2338,10 @@ export async function executeMemberSync( await loadDocumentIdsByExternalId(connectorId, relevantIds), connector.syncIntervalMinutes ) - await withMemberLease(run, (tx) => materializeDocumentAcls(connectorId, affected, tx)) + for (const page of chunkArray([...affected], ACL_CHANGE_BATCH_SIZE)) { + await run.lease.beatIfDue() + await withMemberLease(run, (tx) => materializeDocumentAcls(connectorId, page, tx)) + } } /** A service-owned corpus outlives its last observer; only the content pass removes it. */ @@ -2378,7 +2427,8 @@ export async function executeMemberSync( } if (error instanceof MemberBindingGoneError) { try { - await disableMemberSync(run, error.message) + if (!(await disableMemberSync(run, error.message))) + return await finishDisableLater(run, connector.syncIntervalMinutes) } catch (disableError) { if (!(disableError instanceof SyncLockLostException)) throw disableError logger.warn('Member sync abandoned — lock was reclaimed before it could be disabled', { 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..e92a8de9852 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 @@ -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..d0ee2f9f0d6 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,10 @@ describe('completed listing removal counts', () => { ]) queueTableRows(schemaMock.document, options.revoked ?? []) if (options.revoked?.length) { - queueTableRows(schemaMock.knowledgeConnector, [{ id: 'connector' }]) + /** Every revocation batch is its own lease-proving transaction: one evidence clear, then acl batches. */ + 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) { @@ -322,6 +325,9 @@ describe('completed listing removal counts', () => { ).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).toHaveBeenCalledTimes(3) }) it('does not report a full-sync removal when the guarded delete removed no live rows', async () => { @@ -828,6 +834,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..b3785edaf4d 100644 --- a/apps/sim/lib/knowledge/connectors/sync-content-pass.ts +++ b/apps/sim/lib/knowledge/connectors/sync-content-pass.ts @@ -26,7 +26,11 @@ import { SOURCE_CONTENT_ERROR, SOURCE_PERMISSION_ERROR, } from '@/lib/knowledge/connectors/sync-limits' -import { assertSyncLeaseHeldInTx, type SyncRunLease } from '@/lib/knowledge/connectors/sync-lock' +import { + type LeaseTransaction, + leaseTransaction, + type SyncRunLease, +} from '@/lib/knowledge/connectors/sync-lock' import { type KnowledgeBaseOwner, persistSourceDocumentFailures, @@ -87,11 +91,8 @@ 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) => - db.transaction(async (tx) => { - await assertSyncLeaseHeldInTx(tx, input.connectorId, input.lease) - return fn(tx) - }) + /** Every lease-fenced write of the pass is one short, bounded transaction. */ + const withLease = 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 +181,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( + withLease, + changed.map((item) => item.externalId), + (batch) => + and( + eq(document.connectorId, input.connectorId), + inArray(document.externalId, batch), + isNull(document.archivedAt) + ) ) } const state = createSyncRunState(input.result) @@ -349,7 +348,7 @@ export async function runConnectorContentPass(input: ContentPassInput) { async function reconcileCompletedListing( input: ContentPassInput, checkpoint: ListingCheckpoint, - withLease: (fn: (tx: DbOrTx) => Promise) => Promise + withLease: LeaseTransaction ): Promise<{ finished: boolean; notice: string | null }> { if (checkpoint.unsafe || (checkpoint.listingFailures?.count ?? 0) > 0) return { @@ -426,12 +425,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( + withLease, + 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 6b7e35e018f..09d8f4bc835 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts @@ -2196,7 +2196,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 = @@ -2266,7 +2265,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' }]) @@ -2299,7 +2297,6 @@ describe('completeSuccessfulSync', () => { ) ) dbChainMockFns.returning - .mockResolvedValueOnce([]) .mockResolvedValueOnce([{ id: 'log-1' }]) .mockResolvedValueOnce([{ id: 'c-1' }]) @@ -2332,8 +2329,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' }]) @@ -2373,7 +2368,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' }]) @@ -2420,7 +2414,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.' @@ -2488,7 +2481,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' }]) @@ -3224,6 +3216,59 @@ describe('executeSync heartbeats during the listing phase', () => { } ) + it('proves the lease 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' }]) + const pass = vi + .spyOn(contentPass, 'runConnectorContentPass') + .mockImplementation(async (input) => { + await input.onPage?.( + [ + { + externalId: 'page-1', + title: 'Page', + content: 'Body', + 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) + expect(between(leaseChecks)).toBe(true) + const bounds = dbChainMockFns.execute.mock.calls + .map(([query], 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 faa486faf66..8d53699df30 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.ts @@ -68,6 +68,7 @@ import { createContentSyncLease, holdsSyncLockToken, LOCKABLE_CONNECTOR_STATUSES, + leaseTransaction, RUNNABLE_CONNECTOR_STATUSES, SyncLockLostException, type SyncRunLease, @@ -186,12 +187,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, @@ -448,21 +450,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) @@ -516,7 +503,7 @@ export async function completeSuccessfulSync( completionNotice, result.docsFailed === 0 && (!contentPass || !isContentPassIncomplete(contentPass)) ), - /** Restored above under this same lock, or hidden by the admin pass before the ACLs it wrote. */ + /** Restored before completion under this run's lease, or hidden by the admin pass before the ACLs it wrote. */ accessRewritePending: false, ...(contentPass?.complete ? { listingCheckpoint: null } : {}), ...(contentPass && !isContentPassIncomplete(contentPass) && result.docsFailed === 0 @@ -1332,6 +1319,31 @@ export async function executeSync( lease, }) + /** + * 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. + * Before the completion transaction, one short lease-proving transaction + * per page that also re-checks the mode, so a reclaimed run cannot + * rewrite a connector that has since changed mode, and a large restore + * never holds the connector row across its projection fan-out. Nothing + * can change the mode while this run holds its lease, and the completion + * write proves it still does. + */ + if (accessMode === 'workspace') { + const restoredAcls = await restoreWorkspaceDocumentAcls( + connectorId, + leaseTransaction(connectorId, lease), + lease.beatIfDue + ) + if (restoredAcls > 0) { + logger.warn('Restored workspace access on connector documents that had drifted', { + connectorId, + restoredAcls, + }) + } + } + const completionLanded = await completeSuccessfulSync( connectorId, connector.knowledgeBaseId, diff --git a/apps/sim/lib/knowledge/connectors/sync-limits.ts b/apps/sim/lib/knowledge/connectors/sync-limits.ts index 798258e78b9..e260ff78725 100644 --- a/apps/sim/lib/knowledge/connectors/sync-limits.ts +++ b/apps/sim/lib/knowledge/connectors/sync-limits.ts @@ -155,3 +155,22 @@ 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 + +/** + * 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 most + * documents one connector-lease transaction assigns an ACL to, so no transaction + * holds the connector row across more fan-out than one statement's. + */ +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..80560a0f04a 100644 --- a/apps/sim/lib/knowledge/connectors/sync-lock.ts +++ b/apps/sim/lib/knowledge/connectors/sync-lock.ts @@ -1,6 +1,7 @@ import { db } from '@sim/db' import { knowledgeConnector } from '@sim/db/schema' -import { and, eq, isNull } from 'drizzle-orm' +import { and, eq, isNull, sql } from 'drizzle-orm' +import type { DbOrTx } from '@/lib/db/types' import { SYNC_LOCK_HEARTBEAT_INTERVAL_MS } from '@/lib/knowledge/connectors/sync-limits' /** @@ -235,6 +236,40 @@ export async function assertSyncLeaseHeldInTx( if (!held) throw new SyncLockLostException(connectorId) } +/** + * The bounds of a transaction that holds a connector row lock while it writes documents. The + * `document` ACL trigger rewrites every filled search projection row of a document whose ACL is + * assigned, so one page of ACL writes can outlast the role's statement timeout under I/O pressure. + * A page that waits on a lock or runs long fails in seconds and rolls back only itself, instead of + * holding the connector row, and the table lock a migration queues behind, for a minute. + */ +export async function boundLeaseTransaction(tx: Pick): Promise { + await tx.execute( + sql`SELECT set_config('lock_timeout', '5s', true), set_config('statement_timeout', '30s', 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` before `write` runs, so a run that + * lost its lease between pages writes nothing further. 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) + if (lease) await assertSyncLeaseHeldInTx(tx, connectorId, lease) + return write(tx) + }) +} + /** * 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..adeeffffde4 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!) @@ -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..d8effe913d3 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,6 +60,23 @@ import { const CONNECTOR = 'connector-1' +/** 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) + } +} + /** Each `update(...).where(...)` chain ends in `returning()`; one row per changed document. */ function queueUpdatedCounts(...counts: number[]) { for (const count of counts) { @@ -291,6 +309,46 @@ 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']]) + ) + + /** One transaction per statement: a lease lock held across batches outlasted the statement timeout. */ + it('writes every batch in a transaction of its own', async () => { + await persistDocumentAcls(CONNECTOR, changedGroup(60), direct) + + const assignments = dbChainMockFns.set.mock.calls.filter(([values]) => 'acl' in values) + expect(assignments).toHaveLength(3) + expect(pages).toHaveBeenCalledTimes(dbChainMockFns.set.mock.calls.length) + }) + + it('writes nothing further once the lease is lost between batches', async () => { + await expect( + persistDocumentAcls(CONNECTOR, changedGroup(60), losingLease(2)) + ).rejects.toBeInstanceOf(SyncLockLostException) + + /** The evidence refresh and the first change batch landed; the rest never ran. */ + expect(dbChainMockFns.set).toHaveBeenCalledTimes(2) + }) + + it('bounds each batch in a transaction of its own by default', async () => { + await persistDocumentAcls(CONNECTOR, changedGroup(30)) + + expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(dbChainMockFns.set.mock.calls.length) + const bounds = dbChainMockFns.execute.mock.calls.filter(([query]) => + JSON.stringify(query).includes('lock_timeout') + ) + expect(bounds).toHaveLength(dbChainMockFns.transaction.mock.calls.length) + }) +}) + describe('revokeDocumentAcls', () => { beforeEach(() => { vi.clearAllMocks() @@ -320,7 +378,7 @@ describe('revokeDocumentAcls', () => { * 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) + await revokeDocumentAcls(direct, ['a', 'b'], scope) const writes = statements().filter(({ values }) => 'acl' in values) expect(writes).toHaveLength(1) @@ -329,7 +387,7 @@ describe('revokeDocumentAcls', () => { }) it('clears leftover evidence on an already-empty ACL without assigning acl', async () => { - await revokeDocumentAcls(db, ['a', 'b'], scope) + await revokeDocumentAcls(direct, ['a', 'b'], scope) const clears = statements().filter(({ values }) => !('acl' in values)) expect(clears).toHaveLength(1) @@ -345,7 +403,7 @@ describe('revokeDocumentAcls', () => { it('assigns acl in batches of 25 and clears evidence in batches of 500', async () => { const ids = Array.from({ length: 60 }, (_unused, index) => `doc-${index}`) - await revokeDocumentAcls(db, ids, scope) + await revokeDocumentAcls(direct, ids, scope) const batchSizes = (writesAcl: boolean) => statements() @@ -357,6 +415,25 @@ describe('revokeDocumentAcls', () => { expect(batchSizes(true)).toEqual([25, 25, 10]) expect(batchSizes(false)).toEqual([60]) }) + + it('runs every batch in a transaction of its own', async () => { + const ids = Array.from({ length: 60 }, (_unused, index) => `doc-${index}`) + + await revokeDocumentAcls(direct, ids, scope) + + expect(pages).toHaveBeenCalledTimes(dbChainMockFns.set.mock.calls.length) + expect(pages).toHaveBeenCalledTimes(4) + }) + + it('stops writing at the first batch whose lease is gone', async () => { + const ids = Array.from({ length: 60 }, (_unused, index) => `doc-${index}`) + + await expect(revokeDocumentAcls(losingLease(2), ids, scope)).rejects.toBeInstanceOf( + SyncLockLostException + ) + + expect(dbChainMockFns.set).toHaveBeenCalledTimes(2) + }) }) describe('persistSourceDocumentFailures', () => { diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.ts index ec3799f1dd5..e0c0eb4c7cb 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 { rewriteConnectorDocumentAcls } 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_CHANGE_BATCH_SIZE, + 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,35 @@ 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. */ export async function restoreWorkspaceDocumentAcls( - executor: DbOrTx, - connectorId: string + connectorId: string, + transaction: LeaseTransaction, + beforePage?: () => Promise ): 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') - ) - ) + const { rewritten } = await rewriteConnectorDocumentAcls({ + connectorId, + target: WORKSPACE_ACL, + transaction, + beforePage, + 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 rewritten } -/** - * 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 +101,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 @@ -176,19 +169,23 @@ export async function persistDocumentAcls( ) /** 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 }) + const rows = await transaction((tx) => + tx + .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 }) + const rows = await transaction((tx) => + tx + .update(document) + .set({ acl, aclRequirements: requirements, aclVerifiedAt }) + .where(target(batch, false)) + .returning({ id: document.id }) + ) updated += rows.length } } @@ -202,31 +199,36 @@ export async function persistDocumentAcls( * `acl` assigned, {@link ACL_CHANGE_BATCH_SIZE} at a time: 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)` + await transaction((tx) => + tx + .update(document) + .set({ aclRequirements: [], aclVerifiedAt: null }) + .where( + and( + target(batch), + 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)) + await transaction((tx) => + tx + .update(document) + .set({ acl: [], aclRequirements: [], aclVerifiedAt: null }) + .where(and(target(batch), grants)) + ) } } From 0a7dd0cfac9c6ce88b8ce1976ca5472c7078b414 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 22 Sep 2026 23:41:16 -0700 Subject: [PATCH 2/8] fix(knowledge): isolate stale sweep lock failures, close failed disables, update ACL test callers --- .../connectors/member-sync/route.test.ts | 15 +++ .../knowledge/connectors/member-sync/route.ts | 14 +- .../connector-lease-pages.integration.ts | 120 ++++++++++++++++++ .../access/predicate.postgres.test.ts | 33 +++-- .../connectors/member-observations.test.ts | 40 +++++- .../connectors/member-observations.ts | 37 ++++-- .../connectors/member-sync-engine.ts | 44 ++++--- .../knowledge/connectors/sync-engine.test.ts | 3 +- .../connectors/sync-persistence.test.ts | 4 +- 9 files changed, 266 insertions(+), 44 deletions(-) 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 index 4cef8073038..8f6131e3f58 100644 --- a/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts @@ -9,6 +9,8 @@ import { document, knowledgeConnector, knowledgeConnectorMember, + knowledgeConnectorMemberSyncLog, + knowledgeDocumentObservation, organization, resourcePolicy, user, @@ -47,6 +49,7 @@ import { materializeDocumentAcls, recordMemberObservations, rewriteConnectorAcls, + sweepStaleMemberObservations, } from '@/lib/knowledge/connectors/member-observations' import { executeMemberSync, @@ -395,6 +398,83 @@ describe('connector lease ACL pages in PostgreSQL', () => { }) }) + 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, []) @@ -649,5 +729,45 @@ describe('connector lease ACL pages in PostgreSQL', () => { ) 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/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/member-observations.test.ts b/apps/sim/lib/knowledge/connectors/member-observations.test.ts index ce91f892068..0555ac205d7 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.test.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.test.ts @@ -203,12 +203,44 @@ describe('sweepStaleMemberObservations', () => { expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(2) expect(dbChainMockFns.limit).toHaveBeenCalledWith(25) - const bounds = dbChainMockFns.execute.mock.calls.filter(([query]) => - JSON.stringify(query).includes('lock_timeout') + const bounds = dbChainMockFns.execute.mock.calls.filter((call: unknown[]) => + JSON.stringify(call).includes('lock_timeout') ) expect(bounds).toHaveLength(2) }) + /** 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, + { ...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.knowledgeConnector, [{ id: 'c-2' }]) + queueTableRows(schemaMock.knowledgeConnectorMember, [{ id: 'm-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 @@ -336,8 +368,8 @@ describe('rewriteConnectorAcls', () => { expect(assignedPages()).toHaveLength(3) expect(pageSizes()).toEqual([25, 25, 10]) expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(3) - const bounds = dbChainMockFns.execute.mock.calls.filter(([query]) => - JSON.stringify(query).includes('lock_timeout') + const bounds = dbChainMockFns.execute.mock.calls.filter((call: unknown[]) => + JSON.stringify(call).includes('lock_timeout') ) expect(bounds).toHaveLength(3) }) diff --git a/apps/sim/lib/knowledge/connectors/member-observations.ts b/apps/sim/lib/knowledge/connectors/member-observations.ts index 0d8e0a8ef0c..91795585923 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.ts @@ -5,6 +5,8 @@ import { knowledgeConnectorMember, knowledgeDocumentObservation, } from '@sim/db/schema' +import { createLogger } from '@sim/logger' +import { getPostgresErrorCode } from '@sim/utils/errors' import { chunkArray } from '@sim/utils/helpers' import { and, @@ -49,6 +51,8 @@ import { hardDeleteDocuments, } from '@/lib/knowledge/documents/service' +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. @@ -870,6 +874,9 @@ function memberStillStale(memberId: string, cutoff: Date) { ) } +/** Lock, statement-timeout and serialization failures a sweep page defers instead of failing the tick. */ +const SWEEP_DEFERRABLE_CODES = new Set(['55P03', '57014', '40P01', '40001']) + /** * Stale-member observations one sweep tick removes per member: the same * {@link OBSERVATION_BATCH_SIZE} as before, now in pages of @@ -1044,14 +1051,28 @@ export async function sweepStaleMemberObservations(now: Date): Promise { @@ -2407,7 +2408,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, @@ -2426,18 +2429,27 @@ export async function executeMemberSync( return skipped(result, 'connector_deleted_during_sync') } if (error instanceof MemberBindingGoneError) { + const bindingError = error try { - if (!(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-engine.test.ts b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts index 09d8f4bc835..8f8ecbfcfa0 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts @@ -3230,6 +3230,7 @@ describe('executeSync heartbeats during the listing phase', () => { externalId: 'page-1', title: 'Page', content: 'Body', + contentHash: 'hash-1', mimeType: 'text/plain', acl: ['u:reader@example.com'], }, @@ -3257,7 +3258,7 @@ describe('executeSync heartbeats during the listing phase', () => { .map(({ order }) => order) expect(between(leaseChecks)).toBe(true) const bounds = dbChainMockFns.execute.mock.calls - .map(([query], index) => ({ + .map((query: unknown[], index) => ({ query, order: dbChainMockFns.execute.mock.invocationCallOrder[index], })) diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts index d8effe913d3..b363b321854 100644 --- a/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts @@ -342,8 +342,8 @@ describe('persistDocumentAcls paging', () => { await persistDocumentAcls(CONNECTOR, changedGroup(30)) expect(dbChainMockFns.transaction).toHaveBeenCalledTimes(dbChainMockFns.set.mock.calls.length) - const bounds = dbChainMockFns.execute.mock.calls.filter(([query]) => - JSON.stringify(query).includes('lock_timeout') + const bounds = dbChainMockFns.execute.mock.calls.filter((call: unknown[]) => + JSON.stringify(call).includes('lock_timeout') ) expect(bounds).toHaveLength(dbChainMockFns.transaction.mock.calls.length) }) From 2d3f1f3dcffc8489013531026ffa342b13209d8d Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 23 Sep 2026 00:13:21 -0700 Subject: [PATCH 3/8] fix(knowledge): prove connector leases last, page ACL writes by projection rows, and skip unchanged work --- .../connector-lease-pages.integration.ts | 185 ++++++++- .../__integration__/migration-fixture.ts | 3 +- .../lib/knowledge/connectors/detachment.ts | 2 +- .../connectors/member-observations.test.ts | 155 ++++++-- .../connectors/member-observations.ts | 368 +++++++++++------- .../member-sync-engine.integration.test.ts | 39 ++ .../connectors/member-sync-engine.ts | 219 ++++++----- .../sync-content-pass.postgres.test.ts | 2 +- .../connectors/sync-content-pass.test.ts | 12 +- .../knowledge/connectors/sync-content-pass.ts | 21 +- .../knowledge/connectors/sync-engine.test.ts | 15 +- .../lib/knowledge/connectors/sync-engine.ts | 20 +- .../lib/knowledge/connectors/sync-limits.ts | 24 +- .../sim/lib/knowledge/connectors/sync-lock.ts | 28 +- .../sync-persistence.postgres.test.ts | 2 +- .../connectors/sync-persistence.test.ts | 170 ++++++-- .../knowledge/connectors/sync-persistence.ts | 94 +++-- 17 files changed, 968 insertions(+), 391 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts index 8f6131e3f58..8176aa57428 100644 --- a/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts @@ -64,6 +64,7 @@ import { stillHoldsMemberSyncLock, stillHoldsSyncLock, } from '@/lib/knowledge/connectors/sync-lock' +import * as syncPersistence from '@/lib/knowledge/connectors/sync-persistence' import { persistDocumentAcls, restoreWorkspaceDocumentAcls, @@ -140,6 +141,8 @@ describe('connector lease ACL pages in PostgreSQL', () => { 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) @@ -167,7 +170,7 @@ describe('connector lease ACL pages in PostgreSQL', () => { 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: '5s', statement: '30s' }]) + expect([...bounds]).toEqual([{ lock: '15s', statement: '30s' }]) } /** Takes the lease away once `held` pages have committed, as a reclaim between pages would. */ @@ -236,6 +239,134 @@ describe('connector lease ACL pages in PostgreSQL', () => { }) }) + 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 query ILIKE 'update "document" set "acl"%'`) + 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) + + /** 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 query ILIKE 'update "document" set "acl"%'`) + 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 @@ -269,10 +400,10 @@ describe('connector lease ACL pages in PostgreSQL', () => { ).resolves.toBe(0) }) - it('restores drift before a sync completes, outside the completion transaction', async () => { + it('finishes a pending rewrite before a sync completes, outside the completion transaction', async () => { await db .update(knowledgeConnector) - .set({ status: 'active', syncLockToken: null }) + .set({ status: 'active', syncLockToken: null, accessRewritePending: true }) .where(eq(knowledgeConnector.id, ids.connectorId)) const seeded = await seedDocuments(ids.connectorId, []) await db @@ -316,10 +447,56 @@ describe('connector lease ACL pages in PostgreSQL', () => { .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 }) + expect(connector).toEqual({ + status: 'active', + syncLockToken: null, + accessRewritePending: false, + }) + }) + + /** 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 () => { 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/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 0555ac205d7..b3a615c82f7 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.test.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.test.ts @@ -21,6 +21,8 @@ import { inArray } from 'drizzle-orm' import { applyMemberDocumentLifecycle, materializeDocumentAcls, + pagesByProjectionRows, + rematerializeDocumentAcls, removeUnseenMemberObservations, renewMemberObservationsInScopes, rewriteConnectorAcls, @@ -31,7 +33,7 @@ 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, @@ -174,8 +176,12 @@ describe('sweepStaleMemberObservations', () => { }) expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() - expect(dbChainMockFns.for).toHaveBeenNthCalledWith(1, 'share') - expect(dbChainMockFns.for).toHaveBeenNthCalledWith(2, 'update') + /** The member row first, the connector row last: never held while documents are written. */ + expect(dbChainMockFns.for).toHaveBeenNthCalledWith(1, 'update') + expect(dbChainMockFns.for).toHaveBeenNthCalledWith(2, 'share') + expect(dbChainMockFns.for.mock.invocationCallOrder[1]).toBeGreaterThan( + dbChainMockFns.set.mock.invocationCallOrder.at(-1)! + ) expect(dbChainMockFns.delete).toHaveBeenCalledWith(schemaMock.knowledgeDocumentObservation) expect(dbChainMockFns.set).toHaveBeenLastCalledWith({ deletedAt: NOW }) }) @@ -213,13 +219,15 @@ describe('sweepStaleMemberObservations', () => { 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.knowledgeConnector, [{ id: 'c-2' }]) queueTableRows(schemaMock.knowledgeConnectorMember, [{ id: 'm-2' }]) + queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-2' }]) dbChainMockFns.returning .mockResolvedValueOnce([{ documentId: 'd-1' }]) .mockResolvedValueOnce([{ id: 'd-1' }]) @@ -262,16 +270,29 @@ 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() }) }) @@ -301,15 +322,86 @@ describe('materializeDocumentAcls', () => { }) }) +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('rewriteConnectorAcls', () => { beforeEach(() => { vi.clearAllMocks() resetDbChainMock() }) - const stale = (id: string, aclDiffers = true, evidencePresent = false) => ({ + const stale = (id: string, aclDiffers = true, evidencePresent = false, chunkCount = 10) => ({ id, externalId: `ext-${id}`, + chunkCount, aclDiffers, evidencePresent, }) @@ -342,15 +434,38 @@ describe('rewriteConnectorAcls', () => { expect(dbChainMockFns.update).toHaveBeenCalledWith(schemaMock.document) }) - it('stops without writing once the lease is gone', async () => { - queueTableRows(schemaMock.document, [stale('d-1')]) + /** 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).not.toHaveBeenCalled() + expect(dbChainMockFns.update).toHaveBeenCalledOnce() + expect(dbChainMockFns.for.mock.invocationCallOrder[0]).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. */ @@ -410,20 +525,6 @@ describe('rewriteConnectorAcls', () => { expect(dbChainMockFns.update).not.toHaveBeenCalled() }) - - it('walks again after a pass that wrote, so a row changed behind the walk is reached', async () => { - queueTableRows(schemaMock.document, [stale('d-1')]) - queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) - dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'd-1' }]) - queueTableRows(schemaMock.document, [stale('d-0')]) - queueTableRows(schemaMock.knowledgeConnector, [{ id: 'c-1' }]) - dbChainMockFns.returning.mockResolvedValueOnce([{ id: 'd-0' }]) - queueTableRows(schemaMock.document, []) - - await expect(rewriteConnectorAcls('c-1', [], { lease: held })).resolves.toBe(true) - - expect(assignedPages()).toHaveLength(2) - }) }) describe('applyMemberDocumentLifecycle', () => { diff --git a/apps/sim/lib/knowledge/connectors/member-observations.ts b/apps/sim/lib/knowledge/connectors/member-observations.ts index 91795585923..33e3aea3333 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.ts @@ -34,6 +34,7 @@ import { 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 { boundLeaseTransaction, @@ -64,6 +65,8 @@ const OBSERVATION_BATCH_SIZE = 500 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 @@ -341,18 +344,45 @@ interface ConnectorDocumentCursor { id: string } +/** + * 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. + */ +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 the ACL of every document of the connector to `target`, clearing its permission - * evidence, one short transaction per page. The documents are walked in keyset windows of + * 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, a document whose `acl` differs is - * assigned it {@link ACL_CHANGE_BATCH_SIZE} at a time, each batch its own transaction, because each - * assignment rewrites the document's search projection rows; one whose `acl` already matches only - * has its 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. A pass that wrote - * anything is followed by another, so a document a concurrent writer changed behind the walk is - * still reached. Returns the documents written and whether it finished before `deadlineAt`. + * 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 @@ -371,81 +401,76 @@ export async function rewriteConnectorDocumentAcls(input: { 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 (;;) { - let passWrote = false - let after: ConnectorDocumentCursor | undefined - for (;;) { - if (expired()) return { rewritten, finished: false } - await input.beforePage?.() - const window = await db - .select({ - id: document.id, - externalId: document.externalId, - 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 - ) + if (expired()) return { rewritten, finished: false } + await input.beforePage?.() + 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 - ) + ) + .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 - if (rows.length > 0) passWrote = true - } - const changed = window.filter((row) => row.aclDiffers).map((row) => row.id) - for (const batch of chunkArray(changed, ACL_CHANGE_BATCH_SIZE)) { - if (expired()) return { rewritten, finished: false } - await input.beforePage?.() - const rows = await transaction((tx) => - tx - .update(document) - .set({ acl: [...target], aclRequirements: [], aclVerifiedAt: null }) - .where( - and( - eq(document.connectorId, connectorId), - inArray(document.id, batch), - aclDiffers, - guard - ) + ) + .returning({ id: document.id }) + ) + rewritten += rows.length + } + for (const page of pagesByProjectionRows(window.filter((row) => row.aclDiffers))) { + if (expired()) return { rewritten, finished: false } + await input.beforePage?.() + const rows = await transaction((tx) => + tx + .update(document) + .set({ acl: [...target], aclRequirements: [], aclVerifiedAt: null }) + .where( + and( + eq(document.connectorId, connectorId), + inArray(document.id, page), + aclDiffers, + guard ) - .returning({ id: document.id }) - ) - rewritten += rows.length - if (rows.length > 0) passWrote = true - } - const last = window.at(-1) - if (window.length < ACL_WRITE_BATCH_SIZE || !last?.externalId) break - after = { externalId: last.externalId, id: last.id } + ) + .returning({ id: document.id }) + ) + rewritten += rows.length } - if (!passWrote) return { rewritten, finished: true } + 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 } } } @@ -511,6 +536,39 @@ 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)) { + await beforePage?.() + updated += await transaction((tx) => materializeDocumentAcls(connectorId, page, tx)) + } + } + return updated +} + export interface MemberDocumentLifecycleResult { tombstoned: number resurrected: number @@ -900,70 +958,79 @@ async function sweepStaleMemberPage( documentsRematerialized: number docsTombstoned: number } | null> { - return db.transaction(async (tx) => { - await boundLeaseTransaction(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 + 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 = tx - .select({ documentId: knowledgeDocumentObservation.documentId }) - .from(knowledgeDocumentObservation) - .where(eq(knowledgeDocumentObservation.memberId, member.id)) - .limit(ACL_CHANGE_BATCH_SIZE) - const removed = await tx - .delete(knowledgeDocumentObservation) - .where( - and( - eq(knowledgeDocumentObservation.memberId, member.id), - inArray(knowledgeDocumentObservation.documentId, candidates) + const candidates = tx + .select({ documentId: knowledgeDocumentObservation.documentId }) + .from(knowledgeDocumentObservation) + .where(eq(knowledgeDocumentObservation.memberId, member.id)) + .limit(ACL_CHANGE_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({ 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, - } - }) + .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, + } + }) + } 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. @@ -980,15 +1047,21 @@ async function sweepStaleMemberPage( * lists for them rebuilds their observations. Purging is left to a run holding * the lease. * - * Each page of a member's sweep is 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' @@ -1048,7 +1121,11 @@ 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)) let sweptAny = false try { @@ -1068,6 +1145,7 @@ export async function sweepStaleMemberObservations(now: Date): Promise ({ 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 6c3da37ef33..c3a11573c36 100644 --- a/apps/sim/lib/knowledge/connectors/member-sync-engine.ts +++ b/apps/sim/lib/knowledge/connectors/member-sync-engine.ts @@ -57,6 +57,7 @@ import { applyMemberDocumentLifecycle, materializeDocumentAcls, recordMemberObservations, + rematerializeDocumentAcls, removeMemberObservationsForDocuments, removeUnseenMemberObservations, renewMemberObservationsInScopes, @@ -72,6 +73,7 @@ import { } 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, @@ -91,6 +93,7 @@ import { boundLeaseTransaction, createMemberSyncLease, holdsMemberSyncLockToken, + type LeaseTransaction, MEMBER_LOCKABLE_CONNECTOR_STATUSES, SyncLockLostException, stillHoldsMemberSyncLock, @@ -434,27 +437,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. Bounded - * like every connector-lease transaction, and callers keep the ACLs one call - * writes to a page, so the connector row is never held across unbounded fan-out. + * 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) => { - await boundLeaseTransaction(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 }) } @@ -576,56 +584,62 @@ 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(ACL_CHANGE_BATCH_SIZE) - await materializeDocumentAcls( - run.connectorId, - documents.map((row) => row.documentId), - tx - ) - if (checkpoint.removeMember && run.tombstonesUnobserved && documents.length > 0) { - const tombstoned = await tombstoneDocumentsObservedOnlyBy( - 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, + .orderBy(asc(knowledgeDocumentObservation.documentId)) + .limit(ACL_CHANGE_BATCH_SIZE) + await materializeDocumentAcls( run.connectorId, - documents.map((row) => row.documentId) + documents.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 && documents.length > 0) { + const tombstoned = await tombstoneDocumentsObservedOnlyBy( + 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, + run.connectorId, + documents.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: documents.at(-1)!.documentId }, + }) + .where(eq(knowledgeConnectorMember.id, member.id)) + } + }, + { aclPage: true } + ) } } @@ -1405,16 +1419,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 @@ -1573,18 +1590,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, @@ -1825,19 +1843,20 @@ async function disableMemberSync(run: MemberSyncRun, reason: string): Promise - tx - .update(knowledgeConnector) - .set({ - memberSyncStatus: 'disabled', - lastMemberSyncError: reason, - nextMemberSyncAt: null, - memberSyncLockToken: null, - memberSyncLockLeaseAt: null, - updatedAt: new Date(), - }) - .where(holdsMemberSyncLockToken(run.connectorId, run.runId)) - ) + /** One statement that both proves the lease and hands it back with the disabled flag. */ + const disabled = await db + .update(knowledgeConnector) + .set({ + memberSyncStatus: 'disabled', + lastMemberSyncError: reason, + nextMemberSyncAt: null, + memberSyncLockToken: null, + memberSyncLockLeaseAt: null, + updatedAt: new Date(), + }) + .where(stillHoldsMemberSyncLock(run.connectorId, run.runId)) + .returning({ id: knowledgeConnector.id }) + if (disabled.length === 0) throw new SyncLockLostException(run.connectorId) await failMemberSyncLog(run.runId, run.result, reason) logger.warn('Member sync disabled', { connectorId: run.connectorId, reason }) return true @@ -2102,6 +2121,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, @@ -2138,12 +2159,15 @@ export async function executeMemberSync( ).values(), ] /** - * One lease transaction per page of documents, each observing and - * materialising its page together; the failure checkpoint lands with - * the first, since forcing a later relist is the conservative side. + * 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_CHANGE_BATCH_SIZE) : [[]] + 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( @@ -2152,7 +2176,6 @@ export async function executeMemberSync( page, checkpoint.generationId ) - await materializeDocumentAcls(connectorId, page, tx) if (index === 0 && durableCheckpoint && checkpoint.contentFailures) { await tx .update(knowledgeConnectorMember) @@ -2167,6 +2190,7 @@ export async function executeMemberSync( } }) } + await rematerializeDocumentAcls(connectorId, documentIds, aclPage, run.lease.beatIfDue) result.docsListed += attempted.length } if (!serviceContent) { @@ -2339,10 +2363,7 @@ export async function executeMemberSync( await loadDocumentIdsByExternalId(connectorId, relevantIds), connector.syncIntervalMinutes ) - for (const page of chunkArray([...affected], ACL_CHANGE_BATCH_SIZE)) { - await run.lease.beatIfDue() - await withMemberLease(run, (tx) => materializeDocumentAcls(connectorId, page, tx)) - } + await rematerializeDocumentAcls(connectorId, affected, aclPage, run.lease.beatIfDue) } /** A service-owned corpus outlives its last observer; only the content pass removes it. */ 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 e92a8de9852..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 @@ -157,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 '[]', 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 d0ee2f9f0d6..d8a71cbc0dc 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,12 @@ describe('completed listing removal counts', () => { ]) queueTableRows(schemaMock.document, options.revoked ?? []) if (options.revoked?.length) { - /** Every revocation batch is its own lease-proving transaction: one evidence clear, then acl batches. */ + /** Each window reads what still grants someone; pages are bounded by their chunks' rows. */ + queueTableRows( + schemaMock.document, + options.revoked.map(({ id }) => ({ id, chunkCount: 10 })) + ) + /** 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' }]) @@ -320,7 +325,7 @@ 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) @@ -370,6 +375,9 @@ 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 }]) if (options.readCurrent) { queueTableRows(schemaMock.document, [{ fileUrl: options.existing?.fileUrl ?? '' }]) if ( diff --git a/apps/sim/lib/knowledge/connectors/sync-content-pass.ts b/apps/sim/lib/knowledge/connectors/sync-content-pass.ts index b3785edaf4d..b06e47c8d83 100644 --- a/apps/sim/lib/knowledge/connectors/sync-content-pass.ts +++ b/apps/sim/lib/knowledge/connectors/sync-content-pass.ts @@ -27,6 +27,7 @@ import { SOURCE_PERMISSION_ERROR, } from '@/lib/knowledge/connectors/sync-limits' import { + assertSyncLeaseHeldInTx, type LeaseTransaction, leaseTransaction, type SyncRunLease, @@ -91,8 +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 - /** Every lease-fenced write of the pass is one short, bounded transaction. */ - const withLease = leaseTransaction(input.connectorId, input.lease) + /** 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 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"` @@ -182,7 +190,7 @@ 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 revokeDocumentAcls( - withLease, + withAclPage, changed.map((item) => item.externalId), (batch) => and( @@ -325,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 = @@ -348,7 +356,8 @@ export async function runConnectorContentPass(input: ContentPassInput) { async function reconcileCompletedListing( input: ContentPassInput, checkpoint: ListingCheckpoint, - withLease: LeaseTransaction + withLease: LeaseTransaction, + withAclPage: LeaseTransaction ): Promise<{ finished: boolean; notice: string | null }> { if (checkpoint.unsafe || (checkpoint.listingFailures?.count ?? 0) > 0) return { @@ -426,7 +435,7 @@ async function reconcileCompletedListing( const rows = await loadBatch(and(absent, sql`cardinality(${document.acl}) > 0`), 500, after) if (rows.length === 0) break await revokeDocumentAcls( - withLease, + withAclPage, rows.map((row) => row.id), (batch) => and(absent, inArray(document.id, batch)) ) diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts index 8f8ecbfcfa0..3337222a013 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts @@ -3189,6 +3189,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') @@ -3216,11 +3219,13 @@ describe('executeSync heartbeats during the listing phase', () => { } ) - it('proves the lease inside each bounded transaction that writes mirrored permissions', async () => { + 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) => { @@ -3256,7 +3261,13 @@ describe('executeSync heartbeats during the listing phase', () => { })) .filter(({ mode }) => mode === 'share') .map(({ order }) => order) - expect(between(leaseChecks)).toBe(true) + /** 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, diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.ts b/apps/sim/lib/knowledge/connectors/sync-engine.ts index 8d53699df30..7ab18717cbd 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.ts @@ -1320,17 +1320,17 @@ export async function executeSync( }) /** - * 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. - * Before the completion transaction, one short lease-proving transaction - * per page that also re-checks the mode, so a reclaimed run cannot - * rewrite a connector that has since changed mode, and a large restore - * never holds the connector row across its projection fan-out. Nothing - * can change the mode while this run holds its lease, and the completion - * write proves it still does. + * 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. */ - if (accessMode === 'workspace') { + if (accessMode === 'workspace' && connector.accessRewritePending) { const restoredAcls = await restoreWorkspaceDocumentAcls( connectorId, leaseTransaction(connectorId, lease), diff --git a/apps/sim/lib/knowledge/connectors/sync-limits.ts b/apps/sim/lib/knowledge/connectors/sync-limits.ts index e260ff78725..f4134be13c7 100644 --- a/apps/sim/lib/knowledge/connectors/sync-limits.ts +++ b/apps/sim/lib/knowledge/connectors/sync-limits.ts @@ -164,13 +164,31 @@ export const SOURCE_CONTENT_ERROR = */ 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 most - * documents one connector-lease transaction assigns an ACL to, so no transaction - * holds the connector row across more fan-out than one statement's. + * 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 80560a0f04a..8d2ddc72dce 100644 --- a/apps/sim/lib/knowledge/connectors/sync-lock.ts +++ b/apps/sim/lib/knowledge/connectors/sync-lock.ts @@ -2,7 +2,11 @@ import { db } from '@sim/db' import { knowledgeConnector } from '@sim/db/schema' import { and, eq, isNull, sql } from 'drizzle-orm' import type { DbOrTx } from '@/lib/db/types' -import { SYNC_LOCK_HEARTBEAT_INTERVAL_MS } from '@/lib/knowledge/connectors/sync-limits' +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. @@ -237,15 +241,13 @@ export async function assertSyncLeaseHeldInTx( } /** - * The bounds of a transaction that holds a connector row lock while it writes documents. The - * `document` ACL trigger rewrites every filled search projection row of a document whose ACL is - * assigned, so one page of ACL writes can outlast the role's statement timeout under I/O pressure. - * A page that waits on a lock or runs long fails in seconds and rolls back only itself, instead of - * holding the connector row, and the table lock a migration queues behind, for a minute. + * 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', '5s', true), set_config('statement_timeout', '30s', true)` + sql`SELECT set_config('lock_timeout', ${`${LEASE_PAGE_LOCK_TIMEOUT_MS}ms`}, true), set_config('statement_timeout', ${`${LEASE_PAGE_STATEMENT_TIMEOUT_MS}ms`}, true)` ) } @@ -253,9 +255,12 @@ export async function boundLeaseTransaction(tx: Pick): Promis export type LeaseTransaction = (write: (tx: DbOrTx) => Promise) => Promise /** - * One short, bounded transaction per call that proves `lease` before `write` runs, so a run that - * lost its lease between pages writes nothing further. Without a lease the page is only bounded: - * callers outside a sync run have no lease to prove. + * 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, @@ -265,8 +270,9 @@ export function leaseTransaction( return (write) => executor.transaction(async (tx) => { await boundLeaseTransaction(tx) + const written = await write(tx) if (lease) await assertSyncLeaseHeldInTx(tx, connectorId, lease) - return write(tx) + return written }) } 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 adeeffffde4..b8c29fc5f00 100644 --- a/apps/sim/lib/knowledge/connectors/sync-persistence.postgres.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.postgres.test.ts @@ -65,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 )` diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts index b363b321854..eddccd8780a 100644 --- a/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts @@ -77,13 +77,22 @@ function losingLease(held: number): LeaseTransaction { } } -/** Each `update(...).where(...)` chain ends in `returning()`; one row per changed document. */ -function queueUpdatedCounts(...counts: number[]) { - for (const count of counts) { +/** + * 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 + queueTableRows( + schemaMock.document, + Array.from({ length: changed }, (_unused, index) => ({ id: `doc-${index}`, chunkCount: 1 })) + ) + 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', () => { @@ -99,7 +108,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']]])) @@ -120,7 +129,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']]])) @@ -135,7 +144,7 @@ describe('persistDocumentAcls', () => { }) it('reports how many documents received current permission evidence', async () => { - queueUpdatedCounts(2) + queueGroup(['file-1', 'file-2']) await expect( persistDocumentAcls( @@ -153,7 +162,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, @@ -184,7 +194,7 @@ describe('persistDocumentAcls', () => { }) it('groups ACLs that differ only in order or duplication', async () => { - queueUpdatedCounts(2) + queueGroup([], 2) await persistDocumentAcls( CONNECTOR, @@ -207,7 +217,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']] }]]) @@ -221,7 +231,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([ @@ -249,7 +260,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']]])) @@ -262,7 +273,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({ @@ -277,7 +288,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( @@ -286,7 +297,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( @@ -319,26 +331,81 @@ describe('persistDocumentAcls paging', () => { 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) - /** One transaction per statement: a lease lock held across batches outlasted the statement timeout. */ - it('writes every batch in a transaction of its own', async () => { await persistDocumentAcls(CONNECTOR, changedGroup(60), direct) - const assignments = dbChainMockFns.set.mock.calls.filter(([values]) => 'acl' in values) - expect(assignments).toHaveLength(3) + expect(assignedPages()).toEqual([25, 25, 10]) expect(pages).toHaveBeenCalledTimes(dbChainMockFns.set.mock.calls.length) }) - it('writes nothing further once the lease is lost between batches', async () => { + 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 batch landed; the rest never ran. */ + /** The evidence refresh and the first change page landed; the rest never ran. */ expect(dbChainMockFns.set).toHaveBeenCalledTimes(2) }) - it('bounds each batch in a transaction of its own by default', async () => { + 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) @@ -365,19 +432,32 @@ 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 () => { + queueGranting(1) await revokeDocumentAcls(direct, ['a', 'b'], scope) const writes = statements().filter(({ values }) => 'acl' in values) @@ -386,7 +466,15 @@ 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 () => { + queueGranting(1) await revokeDocumentAcls(direct, ['a', 'b'], scope) const clears = statements().filter(({ values }) => !('acl' in values)) @@ -400,24 +488,25 @@ 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 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]) + 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 batch in a transaction of its own', async () => { + 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(direct, ids, scope) @@ -425,8 +514,9 @@ describe('revokeDocumentAcls', () => { expect(pages).toHaveBeenCalledTimes(4) }) - it('stops writing at the first batch whose lease is gone', async () => { + 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 diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.ts index e0c0eb4c7cb..2b6ea118e6e 100644 --- a/apps/sim/lib/knowledge/connectors/sync-persistence.ts +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.ts @@ -15,13 +15,12 @@ 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 { rewriteConnectorDocumentAcls } from '@/lib/knowledge/connectors/member-observations' -import { resolveSourceModifiedAt } from '@/lib/knowledge/connectors/source-modified-at' import { - ACL_CHANGE_BATCH_SIZE, - ACL_WRITE_BATCH_SIZE, - SOURCE_CONTENT_ERROR, -} from '@/lib/knowledge/connectors/sync-limits' + pagesByProjectionRows, + rewriteConnectorDocumentAcls, +} from '@/lib/knowledge/connectors/member-observations' +import { resolveSourceModifiedAt } from '@/lib/knowledge/connectors/source-modified-at' +import { ACL_WRITE_BATCH_SIZE, SOURCE_CONTENT_ERROR } from '@/lib/knowledge/connectors/sync-limits' import { assertSyncLeaseHeldInTx, type LeaseTransaction, @@ -167,26 +166,41 @@ 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 transaction((tx) => - tx + /** + * 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(batch, true)) - .returning({ id: document.id }) - ) - updated += rows.length - } - for (const batch of chunkArray(externalIds, ACL_CHANGE_BATCH_SIZE)) { - const rows = await transaction((tx) => - tx - .update(document) - .set({ acl, aclRequirements: requirements, aclVerifiedAt }) - .where(target(batch, false)) - .returning({ id: document.id }) - ) - updated += rows.length + .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 rows = await transaction((tx) => + tx + .update(document) + .set({ acl, aclRequirements: requirements, aclVerifiedAt }) + .where(and(inArray(document.id, page), target(window, false))) + .returning({ id: document.id }) + ) + updated += rows.length + } } } @@ -196,7 +210,7 @@ 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. Each batch is its own `transaction`, so a lease lost between batches @@ -208,27 +222,31 @@ export async function revokeDocumentAcls( target: (batch: string[]) => SQL | undefined ): Promise { const grants = sql`cardinality(${document.acl}) > 0` - for (const batch of chunkArray(ids, ACL_WRITE_BATCH_SIZE)) { - await transaction((tx) => - tx + 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(batch), + 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 transaction((tx) => - tx - .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 transaction((tx) => + tx + .update(document) + .set({ acl: [], aclRequirements: [], aclVerifiedAt: null }) + .where(and(target(window), inArray(document.id, page), grants)) + ) + } } } From d0d2bcd1deee35a1f45e6c87ac49aac317729197 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 23 Sep 2026 00:39:30 -0700 Subject: [PATCH 4/8] fix(knowledge): row-bound every observation page, size ACL pages from locked chunk counts, and check the sweep budget per page --- .../connector-lease-pages.integration.ts | 65 ++++++- .../connectors/member-observations.test.ts | 143 +++++++++++++-- .../connectors/member-observations.ts | 170 +++++++++++++----- .../connectors/member-sync-engine.ts | 22 ++- .../connectors/sync-content-pass.test.ts | 18 +- .../connectors/sync-persistence.test.ts | 11 +- .../knowledge/connectors/sync-persistence.ts | 28 +-- 7 files changed, 369 insertions(+), 88 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts index 8176aa57428..359da7c67a5 100644 --- a/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts @@ -273,7 +273,9 @@ describe('connector lease ACL pages in PostgreSQL', () => { 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 query ILIKE 'update "document" set "acl"%'`) + 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)) } @@ -290,6 +292,63 @@ describe('connector lease ACL pages in PostgreSQL', () => { 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) @@ -327,7 +386,9 @@ describe('connector lease ACL pages in PostgreSQL', () => { 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 query ILIKE 'update "document" set "acl"%'`) + 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)) } diff --git a/apps/sim/lib/knowledge/connectors/member-observations.test.ts b/apps/sim/lib/knowledge/connectors/member-observations.test.ts index b3a615c82f7..f4743a6d082 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.test.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.test.ts @@ -48,10 +48,13 @@ 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: 25 }, (_, 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) @@ -66,6 +69,25 @@ describe('removeUnseenMemberObservations', () => { 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', () => { @@ -163,6 +185,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' }]) @@ -176,10 +202,13 @@ describe('sweepStaleMemberObservations', () => { }) expect(dbChainMockFns.transaction).toHaveBeenCalledOnce() - /** The member row first, the connector row last: never held while documents are written. */ - expect(dbChainMockFns.for).toHaveBeenNthCalledWith(1, 'update') - expect(dbChainMockFns.for).toHaveBeenNthCalledWith(2, 'share') - expect(dbChainMockFns.for.mock.invocationCallOrder[1]).toBeGreaterThan( + /** 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) @@ -194,6 +223,10 @@ describe('sweepStaleMemberObservations', () => { 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 }))) @@ -215,6 +248,76 @@ describe('sweepStaleMemberObservations', () => { 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, [ @@ -227,6 +330,7 @@ describe('sweepStaleMemberObservations', () => { 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' }]) @@ -412,14 +516,22 @@ describe('rewriteConnectorAcls', () => { .map(([values], index) => ({ values, where: dbChainMockFns.where.mock.calls[index] })) .filter(({ values }) => 'acl' in values) const pageSizes = () => - dbChainMockFns.where.mock.calls - .map(([condition]) => - flattenMockConditions(condition).find( - (node) => node.type === 'inArray' && node.column === schemaMock.document.id + 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 ) - ) - .filter((node) => node !== undefined) - .map((node) => (node!.values as string[]).length) + 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')]) @@ -447,7 +559,8 @@ describe('rewriteConnectorAcls', () => { ) expect(dbChainMockFns.update).toHaveBeenCalledOnce() - expect(dbChainMockFns.for.mock.invocationCallOrder[0]).toBeGreaterThan( + const leaseCheck = dbChainMockFns.for.mock.calls.findIndex(([mode]) => mode === 'share') + expect(dbChainMockFns.for.mock.invocationCallOrder[leaseCheck]).toBeGreaterThan( dbChainMockFns.update.mock.invocationCallOrder[0] ) }) diff --git a/apps/sim/lib/knowledge/connectors/member-observations.ts b/apps/sim/lib/knowledge/connectors/member-observations.ts index 33e3aea3333..86dd25405c2 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.ts @@ -6,7 +6,7 @@ import { knowledgeDocumentObservation, } from '@sim/db/schema' import { createLogger } from '@sim/logger' -import { getPostgresErrorCode } from '@sim/utils/errors' +import { getPostgresErrorCode, getTransientDatabaseFailure } from '@sim/utils/errors' import { chunkArray } from '@sim/utils/helpers' import { and, @@ -236,9 +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. One page of - * {@link ACL_CHANGE_BATCH_SIZE} per call, because `onRemoved` rematerialises - * the page's ACLs in the caller's lease transaction. + * 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, @@ -250,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(ACL_CHANGE_BATCH_SIZE) - const removed = await executor - .delete(knowledgeDocumentObservation) - .where(and(unseen, inArray(knowledgeDocumentObservation.documentId, candidates))) - .returning({ documentId: knowledgeDocumentObservation.documentId }) + 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 < ACL_CHANGE_BATCH_SIZE } + return { + removed: removed.length, + finished: candidates.length < ACL_CHANGE_BATCH_SIZE && rest.length === 0, + } } /** @@ -370,6 +381,59 @@ export function pagesByProjectionRows( return pages } +/** + * 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. Returns the + * rows `write` reports. + */ +export async function writeProjectionPages( + documentIds: readonly string[], + transaction: LeaseTransaction, + write: (tx: DbOrTx, page: string[]) => Promise, + beforePage?: () => Promise +): Promise { + let pending = [...documentIds] + let written = 0 + while (pending.length > 0) { + await beforePage?.() + 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 +} + /** * 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 @@ -450,22 +514,26 @@ export async function rewriteConnectorDocumentAcls(input: { } for (const page of pagesByProjectionRows(window.filter((row) => row.aclDiffers))) { if (expired()) return { rewritten, finished: false } - await input.beforePage?.() - const rows = await transaction((tx) => - tx - .update(document) - .set({ acl: [...target], aclRequirements: [], aclVerifiedAt: null }) - .where( - and( - eq(document.connectorId, connectorId), - inArray(document.id, page), - aclDiffers, - guard - ) - ) - .returning({ id: document.id }) + rewritten += 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, + input.beforePage ) - rewritten += rows.length } const last = window.at(-1) if (window.length < ACL_WRITE_BATCH_SIZE || !last?.externalId) @@ -562,8 +630,12 @@ export async function rematerializeDocumentAcls( ) ) for (const page of pagesByProjectionRows(stale)) { - await beforePage?.() - updated += await transaction((tx) => materializeDocumentAcls(connectorId, page, tx)) + updated += await writeProjectionPages( + page, + transaction, + (tx, locked) => materializeDocumentAcls(connectorId, locked, tx), + beforePage + ) } } return updated @@ -932,9 +1004,6 @@ function memberStillStale(memberId: string, cutoff: Date) { ) } -/** Lock, statement-timeout and serialization failures a sweep page defers instead of failing the tick. */ -const SWEEP_DEFERRABLE_CODES = new Set(['55P03', '57014', '40P01', '40001']) - /** * Stale-member observations one sweep tick removes per member: the same * {@link OBSERVATION_BATCH_SIZE} as before, now in pages of @@ -957,6 +1026,8 @@ async function sweepStaleMemberPage( 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) => { @@ -968,20 +1039,27 @@ async function sweepStaleMemberPage( .for('update') if (!stale) return null - const candidates = tx + const candidates = await tx .select({ documentId: knowledgeDocumentObservation.documentId }) .from(knowledgeDocumentObservation) .where(eq(knowledgeDocumentObservation.memberId, member.id)) .limit(ACL_CHANGE_BATCH_SIZE) - const removed = await tx - .delete(knowledgeDocumentObservation) - .where( - and( - eq(knowledgeDocumentObservation.memberId, member.id), - inArray(knowledgeDocumentObservation.documentId, candidates) - ) - ) - .returning({ documentId: knowledgeDocumentObservation.documentId }) + 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 = @@ -1020,6 +1098,7 @@ async function sweepStaleMemberPage( observationsRemoved: documentIds.length, documentsRematerialized: rematerialized, docsTombstoned: tombstoned.length, + more: candidates.length === ACL_CHANGE_BATCH_SIZE || rest.length > 0, } }) } catch (error) { @@ -1130,26 +1209,29 @@ export async function sweepStaleMemberObservations( 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.observationsRemoved < ACL_CHANGE_BATCH_SIZE) break + if (!swept.more) break } } 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 code = getPostgresErrorCode(error) - if (!code || !SWEEP_DEFERRABLE_CODES.has(code)) throw error + 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, - code, + failure, + code: getPostgresErrorCode(error), }) } if (sweptAny) result.members += 1 diff --git a/apps/sim/lib/knowledge/connectors/member-sync-engine.ts b/apps/sim/lib/knowledge/connectors/member-sync-engine.ts index 4913a22f50b..64e4de04f02 100644 --- a/apps/sim/lib/knowledge/connectors/member-sync-engine.ts +++ b/apps/sim/lib/knowledge/connectors/member-sync-engine.ts @@ -55,6 +55,7 @@ import { } from '@/lib/knowledge/connectors/member-access' import { applyMemberDocumentLifecycle, + lockProjectionPage, materializeDocumentAcls, recordMemberObservations, rematerializeDocumentAcls, @@ -682,25 +683,34 @@ export async function resumeMembershipRewrites( ) .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, + documents.map((row) => row.documentId) + ) + const pageDocuments = documents.slice(0, page.length) await materializeDocumentAcls( run.connectorId, - documents.map((row) => row.documentId), + pageDocuments.map((row) => row.documentId), tx ) - if (checkpoint.removeMember && run.tombstonesUnobserved && documents.length > 0) { + if (checkpoint.removeMember && run.tombstonesUnobserved && pageDocuments.length > 0) { const tombstoned = await tombstoneDocumentsObservedOnlyBy( tx, run.connectorId, member.id, - documents.map((row) => row.documentId) + pageDocuments.map((row) => row.documentId) ) if (run.result) run.result.docsTombstoned += tombstoned } - if (!checkpoint.removeMember && run.tombstonesUnobserved && documents.length > 0) { + if (!checkpoint.removeMember && run.tombstonesUnobserved && pageDocuments.length > 0) { const resurrected = await resurrectObservedDocuments( tx, run.connectorId, - documents.map((row) => row.documentId) + pageDocuments.map((row) => row.documentId) ) if (run.result) run.result.docsResurrected += resurrected } @@ -715,7 +725,7 @@ export async function resumeMembershipRewrites( listingCheckpoint: documents.length === 0 ? null - : { ...checkpoint, cursor: documents.at(-1)!.documentId }, + : { ...checkpoint, cursor: pageDocuments.at(-1)!.documentId }, }) .where(eq(knowledgeConnectorMember.id, member.id)) } 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 d8a71cbc0dc..90e377ddad8 100644 --- a/apps/sim/lib/knowledge/connectors/sync-content-pass.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-content-pass.test.ts @@ -209,10 +209,11 @@ describe('completed listing removal counts', () => { queueTableRows(schemaMock.document, options.revoked ?? []) if (options.revoked?.length) { /** Each window reads what still grants someone; pages are bounded by their chunks' rows. */ - queueTableRows( - schemaMock.document, - options.revoked.map(({ id }) => ({ id, chunkCount: 10 })) - ) + 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++) @@ -332,7 +333,9 @@ describe('completed listing removal counts', () => { 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).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 () => { @@ -376,8 +379,11 @@ 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') + 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 ( diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts index eddccd8780a..5ba27d1ff6c 100644 --- a/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.test.ts @@ -85,10 +85,13 @@ function losingLease(held: number): LeaseTransaction { function queueGroup(refreshed: string[], changed = 0, unrefreshed = changed > 0) { dbChainMockFns.returning.mockResolvedValueOnce(refreshed.map((externalId) => ({ externalId }))) if (!unrefreshed) return - queueTableRows( - schemaMock.document, - Array.from({ length: changed }, (_unused, index) => ({ id: `doc-${index}`, chunkCount: 1 })) - ) + 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: changed }, (_unused, index) => ({ id: `doc-${index}` })) diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.ts index 2b6ea118e6e..2d2efe96e08 100644 --- a/apps/sim/lib/knowledge/connectors/sync-persistence.ts +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.ts @@ -18,6 +18,7 @@ import type { ConnectorFailureDiagnostic } from '@/lib/knowledge/connectors/conn import { pagesByProjectionRows, rewriteConnectorDocumentAcls, + writeProjectionPages, } from '@/lib/knowledge/connectors/member-observations' import { resolveSourceModifiedAt } from '@/lib/knowledge/connectors/source-modified-at' import { ACL_WRITE_BATCH_SIZE, SOURCE_CONTENT_ERROR } from '@/lib/knowledge/connectors/sync-limits' @@ -192,14 +193,18 @@ export async function persistDocumentAcls( }) updated += refreshed for (const page of pagesByProjectionRows(changed)) { - const rows = await transaction((tx) => - tx - .update(document) - .set({ acl, aclRequirements: requirements, aclVerifiedAt }) - .where(and(inArray(document.id, page), target(window, false))) - .returning({ id: document.id }) + updated += 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 += rows.length } } } @@ -240,12 +245,13 @@ export async function revokeDocumentAcls( .where(and(target(window), grants)) }) for (const page of pagesByProjectionRows(granting)) { - await transaction((tx) => - tx + await writeProjectionPages(page, transaction, async (tx, locked) => { + await tx .update(document) .set({ acl: [], aclRequirements: [], aclVerifiedAt: null }) - .where(and(target(window), inArray(document.id, page), grants)) - ) + .where(and(target(window), inArray(document.id, locked), grants)) + return 0 + }) } } } From c157ef138e108e4ca10cac322cbaeac6b1d60b8b Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 23 Sep 2026 01:04:44 -0700 Subject: [PATCH 5/8] fix(knowledge): stop every ACL page walker at the run budget and keep pending rewrites until they finish --- .../connector-lease-pages.integration.ts | 291 +++++++++++++++++- .../connectors/member-observations.test.ts | 77 +++++ .../connectors/member-observations.ts | 27 +- .../knowledge/connectors/sync-engine.test.ts | 43 +++ .../lib/knowledge/connectors/sync-engine.ts | 68 +++- .../knowledge/connectors/sync-persistence.ts | 18 +- 6 files changed, 488 insertions(+), 36 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts index 359da7c67a5..5bcef987d7d 100644 --- a/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts @@ -7,6 +7,7 @@ import { db } from '@sim/db' import { document, + embedding, knowledgeConnector, knowledgeConnectorMember, knowledgeConnectorMemberSyncLog, @@ -45,6 +46,7 @@ import { } 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, @@ -56,6 +58,7 @@ import { 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, @@ -103,6 +106,26 @@ describe('connector lease ACL pages in PostgreSQL', () => { 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 () => { @@ -114,6 +137,7 @@ describe('connector lease ACL pages in PostgreSQL', () => { 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])) @@ -124,6 +148,12 @@ describe('connector lease ACL pages in PostgreSQL', () => { 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() @@ -239,6 +269,112 @@ describe('connector lease ACL pages in PostgreSQL', () => { }) }) + describe('search projection fan-out', () => { + const CHUNKS = 20 + + /** 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)) + /** + * The keyword projection's own sync trigger ships with the Tin migration, which a database + * without the Tin extension skips; write the rows it would, and its installed ACL trigger + * fills them from the document as it does for every insert. + */ + await db.execute(sql` + INSERT INTO embedding_keyword_tin (id, knowledge_base_id, document_id, enabled, content) + SELECT e.id, e.knowledge_base_id, e.document_id, e.enabled, e.content FROM embedding e + WHERE e.document_id IN (${sql.join( + documents.map((entry) => sql`${entry.id}`), + sql`, ` + )}) + ON CONFLICT (id) DO NOTHING`) + 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.every((row) => row.acl === bob() && row.expected === bob())).toBe(true) + 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.every((row) => row.acl === '' && row.expected === '')).toBe(true) + 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 @@ -444,7 +580,7 @@ describe('connector lease ACL pages in PostgreSQL', () => { ids.connectorId, leaseTransaction(ids.connectorId, adminLease()) ) - ).resolves.toBe(DOCUMENTS) + ).resolves.toEqual({ restored: DOCUMENTS, finished: true }) expect(await writesPerTransaction(ids.connectorId)).toEqual([ PAGE, @@ -458,7 +594,7 @@ describe('connector lease ACL pages in PostgreSQL', () => { ids.connectorId, leaseTransaction(ids.connectorId, adminLease()) ) - ).resolves.toBe(0) + ).resolves.toEqual({ restored: 0, finished: true }) }) it('finishes a pending rewrite before a sync completes, outside the completion transaction', async () => { @@ -519,6 +655,155 @@ describe('connector lease ACL pages in PostgreSQL', () => { }) }) + /** 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, then the first page: the budget passes during that page. */ + if (pages === 2) 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 @@ -572,7 +857,7 @@ describe('connector lease ACL pages in PostgreSQL', () => { ids.connectorId, leaseTransaction(ids.connectorId, adminLease()) ) - ).resolves.toBe(0) + ).resolves.toEqual({ restored: 0, finished: true }) expect((await storedAcls(ids.connectorId)).every((acl) => acl.length === 0)).toBe(true) }) diff --git a/apps/sim/lib/knowledge/connectors/member-observations.test.ts b/apps/sim/lib/knowledge/connectors/member-observations.test.ts index f4743a6d082..68e032903b4 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.test.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.test.ts @@ -28,6 +28,7 @@ import { rewriteConnectorAcls, staleMemberWindowMs, sweepStaleMemberObservations, + writeProjectionPages, } from '@/lib/knowledge/connectors/member-observations' import { MEMBER_OBSERVATION_STALE_AFTER_HOURS, @@ -496,6 +497,54 @@ describe('rematerializeDocumentAcls', () => { }) }) +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() + } + }) + + 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']]) + }) +}) + describe('rewriteConnectorAcls', () => { beforeEach(() => { vi.clearAllMocks() @@ -629,6 +678,34 @@ describe('rewriteConnectorAcls', () => { 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')]) diff --git a/apps/sim/lib/knowledge/connectors/member-observations.ts b/apps/sim/lib/knowledge/connectors/member-observations.ts index 86dd25405c2..b240b246e59 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.ts @@ -411,19 +411,23 @@ export async function lockProjectionPage( /** * 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. Returns the - * rows `write` reports. + * 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, - beforePage?: () => Promise -): Promise { + options: { beforePage?: () => Promise; deadlineAt?: number } = {} +): Promise<{ written: number; finished: boolean }> { let pending = [...documentIds] let written = 0 while (pending.length > 0) { - await beforePage?.() + if (options.deadlineAt !== undefined && Date.now() >= options.deadlineAt) + return { written, finished: false } + await options.beforePage?.() 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 } @@ -431,7 +435,7 @@ export async function writeProjectionPages( written += rows pending = rest } - return written + return { written, finished: true } } /** @@ -514,7 +518,7 @@ export async function rewriteConnectorDocumentAcls(input: { } for (const page of pagesByProjectionRows(window.filter((row) => row.aclDiffers))) { if (expired()) return { rewritten, finished: false } - rewritten += await writeProjectionPages( + const written = await writeProjectionPages( page, transaction, async (tx, locked) => @@ -532,8 +536,10 @@ export async function rewriteConnectorDocumentAcls(input: { ) .returning({ id: document.id }) ).length, - input.beforePage + { 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) @@ -630,12 +636,13 @@ export async function rematerializeDocumentAcls( ) ) for (const page of pagesByProjectionRows(stale)) { - updated += await writeProjectionPages( + const { written } = await writeProjectionPages( page, transaction, (tx, locked) => materializeDocumentAcls(connectorId, locked, tx), - beforePage + { beforePage } ) + updated += written } } return updated diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts index d016a2edf2b..0a933b61738 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.test.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.test.ts @@ -2447,6 +2447,49 @@ 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' }]) diff --git a/apps/sim/lib/knowledge/connectors/sync-engine.ts b/apps/sim/lib/knowledge/connectors/sync-engine.ts index 4a69f777004..dca69a75090 100644 --- a/apps/sim/lib/knowledge/connectors/sync-engine.ts +++ b/apps/sim/lib/knowledge/connectors/sync-engine.ts @@ -390,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 = @@ -458,6 +460,7 @@ export async function completeSuccessfulSync( status: directoryNotice || processingDispatchFailed || + accessRewriteUnfinished || (contentPass && isContentPassIncomplete(contentPass)) ? 'partial' : 'completed', @@ -496,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 before completion under this run's lease, 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) } @@ -1074,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, @@ -1233,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 @@ -1296,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 @@ -1359,16 +1391,19 @@ export async function executeSync( * 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 restoredAcls = await restoreWorkspaceDocumentAcls( + const restore = await restoreWorkspaceDocumentAcls( connectorId, leaseTransaction(connectorId, lease), - lease.beatIfDue + { beforePage: lease.beatIfDue, deadlineAt: runDeadlineAt } ) - if (restoredAcls > 0) { + accessRewriteUnfinished = !restore.finished + if (restore.restored > 0) { logger.warn('Restored workspace access on connector documents that had drifted', { connectorId, - restoredAcls, + restoredAcls: restore.restored, + finished: restore.finished, }) } } @@ -1381,7 +1416,8 @@ export async function executeSync( result, reconciliationHoldNotice, contentPass, - directoryNotice + directoryNotice, + accessRewriteUnfinished ) if (!completionLanded) { diff --git a/apps/sim/lib/knowledge/connectors/sync-persistence.ts b/apps/sim/lib/knowledge/connectors/sync-persistence.ts index 2d2efe96e08..0cbe765e1db 100644 --- a/apps/sim/lib/knowledge/connectors/sync-persistence.ts +++ b/apps/sim/lib/knowledge/connectors/sync-persistence.ts @@ -58,18 +58,21 @@ function updatedDocumentAcl(access: ConnectorAccessMode) { * whatever a mode switch or an interrupted rewrite left behind. Idempotent and * 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. + * 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( connectorId: string, transaction: LeaseTransaction, - beforePage?: () => Promise -): Promise { - const { rewritten } = await rewriteConnectorDocumentAcls({ + options: { beforePage?: () => Promise; deadlineAt?: number } = {} +): Promise<{ restored: number; finished: boolean }> { + const { rewritten, finished } = await rewriteConnectorDocumentAcls({ connectorId, target: WORKSPACE_ACL, transaction, - beforePage, + beforePage: options.beforePage, + deadlineAt: options.deadlineAt, guard: exists( db .select({ one: sql`1` }) @@ -82,7 +85,7 @@ export async function restoreWorkspaceDocumentAcls( ) ), }) - return rewritten + return { restored: rewritten, finished } } export interface DocumentAclWriteResult { @@ -193,7 +196,7 @@ export async function persistDocumentAcls( }) updated += refreshed for (const page of pagesByProjectionRows(changed)) { - updated += await writeProjectionPages( + const { written } = await writeProjectionPages( page, transaction, async (tx, locked) => @@ -205,6 +208,7 @@ export async function persistDocumentAcls( .returning({ id: document.id }) ).length ) + updated += written } } } From 76db3569c3a597b3164e12c6d93e2306c0744330 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 23 Sep 2026 01:18:51 -0700 Subject: [PATCH 6/8] test(knowledge): seed filled projection rows so the fan-out test holds on any provisioned database --- .../connector-lease-pages.integration.ts | 38 +++++++++++++------ 1 file changed, 26 insertions(+), 12 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts index 5bcef987d7d..fdbee143365 100644 --- a/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts @@ -17,6 +17,7 @@ import { 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' @@ -272,6 +273,11 @@ describe('connector lease ACL pages in PostgreSQL', () => { 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) => @@ -292,18 +298,26 @@ describe('connector lease ACL pages in PostgreSQL', () => { for (let offset = 0; offset < rows.length; offset += 200) await db.insert(embedding).values(rows.slice(offset, offset + 200)) /** - * The keyword projection's own sync trigger ships with the Tin migration, which a database - * without the Tin extension skips; write the rows it would, and its installed ACL trigger - * fills them from the document as it does for every insert. + * 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) - SELECT e.id, e.knowledge_base_id, e.document_id, e.enabled, e.content FROM embedding e - WHERE e.document_id IN (${sql.join( - documents.map((entry) => sql`${entry.id}`), - sql`, ` - )}) - ON CONFLICT (id) DO NOTHING`) + 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 }) @@ -348,7 +362,7 @@ describe('connector lease ACL pages in PostgreSQL', () => { ).resolves.toEqual({ updated: 30, rejected: 0 }) const after = [...(await projectionAcls(ids.connectorId))] - expect(after.every((row) => row.acl === bob() && row.expected === bob())).toBe(true) + 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) @@ -368,7 +382,7 @@ describe('connector lease ACL pages in PostgreSQL', () => { const after = [...(await projectionAcls(members.connectorId))] expect(after).toHaveLength(2 * 30 * CHUNKS) - expect(after.every((row) => row.acl === '' && row.expected === '')).toBe(true) + expect(after.filter((row) => row.acl !== '' || row.expected !== '')).toEqual([]) expect(Math.max(...(await projectionRowsPerTransaction()))).toBeLessThanOrEqual( PROJECTION_ROW_BATCH_SIZE ) From 2837d6a8fe0e23b3498cb12c7d0d913fb6a23ccd Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 23 Sep 2026 01:29:37 -0700 Subject: [PATCH 7/8] fix(knowledge): check the page deadline after the lease heartbeat --- .../connectors/member-observations.test.ts | 21 +++++++++++++++++++ .../connectors/member-observations.ts | 4 ++-- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/apps/sim/lib/knowledge/connectors/member-observations.test.ts b/apps/sim/lib/knowledge/connectors/member-observations.test.ts index 68e032903b4..3f161739ff8 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.test.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.test.ts @@ -529,6 +529,27 @@ describe('writeProjectionPages', () => { } }) + /** 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 }, diff --git a/apps/sim/lib/knowledge/connectors/member-observations.ts b/apps/sim/lib/knowledge/connectors/member-observations.ts index b240b246e59..eac00e3567f 100644 --- a/apps/sim/lib/knowledge/connectors/member-observations.ts +++ b/apps/sim/lib/knowledge/connectors/member-observations.ts @@ -425,9 +425,9 @@ export async function writeProjectionPages( 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 } - await options.beforePage?.() 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 } @@ -471,8 +471,8 @@ export async function rewriteConnectorDocumentAcls(input: { let rewritten = 0 let after: ConnectorDocumentCursor | undefined for (;;) { - if (expired()) return { rewritten, finished: false } await input.beforePage?.() + if (expired()) return { rewritten, finished: false } const window = await db .select({ id: document.id, From ac7c2bcea6f36d5598715a51987f7a748eb75619 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 23 Sep 2026 01:35:57 -0700 Subject: [PATCH 8/8] test(knowledge): expire the restore budget at the heartbeat before the next window --- .../__integration__/connector-lease-pages.integration.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts index fdbee143365..b8d61052f05 100644 --- a/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts +++ b/apps/sim/lib/knowledge/__integration__/connector-lease-pages.integration.ts @@ -683,8 +683,8 @@ describe('connector lease ACL pages in PostgreSQL', () => { deadlineAt: clock + 1_000, beforePage: async () => { pages += 1 - /** The window read, then the first page: the budget passes during that page. */ - if (pages === 2) clock += 2_000 + /** The window read and the first page run; the budget passes before the next window. */ + if (pages === 3) clock += 2_000 }, } )