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

.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 = 5;
+
+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,36 +408,41 @@ 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,
+        async ({ taskIdentifier, idempotencyKeys }) => {
+          const rows = await this.runStore.findRunsByIdempotencyKeys(
+            { runtimeEnvironmentId: environment.id, taskIdentifier, idempotencyKeys },
+            this._prisma
+          );
+          return rows.map((row) => ({ ...row, taskIdentifier }));
+        },
+        { 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).
     const expiredRunIds = new Set<string>();
 
     const runs = await Promise.all(
       body.items.map(async (item) => {
-        const cachedRun = cachedRuns.find((r) => r.idempotencyKey === item.options?.idempotencyKey);
+        const cachedRun = cachedRuns.find(
+          (r) => r.taskIdentifier === item.task && r.idempotencyKey === item.options?.idempotencyKey
+        );
 
         if (cachedRun) {
           if (cachedRun.idempotencyKeyExpiresAt && cachedRun.idempotencyKeyExpiresAt < new Date()) {

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

@@ -0,0 +1,120 @@
+import { postgresTest } from "@internal/testcontainers";
+import type { PrismaClient } from "@trigger.dev/database";
+import { describe, expect } from "vitest";
+import { PostgresRunStore } from "./PostgresRunStore.js";
+
+async function seedEnvironment(prisma: PrismaClient) {
+  const organization = await prisma.organization.create({
+    data: { title: "Test Organization", slug: "test-organization" },
+  });
+  const project = await prisma.project.create({
+    data: {
+      name: "Test Project",
+      slug: "test-project",
+      externalRef: "proj_1234",
+      organizationId: organization.id,
+    },
+  });
+  const environment = await prisma.runtimeEnvironment.create({
+    data: {
+      type: "DEVELOPMENT",
+      slug: "dev",
+      projectId: project.id,
+      organizationId: organization.id,
+      apiKey: "tr_dev_apikey",
+      pkApiKey: "pk_dev_apikey",
+      shortcode: "short_code",
+    },
+  });
+  return { organization, project, environment };
+}
+
+async function createRun(
+  prisma: PrismaClient,
+  params: {
+    runtimeEnvironmentId: string;
+    projectId: string;
+    friendlyId: string;
+    taskIdentifier: string;
+    idempotencyKey: string;
+    idempotencyKeyExpiresAt?: Date;
+  }
+) {
+  await prisma.taskRun.create({
+    data: {
+      friendlyId: params.friendlyId,
+      taskIdentifier: params.taskIdentifier,
+      idempotencyKey: params.idempotencyKey,
+      idempotencyKeyExpiresAt: params.idempotencyKeyExpiresAt ?? null,
+      payload: "{}",
+      payloadType: "application/json",
+      runtimeEnvironmentId: params.runtimeEnvironmentId,
+      projectId: params.projectId,
+      queue: `task/${params.taskIdentifier}`,
+      traceId: `trace_${params.friendlyId}`,
+      spanId: `span_${params.friendlyId}`,
+      engine: "V2",
+    },
+  });
+}
+
+describe("PostgresRunStore.findRunsByIdempotencyKeys", () => {
+  postgresTest("resolves multiple keys, scoped to (env, task)", async ({ prisma }) => {
+    const { project, environment } = await seedEnvironment(prisma);
+    const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma });
+
+    const expiresAt = new Date("2999-01-01T00:00:00.000Z");
+    await createRun(prisma, {
+      runtimeEnvironmentId: environment.id,
+      projectId: project.id,
+      friendlyId: "run_a1",
+      taskIdentifier: "task-a",
+      idempotencyKey: "idem-1",
+      idempotencyKeyExpiresAt: expiresAt,
+    });
+    await createRun(prisma, {
+      runtimeEnvironmentId: environment.id,
+      projectId: project.id,
+      friendlyId: "run_a2",
+      taskIdentifier: "task-a",
+      idempotencyKey: "idem-2",
+    });
+    await createRun(prisma, {
+      runtimeEnvironmentId: environment.id,
+      projectId: project.id,
+      friendlyId: "run_b1",
+      taskIdentifier: "task-b",
+      idempotencyKey: "idem-1",
+    });
+
+    const rows = await store.findRunsByIdempotencyKeys({
+      runtimeEnvironmentId: environment.id,
+      taskIdentifier: "task-a",
+      idempotencyKeys: ["idem-1", "idem-2", "does-not-exist"],
+    });
+
+    const byKey = new Map(rows.map((r) => [r.idempotencyKey, r]));
+    expect(rows).toHaveLength(2);
+    expect(byKey.get("idem-1")?.friendlyId).toBe("run_a1");
+    expect(byKey.get("idem-2")?.friendlyId).toBe("run_a2");
+    expect(rows.map((r) => r.friendlyId)).not.toContain("run_b1");
+    expect(byKey.get("idem-1")?.idempotencyKeyExpiresAt).toBeInstanceOf(Date);
+    expect(byKey.get("idem-1")?.idempotencyKeyExpiresAt?.toISOString()).toBe(
+      expiresAt.toISOString()
+    );
+    expect(byKey.get("idem-2")?.idempotencyKeyExpiresAt).toBeNull();
+  });
+
+  postgresTest("short-circuits on an empty key list without querying", async ({ prisma }) => {
+    const { environment } = await seedEnvironment(prisma);
+    const store = new PostgresRunStore({ prisma, readOnlyPrisma: prisma });
+
+    const rows = await store.findRunsByIdempotencyKeys({
+      runtimeEnvironmentId: environment.id,
+      taskIdentifier: "task-a",
+      idempotencyKeys: [],
+    });
+
+    expect(rows).toEqual([]);
+  });
+});

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,30 @@ 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, legacyRows] = await Promise.all([
+      this.#new.findRunsByIdempotencyKeys(args, RoutingRunStore.#ownPrimary(this.#new, client)),
+      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