Automatically retry TriggerTaskService when hitting a unique constraint error on idempotency key

ericallam · ericallam · commit 937d0f7e4930 · 2025-02-02T19:58:04.000Z
diff --git a/apps/webapp/app/routes/api.v1.tasks.$taskId.trigger.ts b/apps/webapp/app/routes/api.v1.tasks.$taskId.trigger.ts
@@ -74,7 +74,7 @@ const { action, loader } = createActionApiRoute(
 
       const idempotencyKeyExpiresAt = resolveIdempotencyKeyTTL(idempotencyKeyTTL);
 
-      const run = await service.call(params.taskId, authentication.environment, body, {
+      const result = await service.call(params.taskId, authentication.environment, body, {
         idempotencyKey: idempotencyKey ?? undefined,
         idempotencyKeyExpiresAt: idempotencyKeyExpiresAt,
         triggerVersion: triggerVersion ?? undefined,
@@ -83,19 +83,20 @@ const { action, loader } = createActionApiRoute(
         oneTimeUseToken,
       });
 
-      if (!run) {
+      if (!result) {
         return json({ error: "Task not found" }, { status: 404 });
       }
 
       const $responseHeaders = await responseHeaders(
-        run,
+        result.run,
         authentication.environment,
         triggerClient
       );
 
       return json(
         {
-          id: run.friendlyId,
+          id: result.run.friendlyId,
+          isCached: result.isCached,
         },
         {
           headers: $responseHeaders,
diff --git a/apps/webapp/app/v3/services/batchTriggerTask.server.ts b/apps/webapp/app/v3/services/batchTriggerTask.server.ts
@@ -104,7 +104,7 @@ export class BatchTriggerTaskService extends BaseService {
 
       for (const item of body.items) {
         try {
-          const run = await triggerTaskService.call(
+          const result = await triggerTaskService.call(
             taskId,
             environment,
             {
@@ -123,16 +123,16 @@ export class BatchTriggerTaskService extends BaseService {
             }
           );
 
-          if (run) {
+          if (result) {
             await this._prisma.batchTaskRunItem.create({
               data: {
                 batchTaskRunId: batch.id,
-                taskRunId: run.id,
-                status: batchTaskRunItemStatusForRunStatus(run.status),
+                taskRunId: result.run.id,
+                status: batchTaskRunItemStatusForRunStatus(result.run.status),
               },
             });
 
-            runs.push(run.friendlyId);
+            runs.push(result.run.friendlyId);
           }
 
           index++;
diff --git a/apps/webapp/app/v3/services/batchTriggerV2.server.ts b/apps/webapp/app/v3/services/batchTriggerV2.server.ts
@@ -72,12 +72,10 @@ export class BatchTriggerV2Service extends BaseService {
         environment,
         async (span) => {
           const existingBatch = options.idempotencyKey
-            ? await this._prisma.batchTaskRun.findUnique({
+            ? await this._prisma.batchTaskRun.findFirst({
                 where: {
-                  runtimeEnvironmentId_idempotencyKey: {
-                    runtimeEnvironmentId: environment.id,
-                    idempotencyKey: options.idempotencyKey,
-                  },
+                  runtimeEnvironmentId: environment.id,
+                  idempotencyKey: options.idempotencyKey,
                 },
               })
             : undefined;
@@ -769,7 +767,7 @@ export class BatchTriggerV2Service extends BaseService {
 
     const triggerTaskService = new TriggerTaskService();
 
-    const run = await triggerTaskService.call(
+    const result = await triggerTaskService.call(
       task.item.task,
       environment,
       {
@@ -790,15 +788,15 @@ export class BatchTriggerV2Service extends BaseService {
       }
     );
 
-    if (!run) {
+    if (!result) {
       throw new Error(`Failed to trigger run ${task.runId} for batch ${batch.friendlyId}`);
     }
 
     await this._prisma.batchTaskRunItem.create({
       data: {
         batchTaskRunId: batch.id,
-        taskRunId: run.id,
-        status: batchTaskRunItemStatusForRunStatus(run.status),
+        taskRunId: result.run.id,
+        status: batchTaskRunItemStatusForRunStatus(result.run.status),
       },
     });
   }
diff --git a/apps/webapp/app/v3/services/replayTaskRun.server.ts b/apps/webapp/app/v3/services/replayTaskRun.server.ts
@@ -87,7 +87,7 @@ export class ReplayTaskRunService extends BaseService {
       });
 
       const triggerTaskService = new TriggerTaskService();
-      return await triggerTaskService.call(
+      const result = await triggerTaskService.call(
         existingTaskRun.taskIdentifier,
         authenticatedEnvironment,
         {
@@ -113,6 +113,8 @@ export class ReplayTaskRunService extends BaseService {
           },
         }
       );
+
+      return result?.run;
     } catch (error) {
       if (error instanceof OutOfEntitlementError) {
         return;
diff --git a/apps/webapp/app/v3/services/testTask.server.ts b/apps/webapp/app/v3/services/testTask.server.ts
@@ -15,13 +15,19 @@ export class TestTaskService extends BaseService {
 
     switch (data.triggerSource) {
       case "STANDARD":
-        return await triggerTaskService.call(data.taskIdentifier, authenticatedEnvironment, {
-          payload: data.payload,
-          options: {
-            test: true,
-            metadata: data.metadata,
-          },
-        });
+        const result = await triggerTaskService.call(
+          data.taskIdentifier,
+          authenticatedEnvironment,
+          {
+            payload: data.payload,
+            options: {
+              test: true,
+              metadata: data.metadata,
+            },
+          }
+        );
+
+        return result?.run;
       case "SCHEDULED": {
         const payload = {
           scheduleId: "sched_1234",
@@ -34,7 +40,7 @@ export class TestTaskService extends BaseService {
         };
         const payloadPacket = await stringifyIO(payload);
 
-        return await triggerTaskService.call(
+        const result = await triggerTaskService.call(
           data.taskIdentifier,
           authenticatedEnvironment,
           {
@@ -43,6 +49,8 @@ export class TestTaskService extends BaseService {
           },
           { customIcon: "scheduled" }
         );
+
+        return result?.run;
       }
       default:
         throw new Error("Invalid trigger source");
diff --git a/apps/webapp/app/v3/services/triggerScheduledTask.server.ts b/apps/webapp/app/v3/services/triggerScheduledTask.server.ts
@@ -131,14 +131,14 @@ export class TriggerScheduledTaskService extends BaseService {
           payloadPacket,
         });
 
-        const run = await triggerTask.call(
+        const result = await triggerTask.call(
           instance.taskSchedule.taskIdentifier,
           instance.environment,
           { payload: payloadPacket.data, options: { payloadType: payloadPacket.dataType } },
           { customIcon: "scheduled" }
         );
 
-        if (!run) {
+        if (!result) {
           logger.error("Failed to trigger task", {
             instanceId,
             scheduleId: instance.taskSchedule.friendlyId,
@@ -147,7 +147,7 @@ export class TriggerScheduledTaskService extends BaseService {
         } else {
           await this._prisma.taskRun.update({
             where: {
-              id: run.id,
+              id: result.run.id,
             },
             data: {
               scheduleId: instance.taskSchedule.id,
diff --git a/apps/webapp/app/v3/services/triggerTask.server.ts b/apps/webapp/app/v3/services/triggerTask.server.ts
@@ -26,7 +26,7 @@ import { ExpireEnqueuedRunService } from "./expireEnqueuedRun.server";
 import { guardQueueSizeLimitsForEnv } from "../queueSizeLimits.server";
 import { clampMaxDuration } from "../utils/maxDuration";
 import { resolveIdempotencyKeyTTL } from "~/utils/idempotencyKeys.server";
-import { Prisma } from "@trigger.dev/database";
+import { Prisma, TaskRun } from "@trigger.dev/database";
 import { sanitizeQueueName } from "~/models/taskQueue.server";
 
 export type TriggerTaskServiceOptions = {
@@ -49,15 +49,30 @@ export class OutOfEntitlementError extends Error {
   }
 }
 
+export type TriggerTaskServiceResult = {
+  run: TaskRun;
+  isCached: boolean;
+};
+
+const MAX_ATTEMPTS = 2;
+
 export class TriggerTaskService extends BaseService {
   public async call(
     taskId: string,
     environment: AuthenticatedEnvironment,
     body: TriggerTaskRequestBody,
-    options: TriggerTaskServiceOptions = {}
-  ) {
+    options: TriggerTaskServiceOptions = {},
+    attempt: number = 0
+  ): Promise<TriggerTaskServiceResult | undefined> {
     return await this.traceWithEnv("call()", environment, async (span) => {
       span.setAttribute("taskId", taskId);
+      span.setAttribute("attempt", attempt);
+
+      if (attempt > MAX_ATTEMPTS) {
+        throw new ServiceValidationError(
+          `Failed to trigger ${taskId} after ${MAX_ATTEMPTS} attempts.`
+        );
+      }
 
       // TODO: Add idempotency key expiring here
       const idempotencyKey = options.idempotencyKey ?? body.options?.idempotencyKey;
@@ -74,13 +89,11 @@ export class TriggerTaskService extends BaseService {
           : body.options?.ttl ?? (environment.type === "DEVELOPMENT" ? "10m" : undefined);
 
       const existingRun = idempotencyKey
-        ? await this._prisma.taskRun.findUnique({
+        ? await this._prisma.taskRun.findFirst({
             where: {
-              runtimeEnvironmentId_taskIdentifier_idempotencyKey: {
-                runtimeEnvironmentId: environment.id,
-                idempotencyKey,
-                taskIdentifier: taskId,
-              },
+              runtimeEnvironmentId: environment.id,
+              idempotencyKey,
+              taskIdentifier: taskId,
             },
           })
         : undefined;
@@ -103,7 +116,7 @@ export class TriggerTaskService extends BaseService {
         } else {
           span.setAttribute("runId", existingRun.friendlyId);
 
-          return existingRun;
+          return { run: existingRun, isCached: true };
         }
       }
 
@@ -572,7 +585,7 @@ export class TriggerTaskService extends BaseService {
               );
             }
 
-            return run;
+            return { run, isCached: false };
           }
         );
       } catch (error) {
@@ -608,6 +621,23 @@ export class TriggerTaskService extends BaseService {
               throw new Error(
                 `Failed to trigger ${taskId} as the queue could not be created do to a unique constraint error, please try again.`
               );
+            } else if (
+              Array.isArray(target) &&
+              target.length == 3 &&
+              typeof target[0] === "string" &&
+              typeof target[1] === "string" &&
+              typeof target[2] === "string" &&
+              target[0] == "runtimeEnvironmentId" &&
+              target[1] == "taskIdentifier" &&
+              target[2] == "idempotencyKey"
+            ) {
+              logger.debug("TriggerTask: Idempotency key violation, retrying...", {
+                taskId,
+                environmentId: environment.id,
+                idempotencyKey,
+              });
+              // We need to retry the task run creation as the idempotency key has been used
+              return await this.call(taskId, environment, body, options, attempt + 1);
             } else {
               throw new ServiceValidationError(
                 `Cannot trigger ${taskId} as it has already been triggered with the same idempotency key.`
