WIP post-merge conflicts

Author: matt-aitken

apps/coordinator/src/checkpointer.ts

@@ -1,5 +1,5 @@
 import { ExponentialBackoff } from "@trigger.dev/core/v3/apps";
-import { testDockerCheckpoint } from "@trigger.dev/core/v3/apps";
+import { testDockerCheckpoint } from "@trigger.dev/core/v3/checkpoints";
 import { nanoid } from "nanoid";
 import fs from "node:fs/promises";
 import { ChaosMonkey } from "./chaosMonkey";

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

@@ -9,6 +9,7 @@ import { BaseService } from "./baseService.server";
 import { CancelAttemptService } from "./cancelAttempt.server";
 import { CancelTaskAttemptDependenciesService } from "./cancelTaskAttemptDependencies.server";
 import { FinalizeTaskRunService } from "./finalizeTaskRun.server";
+import { getTaskEventStoreTableForRun } from "../taskEventStore.server";
 
 type ExtendedTaskRun = Prisma.TaskRunGetPayload<{
   include: {
@@ -83,9 +84,14 @@ export class CancelTaskRunServiceV1 extends BaseService {
       },
     });
 
-    const inProgressEvents = await eventRepository.queryIncompleteEvents({
-      runId: taskRun.friendlyId,
-    });
+    const inProgressEvents = await eventRepository.queryIncompleteEvents(
+      getTaskEventStoreTableForRun(taskRun),
+      {
+        runId: taskRun.friendlyId,
+      },
+      taskRun.createdAt,
+      taskRun.completedAt ?? undefined
+    );
 
     logger.debug("Cancelling in-progress events", {
       inProgressEvents: inProgressEvents.map((event) => event.id),

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

@@ -3,6 +3,8 @@ import {
   packetRequiresOffloading,
   QueueOptions,
   SemanticInternalAttributes,
+  taskRunErrorToString,
+  taskRunErrorEnhancer,
   TriggerTaskRequestBody,
 } from "@trigger.dev/core/v3";
 import {
@@ -39,6 +41,8 @@ import {
   TriggerTaskServiceOptions,
   TriggerTaskServiceResult,
 } from "./triggerTask.server";
+import { getTaskEventStore } from "../taskEventStore.server";
+import { enqueueRun } from "./enqueueRun.server";
 
 /** @deprecated Use TriggerTaskService in `triggerTask.server.ts` instead. */
 export class TriggerTaskServiceV1 extends BaseService {
@@ -168,6 +172,8 @@ export class TriggerTaskServiceV1 extends BaseService {
                   taskIdentifier: true,
                   rootTaskRunId: true,
                   depth: true,
+                  queueTimestamp: true,
+                  queue: true,
                 },
               },
             },
@@ -224,6 +230,8 @@ export class TriggerTaskServiceV1 extends BaseService {
                       taskIdentifier: true,
                       rootTaskRunId: true,
                       depth: true,
+                      queueTimestamp: true,
+                      queue: true,
                     },
                   },
                 },
@@ -276,7 +284,7 @@ export class TriggerTaskServiceV1 extends BaseService {
         : undefined;
 
       try {
-        return await eventRepository.traceEvent(
+        const result = await eventRepository.traceEvent(
           taskId,
           {
             context: options.traceContext,
@@ -349,6 +357,12 @@ export class TriggerTaskServiceV1 extends BaseService {
                   ? dependentBatchRun.dependentTaskAttempt.taskRun.depth + 1
                   : 0;
 
+                const queueTimestamp =
+                  dependentAttempt?.taskRun.queueTimestamp ??
+                  dependentBatchRun?.dependentTaskAttempt?.taskRun.queueTimestamp ??
+                  delayUntil ??
+                  new Date();
+
                 const taskRun = await tx.taskRun.create({
                   data: {
                     status: delayUntil ? "DELAYED" : "PENDING",
@@ -376,7 +390,9 @@ export class TriggerTaskServiceV1 extends BaseService {
                     isTest: body.options?.test ?? false,
                     delayUntil,
                     queuedAt: delayUntil ? undefined : new Date(),
+                    queueTimestamp,
                     maxAttempts: body.options?.maxAttempts,
+                    taskEventStore: getTaskEventStore(),
                     ttl,
                     tags:
                       tagIds.length === 0
@@ -528,44 +544,61 @@ export class TriggerTaskServiceV1 extends BaseService {
               this._prisma
             );
 
-            //release the concurrency for the env and org, if part of a (batch)triggerAndWait
-            if (dependentAttempt) {
-              const isSameTask = dependentAttempt.taskRun.taskIdentifier === taskId;
-              await marqs?.releaseConcurrency(dependentAttempt.taskRun.id, isSameTask);
-            }
-            if (dependentBatchRun?.dependentTaskAttempt) {
-              const isSameTask =
-                dependentBatchRun.dependentTaskAttempt.taskRun.taskIdentifier === taskId;
-              await marqs?.releaseConcurrency(
-                dependentBatchRun.dependentTaskAttempt.taskRun.id,
-                isSameTask
-              );
-            }
-
             if (!run) {
               return;
             }
 
-            // We need to enqueue the task run into the appropriate queue. This is done after the tx completes to prevent a race condition where the task run hasn't been created yet by the time we dequeue.
+            // Now enqueue the run if it's not delayed
             if (run.status === "PENDING") {
-              await marqs?.enqueueMessage(
-                environment,
-                run.queue,
-                run.id,
-                {
-                  type: "EXECUTE",
-                  taskIdentifier: taskId,
-                  projectId: environment.projectId,
-                  environmentId: environment.id,
-                  environmentType: environment.type,
-                },
-                body.options?.concurrencyKey
-              );
+              const enqueueResult = await enqueueRun({
+                env: environment,
+                run,
+                dependentRun:
+                  dependentAttempt?.taskRun ?? dependentBatchRun?.dependentTaskAttempt?.taskRun,
+              });
+
+              if (!enqueueResult.ok) {
+                // Now we need to fail the run with enqueueResult.error and make sure and
+                // set the traced event to failed as well
+                await this._prisma.taskRun.update({
+                  where: { id: run.id },
+                  data: {
+                    status: "SYSTEM_FAILURE",
+                    completedAt: new Date(),
+                    error: enqueueResult.error,
+                  },
+                });
+
+                event.failWithError(enqueueResult.error);
+
+                return {
+                  run,
+                  isCached: false,
+                  error: enqueueResult.error,
+                };
+              }
             }
 
             return { run, isCached: false };
           }
         );
+
+        if (result?.error) {
+          throw new ServiceValidationError(
+            taskRunErrorToString(taskRunErrorEnhancer(result.error))
+          );
+        }
+
+        const run = result?.run;
+
+        if (!run) {
+          return;
+        }
+
+        return {
+          run,
+          isCached: result?.isCached,
+        };
       } catch (error) {
         // Detect a prisma transaction Unique constraint violation
         if (error instanceof Prisma.PrismaClientKnownRequestError) {

apps/webapp/app/v3/taskEventStore.server.ts

@@ -20,6 +20,7 @@ export type TraceEvent = Pick<
   | "level"
   | "events"
   | "environmentType"
+  | "isDebug"
 >;
 
 export type TaskEventStoreTable = "taskEvent" | "taskEventPartitioned";
@@ -122,7 +123,7 @@ export class TaskEventStore {
   ) {
     if (table === "taskEventPartitioned") {
       return await this.readReplica.$queryRaw<TraceEvent[]>`
-        SELECT 
+        SELECT
           "spanId",
           "parentId",
           "runId",
@@ -136,11 +137,12 @@ export class TaskEventStore {
           "isCancelled",
           level,
           events,
-          "environmentType"
+          "environmentType",
+          "isDebug"
         FROM "TaskEventPartitioned"
-        WHERE 
-          "traceId" = ${traceId} 
-          AND "createdAt" >= ${startCreatedAt.toISOString()}::timestamp 
+        WHERE
+          "traceId" = ${traceId}
+          AND "createdAt" >= ${startCreatedAt.toISOString()}::timestamp
           AND "createdAt" < ${(endCreatedAt
             ? new Date(endCreatedAt.getTime() + env.TASK_EVENT_PARTITIONED_WINDOW_IN_SECONDS * 1000)
             : new Date()
@@ -150,7 +152,7 @@ export class TaskEventStore {
       `;
     } else {
       return await this.readReplica.$queryRaw<TraceEvent[]>`
-        SELECT 
+        SELECT
           id,
           "spanId",
           "parentId",
@@ -165,7 +167,8 @@ export class TaskEventStore {
           "isCancelled",
           level,
           events,
-          "environmentType"
+          "environmentType",
+          "isDebug"
         FROM "TaskEvent"
         WHERE "traceId" = ${traceId}
         ORDER BY "startTime" ASC