feat(knowledge): fill the projection source and ACL in shards (#8112)

Author: waleedlatif1

apps/sim/background/projection-source-acl-backfill.ts

@@ -1,8 +1,10 @@
 import { task, tasks } from '@trigger.dev/sdk'
 import { resolveTriggerRegion } from '@/lib/core/async-jobs/region'
 import {
+  PROJECTION_SOURCE_ACL_BACKFILL_SHARDS,
   PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID,
   type ProjectionSourceAclBackfillPayload,
+  projectionSourceAclChainTag,
   runProjectionSourceAclBackfill,
 } from '@/lib/knowledge/search/projection-source-acl-backfill'
 
@@ -13,23 +15,30 @@ const RUN_BUDGET_MS = 60 * 60 * 1000
  * Trigger.dev wrapper around `runProjectionSourceAclBackfill`. A run fills unset rows for up to
  * {@link RUN_BUDGET_MS}, then triggers its continuation from the cursor it reached, so the whole
  * projection is filled across as many bounded runs as it takes. Retry-safe: every run writes only
- * rows still unset, so a retried or restarted run repeats no write. The queue admits one run at a
- * time, so two starts never fill the same pages against each other.
+ * rows still unset, so a retried or restarted run repeats no write. A shard's continuation keeps
+ * its shard, so a sliced fill stays sliced until every slice is done.
  */
 export const projectionSourceAclBackfillTask = task({
   id: PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID,
   machine: 'small-1x',
   retry: { maxAttempts: 3 },
+  /**
+   * One run per shard the id space may be sliced into. Shards fill disjoint ranges, so runs never
+   * fill the same page against each other; an unsliced chain still runs one at a time because each
+   * run triggers its continuation only as it ends.
+   */
   queue: {
     name: PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID,
-    concurrencyLimit: 1,
+    concurrencyLimit: PROJECTION_SOURCE_ACL_BACKFILL_SHARDS,
   },
   run: async (payload: ProjectionSourceAclBackfillPayload) => {
     const cursor = await runProjectionSourceAclBackfill(payload, { budgetMs: RUN_BUDGET_MS })
     if (!cursor) return
     const continuation: ProjectionSourceAclBackfillPayload = { ...payload, cursor }
     await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, continuation, {
       region: await resolveTriggerRegion(),
+      /** The chain's tag rides on every continuation, so a start finds the chain wherever it is. */
+      tags: [projectionSourceAclChainTag(payload.shard)],
     })
   },
 })

apps/sim/lib/knowledge/search/projection-source-acl-backfill.test.ts

@@ -3,12 +3,24 @@
  */
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { mockBackfill, mockEnd, mockPostgres, mockPrewarm, mockTasksTrigger } = vi.hoisted(() => ({
+const {
+  mockBackfill,
+  mockEnd,
+  mockPostgres,
+  mockPrewarm,
+  mockRunsList,
+  mockTasksTrigger,
+  mockUnsafe,
+} = vi.hoisted(() => ({
   mockBackfill: vi.fn(),
   mockEnd: vi.fn(async () => undefined),
   mockPostgres: vi.fn(),
   mockPrewarm: vi.fn(async () => []),
+  mockRunsList: vi.fn(
+    (_query: unknown): AsyncIterable<{ id: string; status: string }> => (async function* () {})()
+  ),
   mockTasksTrigger: vi.fn(async () => ({ id: 'run-1' })),
+  mockUnsafe: vi.fn(async () => [{ unfilled: false }]),
 }))
 
 vi.mock('@sim/db', () => ({ resolveDbUrl: () => 'postgres://localhost:5432/sim' }))
@@ -18,7 +30,10 @@ vi.mock('@sim/db/script-migrations/0021_embedding_search_connector', () => ({
 }))
 vi.mock('postgres', () => ({ default: mockPostgres }))
 vi.mock('@/lib/knowledge/search/prewarm', () => ({ prewarmSearchProjection: mockPrewarm }))
-vi.mock('@trigger.dev/sdk', () => ({ tasks: { trigger: mockTasksTrigger } }))
+vi.mock('@trigger.dev/sdk', () => ({
+  runs: { list: mockRunsList },
+  tasks: { trigger: mockTasksTrigger },
+}))
 vi.mock('@/lib/core/async-jobs/region', () => ({ resolveTriggerRegion: async () => 'us-east-1' }))
 vi.mock('@/lib/core/utils/background', () => ({
   runDetached: (_label: string, work: () => Promise<unknown>) => {
@@ -29,10 +44,11 @@ vi.mock('@/lib/core/utils/background', () => ({
 import {
   enqueueProjectionSourceAclBackfill,
   PROJECTION_PREWARM_BUDGET_MS,
+  projectionSourceAclShardRange,
   runProjectionSourceAclBackfill,
 } from '@/lib/knowledge/search/projection-source-acl-backfill'
 
-const connection = { end: mockEnd }
+const connection = { end: mockEnd, unsafe: mockUnsafe }
 
 describe('runProjectionSourceAclBackfill', () => {
   beforeEach(() => {
@@ -60,8 +76,17 @@ describe('runProjectionSourceAclBackfill', () => {
     expect(mockEnd).toHaveBeenCalledTimes(1)
   })
 
-  it('warms the projections on the same connection once both are filled, before closing it', async () => {
+  it('analyzes and warms the projections on the same connection once both are filled, before closing it', async () => {
     await runProjectionSourceAclBackfill({})
+    /** A row whose document is gone is not the fill's to finish; the probe joins the document. */
+    expect(
+      mockUnsafe.mock.calls.some(([query]) =>
+        String(query).includes('JOIN document d ON d.id = s.document_id WHERE s.acl IS NULL')
+      )
+    ).toBe(true)
+    expect(mockUnsafe.mock.calls.map(([query]) => query)).toEqual(
+      expect.arrayContaining(['ANALYZE embedding_search', 'ANALYZE embedding_keyword_tin'])
+    )
     expect(mockPrewarm).toHaveBeenCalledTimes(1)
     expect(mockPrewarm).toHaveBeenCalledWith(connection, { budgetMs: PROJECTION_PREWARM_BUDGET_MS })
     expect(mockPrewarm.mock.invocationCallOrder[0]).toBeLessThan(
@@ -101,11 +126,65 @@ describe('runProjectionSourceAclBackfill', () => {
     await expect(runProjectionSourceAclBackfill({})).rejects.toThrow('statement timeout')
     expect(mockEnd).toHaveBeenCalledTimes(1)
   })
+
+  it('fills only its shard of the id space in both projections', async () => {
+    await runProjectionSourceAclBackfill({ shard: { index: 1, count: 4 } })
+    for (const [, , options] of mockBackfill.mock.calls) {
+      expect(options).toMatchObject({ afterId: '4', beforeId: '8' })
+    }
+  })
+
+  it('resumes a shard after its cursor and keeps its upper bound', async () => {
+    await runProjectionSourceAclBackfill({
+      shard: { index: 1, count: 4 },
+      cursor: { projection: 'embedding_search', afterId: '5a' },
+    })
+    expect(mockBackfill.mock.calls[0][2]).toMatchObject({ afterId: '5a', beforeId: '8' })
+    expect(mockBackfill.mock.calls[1][2]).toMatchObject({ afterId: '4', beforeId: '8' })
+  })
+
+  it('leaves the analysis and the warm to whoever fills the rows another shard still holds', async () => {
+    mockUnsafe.mockResolvedValueOnce([{ unfilled: true }])
+    await expect(
+      runProjectionSourceAclBackfill({ shard: { index: 0, count: 4 } })
+    ).resolves.toBeNull()
+    expect(mockUnsafe.mock.calls.map(([query]) => query)).not.toContain('ANALYZE embedding_search')
+    expect(mockPrewarm).not.toHaveBeenCalled()
+    expect(mockEnd).toHaveBeenCalledTimes(1)
+  })
+})
+
+describe('projectionSourceAclShardRange', () => {
+  it('slices the hex id space into contiguous ranges', () => {
+    expect(projectionSourceAclShardRange({ index: 0, count: 4 })).toEqual({
+      afterId: '',
+      beforeId: '4',
+    })
+    expect(projectionSourceAclShardRange({ index: 3, count: 4 })).toEqual({
+      afterId: 'c',
+      beforeId: undefined,
+    })
+    expect(projectionSourceAclShardRange({ index: 0, count: 1 })).toEqual({
+      afterId: '',
+      beforeId: undefined,
+    })
+  })
+
+  it.each([
+    [{ index: 0, count: 3 }, 'shard count must divide 16'],
+    [{ index: 0, count: 8 }, 'shard count must be at most 4'],
+    [{ index: 4, count: 4 }, 'shard index must be within 0..3'],
+    [{ index: 0.5, count: 2 }, 'shard index must be within 0..1'],
+  ])('refuses %j', (shard, message) => {
+    expect(() => projectionSourceAclShardRange(shard)).toThrow(message)
+  })
 })
 
 describe('enqueueProjectionSourceAclBackfill', () => {
   beforeEach(() => {
     vi.clearAllMocks()
+    /** No chain in flight unless a case says so. */
+    mockRunsList.mockImplementation(() => (async function* () {})())
     mockPostgres.mockReturnValue(connection)
     mockBackfill.mockResolvedValue({
       projection: 'embedding_search',
@@ -118,13 +197,91 @@ describe('enqueueProjectionSourceAclBackfill', () => {
 
   it('hands the backfill to the Trigger.dev worker when one is configured', async () => {
     await expect(enqueueProjectionSourceAclBackfill({ pageSize: 25 })).resolves.toEqual({
-      runId: 'run-1',
+      runIds: ['run-1'],
+      inFlight: [],
     })
+    expect(mockRunsList).toHaveBeenCalledWith(
+      expect.objectContaining({ tag: 'projection-source-acl-backfill:shard:0/1' })
+    )
     expect(mockTasksTrigger).toHaveBeenCalledWith(
       'projection-source-acl-backfill',
       { pageSize: 25 },
-      { region: 'us-east-1' }
+      {
+        region: 'us-east-1',
+        tags: ['projection-source-acl-backfill:shard:0/1'],
+        idempotencyKey: 'projection-source-acl-backfill:shard:0/1:after:none',
+        idempotencyKeyTTL: '2m',
+      }
     )
     expect(mockBackfill).not.toHaveBeenCalled()
   })
+
+  it('keys a start after a chain that ended on that chain, so a restart is its own start', async () => {
+    mockRunsList.mockImplementation(() =>
+      (async function* () {
+        yield { id: 'run-done', status: 'COMPLETED' }
+      })()
+    )
+    await expect(enqueueProjectionSourceAclBackfill({})).resolves.toEqual({
+      runIds: ['run-1'],
+      inFlight: [],
+    })
+    expect(mockTasksTrigger.mock.calls[0][2].idempotencyKey).toBe(
+      'projection-source-acl-backfill:shard:0/1:after:run-done'
+    )
+  })
+
+  it('refuses a shard the id space cannot be sliced into before starting anything', async () => {
+    await expect(
+      enqueueProjectionSourceAclBackfill({ shard: { index: 5, count: 4 } })
+    ).rejects.toThrow('shard index must be within 0..3')
+    expect(mockTasksTrigger).not.toHaveBeenCalled()
+  })
+
+  it('leaves a range whose chain is still in flight to that chain', async () => {
+    mockRunsList.mockImplementation((query: unknown) =>
+      (async function* () {
+        if ((query as { tag: string }).tag.endsWith(':shard:1/4'))
+          yield { id: 'run-live', status: 'EXECUTING' }
+      })()
+    )
+    await expect(enqueueProjectionSourceAclBackfill({}, 4)).resolves.toEqual({
+      runIds: ['run-1', 'run-1', 'run-1'],
+      inFlight: ['run-live'],
+    })
+    expect(mockTasksTrigger.mock.calls.map(([, payload]) => payload.shard?.index)).toEqual([
+      0, 2, 3,
+    ])
+  })
+
+  it('starts one run per shard, each on its own slice under its own chain tag', async () => {
+    await expect(enqueueProjectionSourceAclBackfill({ pageSize: 25 }, 4)).resolves.toEqual({
+      runIds: ['run-1', 'run-1', 'run-1', 'run-1'],
+      inFlight: [],
+    })
+    expect(mockTasksTrigger.mock.calls.map(([, payload]) => payload)).toEqual(
+      [0, 1, 2, 3].map((index) => ({ pageSize: 25, shard: { index, count: 4 } }))
+    )
+    expect(mockTasksTrigger.mock.calls.map(([, , options]) => options.tags)).toEqual(
+      [0, 1, 2, 3].map((index) => [`projection-source-acl-backfill:shard:${index}/4`])
+    )
+  })
+
+  it.each([
+    [3, 'must divide 16'],
+    [8, 'must be at most 4'],
+  ])('refuses %s shards before starting anything', async (shards, message) => {
+    await expect(enqueueProjectionSourceAclBackfill({}, shards)).rejects.toThrow(message)
+    expect(mockTasksTrigger).not.toHaveBeenCalled()
+  })
+
+  it('refuses to slice a start that carries a cursor, which belongs to one chain', async () => {
+    await expect(
+      enqueueProjectionSourceAclBackfill(
+        { cursor: { projection: 'embedding_search', afterId: '5a' } },
+        4
+      )
+    ).rejects.toThrow('cannot start from a cursor')
+    expect(mockTasksTrigger).not.toHaveBeenCalled()
+  })
 })