fix(knowledge): row-bound every observation page, size ACL pages from locked chunk counts, and check the sweep budget per page

Author: waleedlatif1

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<void>((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<void>((resolve) => setImmediate(resolve))
         }

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]
     )
   })