Various run execution fixes (incl massive memory bloat issue)

This commit fixes various issues with run executions, including a pretty gnarly memory bloat issue when resuming a run that had a decent number of completed tasks (e.g. anything over a few). Other issues fixed:

- Run executions no longer are bound to a queue, which will allow more parallel runs in a single job (instead of 1).
- Serverless function timeouts (504) errors are now handled better, and no longer are retried using the graphile worker failure/retry mechanism (causing massive delays).
- Fixed the job_key design of the performRunExecutionV2 task, which will ensure resumed runs are executed
- Added a mechanism to measure the amount of execution time a given run has accrued, and added a maximum execution duration on the org to be able to limit total execution time for a single run

apps/webapp/app/components/run/RunCompletedDetail.tsx

@@ -54,7 +54,7 @@ export function RunCompletedDetail({ run }: { run: MatchedRun }) {
         <RunPanelDivider />
         {run.error && <RunPanelError text={run.error.message} stackTrace={run.error.stack} />}
         {run.output ? (
-          <CodeBlock language="json" code={run.output} maxLines={8} />
+          <CodeBlock language="json" code={run.output} maxLines={36} />
         ) : (
           run.output === null && <Paragraph variant="small">This run returned nothing</Paragraph>
         )}

apps/webapp/app/models/jobRunExecution.server.ts

@@ -28,6 +28,7 @@ export type EnqueueRunExecutionV2Options = {
   resumeTaskId?: string;
   isRetry?: boolean;
   skipRetrying?: boolean;
+  executionCount?: number;
 };
 
 export async function enqueueRunExecutionV2(
@@ -44,10 +45,11 @@ export async function enqueueRunExecutionV2(
       isRetry: typeof options.isRetry === "boolean" ? options.isRetry : false,
     },
     {
-      queueName: `job:${run.jobId}:env:${run.environmentId}`,
       tx,
       runAt: options.runAt,
-      jobKey: `job_run:${run.id}`,
+      jobKey: `job_run:${run.id}:${options.executionCount ?? 0}${
+        options.resumeTaskId ? `:task:${options.resumeTaskId}` : ""
+      }`,
       maxAttempts: options.skipRetrying ? 1 : undefined,
     }
   );

apps/webapp/app/services/endpointApi.server.ts

@@ -18,6 +18,7 @@ import {
 } from "@trigger.dev/core";
 import { safeBodyFromResponse, safeParseBodyFromResponse } from "~/utils/json";
 import { logger } from "./logger.server";
+import { performance } from "node:perf_hooks";
 
 export class EndpointApiError extends Error {
   constructor(message: string, stack?: string) {
@@ -28,10 +29,7 @@ export class EndpointApiError extends Error {
 }
 
 export class EndpointApi {
-  constructor(
-    private apiKey: string,
-    private url: string
-  ) {}
+  constructor(private apiKey: string, private url: string) {}
 
   async ping(endpointId: string): Promise<PongResponse> {
     const response = await safeFetch(this.url, {
@@ -165,9 +163,7 @@ export class EndpointApi {
   }
 
   async executeJobRequest(options: RunJobBody) {
-    logger.debug("executeJobRequest()", {
-      options,
-    });
+    const startTimeInMs = performance.now();
 
     const response = await safeFetch(this.url, {
       method: "POST",
@@ -183,6 +179,7 @@ export class EndpointApi {
       response,
       parser: RunJobResponseSchema,
       errorParser: ErrorWithStackSchema,
+      durationInMs: Math.floor(performance.now() - startTimeInMs),
     };
   }
 

apps/webapp/app/services/runs/performRunExecutionV2.server.ts

@@ -20,6 +20,13 @@ import { logger } from "../logger.server";
 type FoundRun = NonNullable<Awaited<ReturnType<typeof findRun>>>;
 type FoundTask = FoundRun["tasks"][number];
 
+export type PerformRunExecutionV2Input = {
+  id: string;
+  reason: "PREPROCESS" | "EXECUTE_JOB";
+  isRetry: boolean;
+  resumeTaskId?: string;
+};
+
 export class PerformRunExecutionV2Service {
   #prismaClient: PrismaClient;
 
@@ -27,25 +34,20 @@ export class PerformRunExecutionV2Service {
     this.#prismaClient = prismaClient;
   }
 
-  public async call(
-    id: string,
-    reason: "PREPROCESS" | "EXECUTE_JOB",
-    isRetry: boolean = false,
-    resumeTaskId?: string
-  ) {
-    const run = await findRun(this.#prismaClient, id);
+  public async call(input: PerformRunExecutionV2Input) {
+    const run = await findRun(this.#prismaClient, input.id);
 
     if (!run) {
       return;
     }
 
-    switch (reason) {
+    switch (input.reason) {
       case "PREPROCESS": {
         await this.#executePreprocessing(run);
         break;
       }
       case "EXECUTE_JOB": {
-        await this.#executeJob(run, isRetry, resumeTaskId);
+        await this.#executeJob(run, input);
         break;
       }
     }
@@ -141,7 +143,9 @@ export class PerformRunExecutionV2Service {
       });
     }
   }
-  async #executeJob(run: FoundRun, isRetry: boolean, resumeTaskId?: string) {
+  async #executeJob(run: FoundRun, input: PerformRunExecutionV2Input) {
+    const { isRetry, resumeTaskId } = input;
+
     if (run.status === "CANCELED") {
       await this.#cancelExecution(run);
       return;
@@ -152,21 +156,27 @@ export class PerformRunExecutionV2Service {
 
     const startedAt = new Date();
 
-    await this.#prismaClient.jobRun.update({
+    const { executionCount } = await this.#prismaClient.jobRun.update({
       where: {
         id: run.id,
       },
       data: {
         status: run.status === "QUEUED" ? "STARTED" : run.status,
         startedAt: run.startedAt ?? new Date(),
+        executionCount: {
+          increment: 1,
+        },
+      },
+      select: {
+        executionCount: true,
       },
     });
 
     const connections = await resolveRunConnections(run.runConnections);
 
     if (!connections.success) {
-      return this.#failRunExecutionWithRetry({
-        message: `Could not resolve all connections for run ${run.id}, attempting to retry`,
+      return this.#failRunExecution(this.#prismaClient, "EXECUTE_JOB", run, {
+        message: `Could not resolve all connections for run ${run.id}. This should not happen`,
       });
     }
 
@@ -195,7 +205,7 @@ export class PerformRunExecutionV2Service {
 
     const sourceContext = RunSourceContextSchema.safeParse(run.event.sourceContext);
 
-    const { response, parser, errorParser } = await client.executeJobRequest({
+    const { response, parser, errorParser, durationInMs } = await client.executeJobRequest({
       event,
       job: {
         id: run.version.job.slug,
@@ -261,50 +271,82 @@ export class PerformRunExecutionV2Service {
 
       // Only retry if the error isn't a 4xx
       if (response.status >= 400 && response.status <= 499 && response.status !== 408) {
-        return await this.#failRunExecution(this.#prismaClient, "EXECUTE_JOB", run, {
-          message: `Endpoint responded with ${response.status} status code`,
-        });
+        return await this.#failRunExecution(
+          this.#prismaClient,
+          "EXECUTE_JOB",
+          run,
+          {
+            message: `Endpoint responded with ${response.status} status code`,
+          },
+          "FAILURE",
+          durationInMs
+        );
       } else {
-        return await this.#failRunExecutionWithRetry({
-          message: `Endpoint responded with ${response.status} status code`,
-        });
+        // If the error is a 504 timeout, we should mark this execution as succeeded (by not throwing an error) and enqueue a new execution
+        if (response.status === 504) {
+          return await this.#resumeRunExecutionAfterTimeout(
+            this.#prismaClient,
+            run,
+            input,
+            durationInMs,
+            executionCount
+          );
+        } else {
+          return await this.#failRunExecutionWithRetry({
+            message: `Endpoint responded with ${response.status} status code`,
+          });
+        }
       }
     }
 
     const safeBody = safeJsonZodParse(parser, rawBody);
 
     if (!safeBody) {
-      return await this.#failRunExecution(this.#prismaClient, "EXECUTE_JOB", run, {
-        message: "Endpoint responded with invalid JSON",
-      });
+      return await this.#failRunExecution(
+        this.#prismaClient,
+        "EXECUTE_JOB",
+        run,
+        {
+          message: "Endpoint responded with invalid JSON",
+        },
+        "FAILURE",
+        durationInMs
+      );
     }
 
     if (!safeBody.success) {
-      return await this.#failRunExecution(this.#prismaClient, "EXECUTE_JOB", run, {
-        message: generateErrorMessage(safeBody.error.issues),
-      });
+      return await this.#failRunExecution(
+        this.#prismaClient,
+        "EXECUTE_JOB",
+        run,
+        {
+          message: generateErrorMessage(safeBody.error.issues),
+        },
+        "FAILURE",
+        durationInMs
+      );
     }
 
     const status = safeBody.data.status;
 
     switch (status) {
       case "SUCCESS": {
-        await this.#completeRunWithSuccess(run, safeBody.data);
+        await this.#completeRunWithSuccess(run, safeBody.data, durationInMs);
 
         break;
       }
       case "RESUME_WITH_TASK": {
-        await this.#resumeRunWithTask(run, safeBody.data, isRetry);
+        await this.#resumeRunWithTask(run, safeBody.data, isRetry, durationInMs, executionCount);
 
         break;
       }
       case "ERROR": {
-        await this.#failRunWithError(run, safeBody.data);
+        await this.#failRunWithError(run, safeBody.data, durationInMs);
 
         break;
       }
       case "RETRY_WITH_TASK": {
-        await this.#retryRunWithTask(run, safeBody.data, isRetry);
+        await this.#retryRunWithTask(run, safeBody.data, isRetry, durationInMs, executionCount);
 
         break;
       }
@@ -319,19 +361,37 @@ export class PerformRunExecutionV2Service {
     }
   }
 
-  async #completeRunWithSuccess(run: FoundRun, data: RunJobSuccess) {
+  async #completeRunWithSuccess(run: FoundRun, data: RunJobSuccess, durationInMs: number) {
     await this.#prismaClient.jobRun.update({
       where: { id: run.id },
       data: {
         completedAt: new Date(),
         status: "SUCCESS",
         output: data.output ?? undefined,
+        executionDuration: {
+          increment: durationInMs,
+        },
       },
     });
   }
 
-  async #resumeRunWithTask(run: FoundRun, data: RunJobResumeWithTask, isRetry: boolean) {
+  async #resumeRunWithTask(
+    run: FoundRun,
+    data: RunJobResumeWithTask,
+    isRetry: boolean,
+    durationInMs: number,
+    executionCount: number
+  ) {
     return await $transaction(this.#prismaClient, async (tx) => {
+      await tx.jobRun.update({
+        where: { id: run.id },
+        data: {
+          executionDuration: {
+            increment: durationInMs,
+          },
+        },
+      });
+
       // If the task has an operation, then the next performRunExecution will occur
       // when that operation has finished
       if (!data.task.operation) {
@@ -340,12 +400,13 @@ export class PerformRunExecutionV2Service {
           resumeTaskId: data.task.id,
           isRetry,
           skipRetrying: run.environment.type === RuntimeEnvironmentType.DEVELOPMENT,
+          executionCount,
         });
       }
     });
   }
 
-  async #failRunWithError(execution: FoundRun, data: RunJobError) {
+  async #failRunWithError(execution: FoundRun, data: RunJobError, durationInMs: number) {
     return await $transaction(this.#prismaClient, async (tx) => {
       if (data.task) {
         await tx.task.update({
@@ -360,11 +421,24 @@ export class PerformRunExecutionV2Service {
         });
       }
 
-      await this.#failRunExecution(tx, "EXECUTE_JOB", execution, data.error ?? undefined);
+      await this.#failRunExecution(
+        tx,
+        "EXECUTE_JOB",
+        execution,
+        data.error ?? undefined,
+        "FAILURE",
+        durationInMs
+      );
     });
   }
 
-  async #retryRunWithTask(run: FoundRun, data: RunJobRetryWithTask, isRetry: boolean) {
+  async #retryRunWithTask(
+    run: FoundRun,
+    data: RunJobRetryWithTask,
+    isRetry: boolean,
+    durationInMs: number,
+    executionCount: number
+  ) {
     return await $transaction(this.#prismaClient, async (tx) => {
       // We need to check for an existing task attempt
       const existingAttempt = await tx.taskAttempt.findFirst({
@@ -405,6 +479,13 @@ export class PerformRunExecutionV2Service {
         },
         data: {
           status: "WAITING",
+          run: {
+            update: {
+              executionDuration: {
+                increment: durationInMs,
+              },
+            },
+          },
         },
       });
 
@@ -413,6 +494,55 @@ export class PerformRunExecutionV2Service {
         resumeTaskId: data.task.id,
         isRetry,
         skipRetrying: run.environment.type === RuntimeEnvironmentType.DEVELOPMENT,
+        executionCount,
+      });
+    });
+  }
+
+  async #resumeRunExecutionAfterTimeout(
+    prisma: PrismaClientOrTransaction,
+    run: FoundRun,
+    input: PerformRunExecutionV2Input,
+    durationInMs: number,
+    executionCount: number
+  ) {
+    await $transaction(prisma, async (tx) => {
+      const executionDuration = run.executionDuration + durationInMs;
+
+      // If the execution duration is greater than the maximum execution time, we need to fail the run
+      if (executionDuration >= run.organization.maximumExecutionTimePerRunInMs) {
+        await this.#failRunExecution(
+          tx,
+          "EXECUTE_JOB",
+          run,
+          {
+            message: `Execution timed out after ${
+              run.organization.maximumExecutionTimePerRunInMs / 1000
+            } seconds`,
+          },
+          "TIMED_OUT",
+          durationInMs
+        );
+        return;
+      }
+
+      await tx.jobRun.update({
+        where: {
+          id: run.id,
+        },
+        data: {
+          executionDuration: {
+            increment: durationInMs,
+          },
+        },
+      });
+
+      // The run has timed out, so we need to enqueue a new execution
+      await enqueueRunExecutionV2(run, tx, {
+        resumeTaskId: input.resumeTaskId,
+        isRetry: input.isRetry,
+        skipRetrying: run.environment.type === RuntimeEnvironmentType.DEVELOPMENT,
+        executionCount,
       });
     });
   }
@@ -426,7 +556,8 @@ export class PerformRunExecutionV2Service {
     reason: "EXECUTE_JOB" | "PREPROCESS",
     run: FoundRun,
     output: Record<string, any>,
-    status: "FAILURE" | "ABORTED" = "FAILURE"
+    status: "FAILURE" | "ABORTED" | "TIMED_OUT" = "FAILURE",
+    durationInMs: number = 0
   ): Promise<void> {
     await $transaction(prisma, async (tx) => {
       switch (reason) {
@@ -438,6 +569,9 @@ export class PerformRunExecutionV2Service {
               completedAt: new Date(),
               status,
               output,
+              executionDuration: {
+                increment: durationInMs,
+              },
             },
           });
 
@@ -510,38 +644,23 @@ function prepareTasksForRun(possibleTasks: FoundTask[]): CachedTask[] {
     return cachedTaskSizes.get(taskId)!;
   }
 
-  // Create a dynamic programming array to store intermediate results
-  const dp: number[][] = [];
-  for (let i = 0; i <= tasks.length; i++) {
-    dp[i] = [];
-    for (let j = 0; j <= TOTAL_CACHED_TASK_BYTE_LIMIT; j++) {
-      dp[i][j] = 0;
-    }
-  }
-
-  // Fill the dynamic programming array
-  for (let i = 1; i <= tasks.length; i++) {
-    const task = tasks[i - 1];
+  // Prepare tasks and calculate their sizes
+  const availableTasks = tasks.map((task) => {
     const cachedTask = getCachedTask(task);
-    const taskSize = getCachedTaskSize(cachedTask);
-    for (let j = 0; j <= TOTAL_CACHED_TASK_BYTE_LIMIT; j++) {
-      if (taskSize <= j) {
-        dp[i][j] = Math.max(dp[i - 1][j], dp[i - 1][j - taskSize] + taskSize);
-      } else {
-        dp[i][j] = dp[i - 1][j];
-      }
-    }
-  }
+    return { task: cachedTask, size: getCachedTaskSize(cachedTask) };
+  });
 
-  // Traverse the dynamic programming array to find the included tasks
+  // Sort tasks in ascending order by size
+  availableTasks.sort((a, b) => a.size - b.size);
+
+  // Select tasks using greedy approach
   const tasksToRun: CachedTask[] = [];
-  let j = TOTAL_CACHED_TASK_BYTE_LIMIT;
-  for (let i = tasks.length; i > 0 && j > 0; i--) {
-    if (dp[i][j] !== dp[i - 1][j]) {
-      const task = tasks[i - 1];
-      const cachedTask = getCachedTask(task);
-      tasksToRun.unshift(cachedTask);
-      j -= getCachedTaskSize(cachedTask);
+  let remainingSize = TOTAL_CACHED_TASK_BYTE_LIMIT;
+
+  for (const { task, size } of availableTasks) {
+    if (size <= remainingSize) {
+      tasksToRun.push(task);
+      remainingSize -= size;
     }
   }
 
@@ -571,7 +690,6 @@ async function findRun(prisma: PrismaClientOrTransaction, id: string) {
       endpoint: true,
       organization: true,
       externalAccount: true,
-      queue: true,
       runConnections: {
         include: {
           integration: true,
@@ -588,6 +706,14 @@ async function findRun(prisma: PrismaClientOrTransaction, id: string) {
             in: ["COMPLETED"],
           },
         },
+        select: {
+          id: true,
+          idempotencyKey: true,
+          status: true,
+          noop: true,
+          output: true,
+          parentId: true,
+        },
       },
       event: true,
       version: {

apps/webapp/app/services/worker.server.ts

@@ -292,7 +292,12 @@ function getExecutionWorkerQueue() {
         handler: async (payload, job) => {
           const service = new PerformRunExecutionV2Service();
 
-          await service.call(payload.id, payload.reason, payload.isRetry, payload.resumeTaskId);
+          await service.call({
+            id: payload.id,
+            reason: payload.reason,
+            resumeTaskId: payload.resumeTaskId,
+            isRetry: payload.isRetry,
+          });
         },
       },
     },

packages/database/prisma/migrations/20230904145326_add_execution_columns_to_runs/migration.sql

@@ -0,0 +1,3 @@
+-- AlterTable
+ALTER TABLE "JobRun" ADD COLUMN     "executionCount" INTEGER NOT NULL DEFAULT 0,
+ADD COLUMN     "executionDuration" INTEGER NOT NULL DEFAULT 0;

packages/database/prisma/migrations/20230904205457_add_max_run_execution_time_to_orgs/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "Organization" ADD COLUMN     "maximumExecutionTimePerRunInMs" INTEGER NOT NULL DEFAULT 900000;

packages/database/prisma/schema.prisma

@@ -63,6 +63,8 @@ model Organization {
   slug  String @unique
   title String
 
+  maximumExecutionTimePerRunInMs Int @default(900000) // 15 minutes
+
   createdAt DateTime @default(now())
   updatedAt DateTime @updatedAt
 
@@ -704,6 +706,9 @@ model JobRun {
   timedOutAt     DateTime?
   timedOutReason String?
 
+  executionCount    Int @default(0)
+  executionDuration Int @default(0)
+
   isTest     Boolean @default(false)
   preprocess Boolean @default(false)
 

perf/src/index.ts

@@ -77,7 +77,25 @@ async function mainLong() {
   }
 }
 
-mainLong().catch((err) => {
+async function mainSerial() {
+  console.log("Preparing serial perf tests...");
+
+  // wait for 10 seconds
+  await new Promise((resolve) => setTimeout(resolve, 10000));
+
+  console.log("Starting serial perf tests in 1 second...");
+
+  // wait for 1 seconds
+  await new Promise((resolve) => setTimeout(resolve, 1000));
+
+  // Send 25 events
+  for (let i = 0; i < 25; i++) {
+    await sendEvent();
+    await new Promise((resolve) => setTimeout(resolve, 100));
+  }
+}
+
+mainSerial().catch((err) => {
   console.error(err);
   process.exit(1);
 });

perf/src/trigger.ts

@@ -1,5 +1,4 @@
 import { TriggerClient, eventTrigger } from "@trigger.dev/sdk";
-import { z } from "zod";
 
 export const triggerClient = new TriggerClient({
   id: "perf",
@@ -7,31 +6,26 @@ export const triggerClient = new TriggerClient({
   apiUrl: process.env.TRIGGER_API_URL!,
 });
 
-// Define 10 jobs in a for loop
-for (let i = 0; i < 10; i++) {
-  triggerClient.defineJob({
-    id: `perf-test-${i + 1}`,
-    name: `Perf Test ${i + 1}`,
-    version: "1.0.0",
-    trigger: eventTrigger({
-      name: "perf.test",
-    }),
-    queue: {
-      name: "perf-test",
-      maxConcurrent: 50,
-    },
-    run: async (payload, io, ctx) => {
-      await io.runTask("task-1", { name: "task 1" }, async (task) => {
-        return {
-          value: Math.random(),
-        };
-      });
+triggerClient.defineJob({
+  id: `perf-test-1`,
+  name: `Perf Test 1`,
+  version: "1.0.0",
+  trigger: eventTrigger({
+    name: "perf.test",
+  }),
+  run: async (payload, io, ctx) => {
+    await io.runTask("task-1", { name: "task 1" }, async (task) => {
+      await new Promise((resolve) => setTimeout(resolve, 2000));
 
-      await io.runTask("task-2", { name: "task 2" }, async (task) => {
-        return {
-          value: Math.random(),
-        };
-      });
-    },
-  });
-}
+      return {
+        value: Math.random(),
+      };
+    });
+
+    await io.runTask("task-2", { name: "task 2" }, async (task) => {
+      return {
+        value: Math.random(),
+      };
+    });
+  },
+});