perf(webapp,run-store): point-lookup batch idempotency keys

Batch triggers that pass per-item idempotency keys could take seconds when the
target task had a large run history, because the bulk `idempotencyKey IN (...)`
lookup degraded to scanning every run for the (environment, task) pair and
filtering in memory. Look each key up individually via a chunked UNION ALL of
point lookups so the planner always uses a per-key index probe.

Author: ericallam

.server-changes/fix-batch-idempotency-point-lookups.md

@@ -0,0 +1,6 @@
+---
+area: webapp
+type: fix
+---
+
+Speed up idempotency checks on `batchTrigger` calls that use idempotency keys. Large batches against a task with a big run history no longer degrade to multi-second lookups.

apps/webapp/app/v3/services/batchTriggerV3.server.ts

@@ -7,6 +7,7 @@ import { packetRequiresOffloading, parsePacket } from "@trigger.dev/core/v3";
 import type { BatchTaskRun, TaskRunAttempt } from "@trigger.dev/database";
 import { isUniqueConstraintError, Prisma } from "@trigger.dev/database";
 import type { RunStore } from "@internal/run-store";
+import pMap from "p-map";
 import { z } from "zod";
 import type { PrismaClientOrTransaction } from "~/db.server";
 import { prisma } from "~/db.server";
@@ -32,6 +33,16 @@ import { BaseService, ServiceValidationError } from "./baseService.server";
 import { OutOfEntitlementError, TriggerTaskService } from "./triggerTask.server";
 
 const PROCESSING_BATCH_SIZE = 50;
+const IDEMPOTENCY_KEY_LOOKUP_CHUNK_SIZE = 50;
+const IDEMPOTENCY_KEY_LOOKUP_CONCURRENCY = 10;
+
+function chunkArray<T>(items: T[], size: number): T[][] {
+  const chunks: T[][] = [];
+  for (let i = 0; i < items.length; i += size) {
+    chunks.push(items.slice(i, i + size));
+  }
+  return chunks;
+}
 const ASYNC_BATCH_PROCESS_SIZE_THRESHOLD = 20;
 const MAX_ATTEMPTS = 10;
 
@@ -397,28 +408,29 @@ export class BatchTriggerV3Service extends BaseService {
       itemsByTask,
     });
 
-    // Fetch cached runs for each task identifier separately to make use of the index
-    const cachedRuns = await Promise.all(
-      Object.entries(itemsByTask).map(([taskIdentifier, items]) =>
-        this.runStore.findRuns(
-          {
-            where: {
-              runtimeEnvironmentId: environment.id,
-              taskIdentifier,
-              idempotencyKey: {
-                in: items.map((i) => i.options?.idempotencyKey).filter(Boolean),
-              },
-            },
-            select: {
-              friendlyId: true,
-              idempotencyKey: true,
-              idempotencyKeyExpiresAt: true,
-            },
-          },
-          this._prisma
+    const idempotencyKeyLookups = Object.entries(itemsByTask).flatMap(([taskIdentifier, items]) => {
+      const idempotencyKeys = Array.from(
+        new Set(
+          items.map((i) => i.options?.idempotencyKey).filter((key): key is string => Boolean(key))
         )
+      );
+      return chunkArray(idempotencyKeys, IDEMPOTENCY_KEY_LOOKUP_CHUNK_SIZE).map((chunk) => ({
+        taskIdentifier,
+        idempotencyKeys: chunk,
+      }));
+    });
+
+    const cachedRuns = (
+      await pMap(
+        idempotencyKeyLookups,
+        ({ taskIdentifier, idempotencyKeys }) =>
+          this.runStore.findRunsByIdempotencyKeys(
+            { runtimeEnvironmentId: environment.id, taskIdentifier, idempotencyKeys },
+            this._prisma
+          ),
+        { concurrency: IDEMPOTENCY_KEY_LOOKUP_CONCURRENCY }
       )
-    ).then((results) => results.flat());
+    ).flat();
 
     // Build the run IDs in order: reuse an unexpired cached id, else mint a new id (and record any
     // expired cached id so its idempotency key can be cleared below).

internal-packages/run-store/src/PostgresRunStore.ts

@@ -19,6 +19,7 @@ import type {
   ExpireSnapshotInput,
   FinalizeRunData,
   ForWaitpointCompletionContext,
+  IdempotencyKeyRunMatch,
   LockRunData,
   ReadClient,
   RescheduleSnapshotInput,
@@ -1682,6 +1683,21 @@ export class PostgresRunStore implements RunStore {
     return byId;
   }
 
+  async findRunsByIdempotencyKeys(
+    args: { runtimeEnvironmentId: string; taskIdentifier: string; idempotencyKeys: string[] },
+    client?: ReadClient
+  ): Promise<IdempotencyKeyRunMatch[]> {
+    if (args.idempotencyKeys.length === 0) {
+      return [];
+    }
+    const prisma = (client ?? this.readOnlyPrisma) as RunOpsCapableClient;
+    const branches = args.idempotencyKeys.map(
+      (key) =>
+        Prisma.sql`SELECT "friendlyId", "idempotencyKey", "idempotencyKeyExpiresAt" FROM "TaskRun" WHERE "runtimeEnvironmentId" = ${args.runtimeEnvironmentId} AND "taskIdentifier" = ${args.taskIdentifier} AND "idempotencyKey" = ${key}`
+    );
+    return prisma.$queryRaw<IdempotencyKeyRunMatch[]>(Prisma.join(branches, " UNION ALL "));
+  }
+
   // --- run-ops persistence ---
 
   async findLatestExecutionSnapshot(

internal-packages/run-store/src/runOpsStore.ts

@@ -20,6 +20,7 @@ import type {
   ExpireSnapshotInput,
   FinalizeRunData,
   ForWaitpointCompletionContext,
+  IdempotencyKeyRunMatch,
   LockRunData,
   ReadClient,
   RescheduleSnapshotInput,
@@ -460,6 +461,31 @@ export class RoutingRunStore implements RunStore {
     return byId;
   }
 
+  async findRunsByIdempotencyKeys(
+    args: { runtimeEnvironmentId: string; taskIdentifier: string; idempotencyKeys: string[] },
+    client?: ReadClient
+  ): Promise<IdempotencyKeyRunMatch[]> {
+    if (args.idempotencyKeys.length === 0) {
+      return [];
+    }
+    const newRows = await this.#new.findRunsByIdempotencyKeys(
+      args,
+      RoutingRunStore.#ownPrimary(this.#new, client)
+    );
+    const legacyRows = await this.#legacy.findRunsByIdempotencyKeys(
+      args,
+      RoutingRunStore.#ownPrimary(this.#legacy, client)
+    );
+    const byKey = new Map<string, IdempotencyKeyRunMatch>();
+    for (const row of legacyRows) {
+      if (row.idempotencyKey != null) byKey.set(row.idempotencyKey, row);
+    }
+    for (const row of newRows) {
+      if (row.idempotencyKey != null) byKey.set(row.idempotencyKey, row);
+    }
+    return [...byKey.values()];
+  }
+
   // ---------------------------------------------------------------------------
   // TaskRun-core: update-family — route by run id in params
   // ---------------------------------------------------------------------------

internal-packages/run-store/src/types.ts

@@ -22,6 +22,12 @@ import type { Residency } from "@trigger.dev/core/v3/isomorphic";
  */
 export type ReadClient = PrismaClientOrTransaction | PrismaReplicaClient;
 
+export type IdempotencyKeyRunMatch = {
+  friendlyId: string;
+  idempotencyKey: string | null;
+  idempotencyKeyExpiresAt: Date | null;
+};
+
 export type CreateRunSnapshotInput = {
   engine: "V2";
   executionStatus: TaskRunExecutionStatus;
@@ -624,6 +630,17 @@ export interface RunStore {
   ): Promise<Map<string, Prisma.TaskRunGetPayload<{ include: I }>>>;
   findRunsByIds(ids: string[], client?: ReadClient): Promise<Map<string, TaskRun>>;
 
+  /**
+   * Point-lookup a set of idempotency keys within one (runtimeEnvironmentId, taskIdentifier).
+   * Each key is matched by full unique-key equality so the planner always does a per-key index
+   * probe and never falls back to scanning the whole (env, task) range and filtering in memory.
+   * Callers chunk large key sets; this resolves one chunk.
+   */
+  findRunsByIdempotencyKeys(
+    args: { runtimeEnvironmentId: string; taskIdentifier: string; idempotencyKeys: string[] },
+    client?: ReadClient
+  ): Promise<IdempotencyKeyRunMatch[]>;
+
   // --- run-ops persistence ---
   // Snapshots, waitpoints, implicit M:N joins, dependents, attempts and checkpoints. The
   // generic model wrappers are thin generics over the Prisma `*Args` types so include/select