feat(knowledge): find a fill chain in flight by its tag before starting another

A start no longer relies on an idempotency key, which would have blocked a legitimate restart for its window and collided across options. Every run of a chain carries the chain's tag, continuations included, and a start lists in-flight runs by that tag first: a range whose chain is still running is left to it, whether the start is sliced or not, and a start after a chain ended starts anew.

Author: waleedlatif1

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

@@ -4,6 +4,7 @@ import {
   PROJECTION_SOURCE_ACL_BACKFILL_SHARDS,
   PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID,
   type ProjectionSourceAclBackfillPayload,
+  projectionSourceAclChainTag,
   runProjectionSourceAclBackfill,
 } from '@/lib/knowledge/search/projection-source-acl-backfill'
 
@@ -36,6 +37,8 @@ export const projectionSourceAclBackfillTask = task({
     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,15 +3,25 @@
  */
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
-const { mockBackfill, mockEnd, mockPostgres, mockPrewarm, mockTasksTrigger, mockUnsafe } =
-  vi.hoisted(() => ({
-    mockBackfill: vi.fn(),
-    mockEnd: vi.fn(async () => undefined),
-    mockPostgres: vi.fn(),
-    mockPrewarm: vi.fn(async () => []),
-    mockTasksTrigger: vi.fn(async () => ({ id: 'run-1' })),
-    mockUnsafe: vi.fn(async () => [{ unfilled: false }]),
-  }))
+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 }> => (async function* () {})()
+  ),
+  mockTasksTrigger: vi.fn(async () => ({ id: 'run-1' })),
+  mockUnsafe: vi.fn(async () => [{ unfilled: false }]),
+}))
 
 vi.mock('@sim/db', () => ({ resolveDbUrl: () => 'postgres://localhost:5432/sim' }))
 vi.mock('@sim/db/script-migrations/0021_embedding_search_connector', () => ({
@@ -20,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>) => {
@@ -164,6 +177,8 @@ describe('projectionSourceAclShardRange', () => {
 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',
@@ -177,28 +192,44 @@ describe('enqueueProjectionSourceAclBackfill', () => {
   it('hands the backfill to the Trigger.dev worker when one is configured', async () => {
     await expect(enqueueProjectionSourceAclBackfill({ pageSize: 25 })).resolves.toEqual({
       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',
-        idempotencyKey: 'projection-source-acl-backfill:1:0',
-        idempotencyKeyTTL: '1h',
-      }
+      { region: 'us-east-1', tags: ['projection-source-acl-backfill:shard:0/1'] }
     )
     expect(mockBackfill).not.toHaveBeenCalled()
   })
 
-  it('starts one run per shard, each on its own slice with its own idempotency key', async () => {
+  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' }
+      })()
+    )
+    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.idempotencyKey)).toEqual(
-      [0, 1, 2, 3].map((index) => `projection-source-acl-backfill:4:${index}`)
+    expect(mockTasksTrigger.mock.calls.map(([, , options]) => options.tags)).toEqual(
+      [0, 1, 2, 3].map((index) => [`projection-source-acl-backfill:shard:${index}/4`])
     )
   })
 

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

@@ -156,25 +156,38 @@ async function projectionsFilled(sql: postgres.Sql): Promise<boolean> {
   return true
 }
 
-/** How long a start stays idempotent: long enough that a retried command finds its runs, not a second set. */
-const ENQUEUE_IDEMPOTENCY_TTL = '1h'
+/** The tag every run of one chain carries, so a chain in flight is found before another is started. */
+export function projectionSourceAclChainTag(shard?: ProjectionSourceAclBackfillShard): string {
+  const { index, count } = shard ?? { index: 0, count: 1 }
+  return `${PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID}:shard:${index}/${count}`
+}
+
+/** A run that has not ended: it, or the continuation it triggers, still owns its range. */
+const IN_FLIGHT_RUN_STATUSES = [
+  'PENDING_VERSION',
+  'QUEUED',
+  'DEQUEUED',
+  'EXECUTING',
+  'WAITING',
+  'DELAYED',
+] as const
 
 /**
  * Starts the backfill on the deployment's Trigger.dev worker, where bounded runs chain until the
  * projections are filled: one chain over the whole id space, or one per shard, each filling its
  * own slice at the same time. Safe to call again at any time: a run only fills rows still unset,
- * and a start repeated within the hour finds the runs it already made rather than making more. A
- * cursor belongs to one chain, so a sliced start takes none: each slice begins at its own bound.
+ * and a range whose chain is still in flight is left to that chain rather than given a second.
+ * A cursor belongs to one chain, so a sliced start takes none: each slice begins at its own bound.
  */
 export async function enqueueProjectionSourceAclBackfill(
   payload: ProjectionSourceAclBackfillPayload = {},
   shards = 1
-): Promise<{ runIds: string[] }> {
+): Promise<{ runIds: string[]; inFlight: string[] }> {
   if (shards !== 1) {
     assertProjectionSourceAclShard({ index: 0, count: shards })
     if (payload.cursor) throw new Error('A sliced projection backfill cannot start from a cursor')
   }
-  const { tasks } = await import('@trigger.dev/sdk')
+  const { runs, tasks } = await import('@trigger.dev/sdk')
   const region = await resolveTriggerRegion()
   const payloads: ProjectionSourceAclBackfillPayload[] =
     shards === 1
@@ -184,14 +197,28 @@ export async function enqueueProjectionSourceAclBackfill(
           shard: { index, count: shards },
         }))
   const runIds: string[] = []
-  for (const [index, shardPayload] of payloads.entries()) {
+  const inFlight: string[] = []
+  for (const shardPayload of payloads) {
+    const tag = projectionSourceAclChainTag(shardPayload.shard)
+    let running: string | undefined
+    for await (const run of runs.list({
+      taskIdentifier: PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID,
+      tag,
+      status: [...IN_FLIGHT_RUN_STATUSES],
+      limit: 1,
+    })) {
+      running = run.id
+    }
+    if (running) {
+      inFlight.push(running)
+      continue
+    }
     const handle = await tasks.trigger(PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID, shardPayload, {
       region,
-      idempotencyKey: `${PROJECTION_SOURCE_ACL_BACKFILL_TASK_ID}:${shards}:${index}`,
-      idempotencyKeyTTL: ENQUEUE_IDEMPOTENCY_TTL,
+      tags: [tag],
     })
     runIds.push(handle.id)
   }
-  logger.info('Projection source and ACL backfill enqueued', { runIds, shards })
-  return { runIds }
+  logger.info('Projection source and ACL backfill enqueued', { runIds, inFlight, shards })
+  return { runIds, inFlight }
 }

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

@@ -41,8 +41,8 @@ function shardsFlag(): number {
 async function main(): Promise<void> {
   const shards = shardsFlag()
   if (isTriggerDevEnabled && env.TRIGGER_SECRET_KEY) {
-    const handle = await enqueueProjectionSourceAclBackfill({}, shards)
-    logger.info('Backfill enqueued on the Trigger.dev worker', handle)
+    const started = await enqueueProjectionSourceAclBackfill({}, shards)
+    logger.info('Backfill enqueued on the Trigger.dev worker', started)
     return
   }
   if (shards !== 1)