Handle unique constraint error on BatchTaskRunItem creation and allow different limits for batchTrigger and batchTriggerAndWait

Author: ericallam

apps/webapp/app/env.server.ts

@@ -362,6 +362,7 @@ const EnvironmentSchema = z.object({
   MAXIMUM_DEV_QUEUE_SIZE: z.coerce.number().int().optional(),
   MAXIMUM_DEPLOYED_QUEUE_SIZE: z.coerce.number().int().optional(),
   MAX_BATCH_V2_TRIGGER_ITEMS: z.coerce.number().int().default(500),
+  MAX_BATCH_AND_WAIT_V2_TRIGGER_ITEMS: z.coerce.number().int().default(500),
 
   REALTIME_STREAM_VERSION: z.enum(["v1", "v2"]).default("v1"),
   BATCH_METADATA_OPERATIONS_FLUSH_INTERVAL_MS: z.coerce.number().int().default(1000),

apps/webapp/app/routes/api.v1.tasks.batch.ts

@@ -40,13 +40,24 @@ const { action, loader } = createActionApiRoute(
     }
 
     // Check the there are fewer than MAX_BATCH_V2_TRIGGER_ITEMS items
-    if (body.items.length > env.MAX_BATCH_V2_TRIGGER_ITEMS) {
-      return json(
-        {
-          error: `Batch size of ${body.items.length} is too large. Maximum allowed batch size is ${env.MAX_BATCH_V2_TRIGGER_ITEMS}.`,
-        },
-        { status: 400 }
-      );
+    if (body.dependentAttempt) {
+      if (body.items.length > env.MAX_BATCH_AND_WAIT_V2_TRIGGER_ITEMS) {
+        return json(
+          {
+            error: `Batch size of ${body.items.length} is too large. Maximum allowed batch size is ${env.MAX_BATCH_AND_WAIT_V2_TRIGGER_ITEMS} when batchTriggerAndWait.`,
+          },
+          { status: 400 }
+        );
+      }
+    } else {
+      if (body.items.length > env.MAX_BATCH_V2_TRIGGER_ITEMS) {
+        return json(
+          {
+            error: `Batch size of ${body.items.length} is too large. Maximum allowed batch size is ${env.MAX_BATCH_V2_TRIGGER_ITEMS}.`,
+          },
+          { status: 400 }
+        );
+      }
     }
 
     const {

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

@@ -5,7 +5,13 @@ import {
   packetRequiresOffloading,
   parsePacket,
 } from "@trigger.dev/core/v3";
-import { BatchTaskRun, Prisma, TaskRunAttempt } from "@trigger.dev/database";
+import {
+  BatchTaskRun,
+  isUniqueConstraintError,
+  Prisma,
+  TaskRunAttempt,
+} from "@trigger.dev/database";
+import { z } from "zod";
 import { $transaction, prisma, PrismaClientOrTransaction } from "~/db.server";
 import { env } from "~/env.server";
 import { batchTaskRunItemStatusForRunStatus } from "~/models/taskRun.server";
@@ -20,9 +26,8 @@ import { downloadPacketFromObjectStore, uploadPacketToObjectStore } from "../r2.
 import { isFinalAttemptStatus, isFinalRunStatus } from "../taskStatus";
 import { startActiveSpan } from "../tracer.server";
 import { BaseService, ServiceValidationError } from "./baseService.server";
-import { OutOfEntitlementError, TriggerTaskService } from "./triggerTask.server";
-import { z } from "zod";
 import { ResumeBatchRunService } from "./resumeBatchRun.server";
+import { OutOfEntitlementError, TriggerTaskService } from "./triggerTask.server";
 
 const PROCESSING_BATCH_SIZE = 50;
 const ASYNC_BATCH_PROCESS_SIZE_THRESHOLD = 20;
@@ -819,20 +824,39 @@ export class BatchTriggerV3Service extends BaseService {
     }
 
     if (!result.isCached) {
-      await $transaction(this._prisma, async (tx) => {
-        await tx.batchTaskRunItem.create({
-          data: {
-            batchTaskRunId: batch.id,
-            taskRunId: result.run.id,
-            status: batchTaskRunItemStatusForRunStatus(result.run.status),
-          },
-        });
+      try {
+        await $transaction(this._prisma, async (tx) => {
+          // [batchTaskRunId, taskRunId] is a unique index
+          await tx.batchTaskRunItem.create({
+            data: {
+              batchTaskRunId: batch.id,
+              taskRunId: result.run.id,
+              status: batchTaskRunItemStatusForRunStatus(result.run.status),
+            },
+          });
 
-        await tx.batchTaskRun.update({
-          where: { id: batch.id },
-          data: { expectedCount: { increment: 1 } },
+          await tx.batchTaskRun.update({
+            where: { id: batch.id },
+            data: { expectedCount: { increment: 1 } },
+          });
         });
-      });
+      } catch (error) {
+        if (isUniqueConstraintError(error, ["batchTaskRunId", "taskRunId"])) {
+          // This means there is already a batchTaskRunItem for this batch and taskRun
+          logger.debug(
+            "[BatchTriggerV2][processBatchTaskRunItem] BatchTaskRunItem already exists",
+            {
+              batchId: batch.friendlyId,
+              runId: task.runId,
+              currentIndex,
+            }
+          );
+
+          return;
+        }
+
+        throw error;
+      }
     }
   }
 

internal-packages/database/src/transaction.ts

@@ -56,3 +56,34 @@ export async function $transaction<R>(
     throw error;
   }
 }
+
+export function isUniqueConstraintError<T extends readonly string[]>(
+  error: unknown,
+  columns: T
+): boolean {
+  if (!isPrismaKnownError(error)) {
+    return false;
+  }
+
+  if (error.code !== "P2002") {
+    return false;
+  }
+
+  const target = error.meta?.target;
+
+  if (!Array.isArray(target)) {
+    return false;
+  }
+
+  if (target.length !== columns.length) {
+    return false;
+  }
+
+  for (let i = 0; i < columns.length; i++) {
+    if (target[i] !== columns[i]) {
+      return false;
+    }
+  }
+
+  return true;
+}