Job run performance improvements and adding "worker only" mode (#360)

* WIP job run performance improvements

- Added a `perf` tool to better measure job run performance under heavy load
- Removed `runFinished` job (not really needed)
- startQueuedRuns now uses a jobKey with replace
- Fixed an issue with ZodWorker when using jobKey

* Publish improvement docker images

* fixed the improvement docker publishing

* Downgrade back to prisma 4.16.0 because 5.1.x broke docker builds

* Changes to how queued runs work

- Split the worker into two different workers, one dedicated to performRunExecution
- Schedule performRunExecution in a single place, with a queue and using a round robin manually controlled concurrency
- Remove startQueuedRuns
- All runs are queued before they are started
- Setting the worker maxPoolSize to the same as the worker concurrency
- Starting to be able to split the docker image

* Remove queue name from startRun graphile job

* Make the prisma connection pool stuff configurable through env vars

* Hardcode (for now) the max concurrent runs limit

* Rewrite performRunExecution to be more performant

PerformRunExecutionV2:

- Does not create and manage jobRunExecution records
- Does not reimplement retrying, uses graphile worker retrying instead

I’ve kept around PerformRunExecutionV1 so this works when deploying. Definitely needs LOTS of testing

* Fix issues with cached tasks

- Limit the size of the cached tasks sent when executing a run, using the knapsack problem dynamic programming approach
- Actually USE the cached tasks in IO by using the idempotencyKey instead of the task ID
- Remove output from all logs
- Added a stress test job catalog

* Forgot to commit the logger updates

* Never log connectionString

* Login to docker hub to get around rate limits

* Add additional logging to the graphile workers

* Fix the *_ENABLED  env vars

* Allow adding and removing jobs to be done from the webapp

* Don’t set the job to failed if it’s being retried

* Deprecated queue options in the job and removed startPosition. Now using the job/env combo as the job queue name

* Dequeung jobs doesn’t check if the runner is initialized

* Fixed issues with retrying a run getting stuck on a cancelled task, and errors from parsing the results of dequeing a job

* Remove queued round robin thing that isn’t used anymore

* Added slack to job catalog

* Better forwards compat

* Added long delay

* Fixed lock file

Author: ericallam

.changeset/fuzzy-trees-lay.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/sdk": patch
+---
+
+Fixed IO not setting the cached task key correctly, resulting in unnecessary API calls to trigger.dev

.changeset/happy-foxes-play.md

@@ -0,0 +1,6 @@
+---
+"@trigger.dev/sdk": patch
+"@trigger.dev/core": patch
+---
+
+Deprecated queue options in the job and removed startPosition

.github/workflows/publish.yml

@@ -4,6 +4,7 @@ on:
   push:
     branches:
       - main
+      - improvements/*
     tags:
       - "v.docker.*"
     paths:
@@ -95,6 +96,12 @@ jobs:
     name: e2e Tests
     runs-on: buildjet-4vcpu-ubuntu-2204
     steps:
+      - name: 🐳 Login to Docker Hub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+
       - name: ⬇️ Checkout repo
         uses: actions/checkout@v3
         with:
@@ -154,6 +161,11 @@ jobs:
       version: ${{ steps.get_version.outputs.version }}
       short_sha: ${{ steps.get_commit.outputs.sha_short }}
     steps:
+      - name: 🐳 Login to Docker Hub
+        uses: docker/login-action@v2
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
       - name: ⬇️ Checkout repo
         uses: actions/checkout@v3
 
@@ -167,6 +179,10 @@ jobs:
               IMAGE_TAG="v${ORIGINAL_VERSION}"
             fi
             echo "IMAGE_TAG=${IMAGE_TAG}"
+          elif [[ $GITHUB_REF == refs/heads/improvements/* ]]; then
+            ORIGINAL_VERSION="${GITHUB_REF#refs/heads/improvements/}"
+            IMAGE_TAG="${ORIGINAL_VERSION}.rc"
+            echo "IMAGE_TAG=${IMAGE_TAG}"
           elif [[ $GITHUB_REF == refs/heads/* ]]; then
             IMAGE_TAG="${GITHUB_REF#refs/heads/}"
             echo "IMAGE_TAG=${IMAGE_TAG}"

apps/webapp/app/consts.ts

@@ -1,6 +1,7 @@
 export const LIVE_ENVIRONMENT = "live";
 export const DEV_ENVIRONMENT = "development";
 export const MAX_LIVE_PROJECTS = 1;
-export const DEFAULT_MAX_CONCURRENT_RUNS = 10000;
+export const DEFAULT_MAX_CONCURRENT_RUNS = 10;
+export const MAX_CONCURRENT_RUNS_LIMIT = 20;
 export const PREPROCESS_RETRY_LIMIT = 2;
 export const EXECUTE_JOB_RETRY_LIMIT = 10;

apps/webapp/app/db.server.ts

@@ -2,6 +2,7 @@ import { PrismaClient, Prisma } from "@trigger.dev/database";
 import invariant from "tiny-invariant";
 import { z } from "zod";
 import { logger } from "./services/logger.server";
+import { env } from "./env.server";
 
 export type PrismaTransactionClient = Omit<
   PrismaClient,
@@ -84,17 +85,24 @@ function getClient() {
   const { DATABASE_URL } = process.env;
   invariant(typeof DATABASE_URL === "string", "DATABASE_URL env var not set");
 
+  const databaseUrl = new URL(DATABASE_URL);
+
+  // We need to add the connection_limit and pool_timeout query params to the url, in a way that works if the DATABASE_URL already has query params
+  const query = databaseUrl.searchParams;
+  query.set("connection_limit", env.DATABASE_CONNECTION_LIMIT.toString());
+  query.set("pool_timeout", env.DATABASE_POOL_TIMEOUT.toString());
+  databaseUrl.search = query.toString();
+
   // Remove the username:password in the url and print that to the console
-  const urlWithoutCredentials = new URL(DATABASE_URL);
+  const urlWithoutCredentials = new URL(databaseUrl.href);
   urlWithoutCredentials.password = "";
 
   console.log(`🔌 setting up prisma client to ${urlWithoutCredentials.toString()}`);
 
   const client = new PrismaClient({
     datasources: {
       db: {
-        url: DATABASE_URL,
-        // We can't set directUrl here, and we don't have to
+        url: databaseUrl.href,
       },
     },
     log: [

apps/webapp/app/env.server.ts

@@ -4,6 +4,8 @@ import { SecretStoreOptionsSchema } from "./services/secrets/secretStore.server"
 const EnvironmentSchema = z.object({
   NODE_ENV: z.union([z.literal("development"), z.literal("production"), z.literal("test")]),
   DATABASE_URL: z.string(),
+  DATABASE_CONNECTION_LIMIT: z.coerce.number().int().default(10),
+  DATABASE_POOL_TIMEOUT: z.coerce.number().int().default(60),
   DIRECT_URL: z.string(),
   SESSION_SECRET: z.string(),
   MAGIC_LINK_SECRET: z.string(),
@@ -31,6 +33,13 @@ const EnvironmentSchema = z.object({
   RESEND_API_KEY: z.string().optional(),
   PLAIN_API_KEY: z.string().optional(),
   RUNTIME_PLATFORM: z.enum(["docker-compose", "ecs", "local"]).default("local"),
+  WORKER_SCHEMA: z.string().default("graphile_worker"),
+  WORKER_CONCURRENCY: z.coerce.number().int().default(10),
+  WORKER_POLL_INTERVAL: z.coerce.number().int().default(1000),
+  EXECUTION_WORKER_CONCURRENCY: z.coerce.number().int().default(10),
+  EXECUTION_WORKER_POLL_INTERVAL: z.coerce.number().int().default(1000),
+  WORKER_ENABLED: z.string().default("true"),
+  EXECUTION_WORKER_ENABLED: z.string().default("true"),
 });
 
 export type Environment = z.infer<typeof EnvironmentSchema>;

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

@@ -0,0 +1,58 @@
+import { JobRun, JobRunExecution } from "@trigger.dev/database";
+import { PrismaClientOrTransaction } from "~/db.server";
+import { executionWorker } from "~/services/worker.server";
+
+export async function enqueueRunExecutionV1(
+  execution: JobRunExecution,
+  queueId: string,
+  concurrency: number,
+  tx: PrismaClientOrTransaction,
+  runAt?: Date
+) {
+  const job = await executionWorker.enqueue(
+    "performRunExecution",
+    {
+      id: execution.id,
+    },
+    {
+      queueName: `job:queue:${queueId}`,
+      tx,
+      runAt,
+      jobKey: `execution:${execution.runId}`,
+    }
+  );
+}
+
+export type EnqueueRunExecutionV2Options = {
+  runAt?: Date;
+  resumeTaskId?: string;
+  isRetry?: boolean;
+};
+
+export async function enqueueRunExecutionV2(
+  run: JobRun,
+  tx: PrismaClientOrTransaction,
+  options: EnqueueRunExecutionV2Options = {}
+) {
+  const job = await executionWorker.enqueue(
+    "performRunExecutionV2",
+    {
+      id: run.id,
+      reason: run.status === "PREPROCESSING" ? "PREPROCESS" : "EXECUTE_JOB",
+      resumeTaskId: options.resumeTaskId,
+      isRetry: typeof options.isRetry === "boolean" ? options.isRetry : false,
+    },
+    {
+      queueName: `job:${run.jobId}:env:${run.environmentId}`,
+      tx,
+      runAt: options.runAt,
+      jobKey: `job_run:${run.id}`,
+    }
+  );
+}
+
+export async function dequeueRunExecutionV2(run: JobRun, tx: PrismaClientOrTransaction) {
+  return await executionWorker.dequeue(`job_run:${run.id}`, {
+    tx,
+  });
+}

apps/webapp/app/platform/zodWorker.server.ts

@@ -51,6 +51,7 @@ const AddJobResultsSchema = z.array(GraphileJobSchema);
 export type ZodTasks<TConsumerSchema extends MessageCatalogSchema> = {
   [K in keyof TConsumerSchema]: {
     queueName?: string | ((payload: z.infer<TConsumerSchema[K]>) => string);
+    jobKey?: string | ((payload: z.infer<TConsumerSchema[K]>) => string | undefined);
     priority?: number;
     maxAttempts?: number;
     jobKeyMode?: "replace" | "preserve_run_at" | "unsafe_dedupe";
@@ -76,7 +77,12 @@ export type ZodWorkerEnqueueOptions = TaskSpec & {
   tx?: PrismaClientOrTransaction;
 };
 
+export type ZodWorkerDequeueOptions = {
+  tx?: PrismaClientOrTransaction;
+};
+
 export type ZodWorkerOptions<TMessageCatalog extends MessageCatalogSchema> = {
+  name: string;
   runnerOptions: RunnerOptions;
   prisma: PrismaClient;
   schema: TMessageCatalog;
@@ -85,6 +91,7 @@ export type ZodWorkerOptions<TMessageCatalog extends MessageCatalogSchema> = {
 };
 
 export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
+  #name: string;
   #schema: TMessageCatalog;
   #prisma: PrismaClient;
   #runnerOptions: RunnerOptions;
@@ -93,6 +100,7 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
   #runner?: GraphileRunner;
 
   constructor(options: ZodWorkerOptions<TMessageCatalog>) {
+    this.#name = options.name;
     this.#schema = options.schema;
     this.#prisma = options.prisma;
     this.#runnerOptions = options.runnerOptions;
@@ -105,7 +113,7 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
       return true;
     }
 
-    logger.debug("Initializing worker queue with options", {
+    this.#logDebug("Initializing worker queue with options", {
       runnerOptions: this.#runnerOptions,
     });
 
@@ -121,6 +129,54 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
       throw new Error("Failed to initialize worker queue");
     }
 
+    this.#runner?.events.on("pool:create", ({ workerPool }) => {
+      this.#logDebug("pool:create");
+    });
+
+    this.#runner?.events.on("pool:listen:connecting", ({ workerPool, attempts }) => {
+      this.#logDebug("pool:create", { attempts });
+    });
+
+    this.#runner?.events.on("pool:listen:success", ({ workerPool, client }) => {
+      this.#logDebug("pool:listen:success");
+    });
+
+    this.#runner?.events.on("pool:listen:error", ({ error }) => {
+      this.#logDebug("pool:listen:error", { error });
+    });
+
+    this.#runner?.events.on("pool:gracefulShutdown", ({ message }) => {
+      this.#logDebug("pool:gracefulShutdown", { workerMessage: message });
+    });
+
+    this.#runner?.events.on("pool:gracefulShutdown:error", ({ error }) => {
+      this.#logDebug("pool:gracefulShutdown:error", { error });
+    });
+
+    this.#runner?.events.on("worker:create", ({ worker }) => {
+      this.#logDebug("worker:create", { workerId: worker.workerId });
+    });
+
+    this.#runner?.events.on("worker:release", ({ worker }) => {
+      this.#logDebug("worker:release", { workerId: worker.workerId });
+    });
+
+    this.#runner?.events.on("worker:stop", ({ worker, error }) => {
+      this.#logDebug("worker:stop", { workerId: worker.workerId, error });
+    });
+
+    this.#runner?.events.on("worker:fatalError", ({ worker, error, jobError }) => {
+      this.#logDebug("worker:fatalError", { workerId: worker.workerId, error, jobError });
+    });
+
+    this.#runner?.events.on("gracefulShutdown", ({ signal }) => {
+      this.#logDebug("gracefulShutdown", { signal });
+    });
+
+    this.#runner?.events.on("stop", () => {
+      this.#logDebug("stop");
+    });
+
     return true;
   }
 
@@ -133,23 +189,34 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
     payload: z.infer<TMessageCatalog[K]>,
     options?: ZodWorkerEnqueueOptions
   ): Promise<GraphileJob> {
-    if (!this.#runner) {
-      throw new Error("Worker not initialized");
-    }
-
     const task = this.#tasks[identifier];
 
     const optionsWithoutTx = omit(options ?? {}, ["tx"]);
+    const taskWithoutJobKey = omit(task, ["jobKey"]);
 
     const spec = {
       ...optionsWithoutTx,
-      ...task,
+      ...taskWithoutJobKey,
     };
 
     if (typeof task.queueName === "function") {
       spec.queueName = task.queueName(payload);
     }
 
+    if (typeof task.jobKey === "function") {
+      const jobKey = task.jobKey(payload);
+
+      if (jobKey) {
+        spec.jobKey = jobKey;
+      }
+    }
+
+    logger.debug("Enqueuing worker task", {
+      identifier,
+      payload,
+      spec,
+    });
+
     const job = await this.#addJob(
       identifier as string,
       payload,
@@ -167,6 +234,17 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
     return job;
   }
 
+  public async dequeue(
+    jobKey: string,
+    option?: ZodWorkerDequeueOptions
+  ): Promise<GraphileJob | undefined> {
+    const results = await this.#removeJob(jobKey, option?.tx ?? this.#prisma);
+
+    logger.debug("dequeued worker task", { results, jobKey });
+
+    return results;
+  }
+
   async #addJob(
     identifier: string,
     payload: unknown,
@@ -192,8 +270,8 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
       spec.maxAttempts || null,
       spec.jobKey || null,
       spec.priority || null,
-      spec.jobKeyMode || null,
-      spec.flags || null
+      spec.flags || null,
+      spec.jobKeyMode || null
     );
 
     const rows = AddJobResultsSchema.safeParse(results);
@@ -209,6 +287,32 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
     return job as GraphileJob;
   }
 
+  async #removeJob(jobKey: string, tx: PrismaClientOrTransaction) {
+    try {
+      const result = await tx.$queryRawUnsafe(
+        `SELECT * FROM graphile_worker.remove_job(
+          job_key => $1::text
+        )`,
+        jobKey
+      );
+      const job = AddJobResultsSchema.safeParse(result);
+
+      if (!job.success) {
+        logger.debug("results returned from remove_job could not be parsed", {
+          error: job.error.flatten(),
+          result,
+          jobKey,
+        });
+
+        return;
+      }
+
+      return job.data[0] as GraphileJob;
+    } catch (e) {
+      throw new Error(`Failed to remove job from queue, ${e}}`);
+    }
+  }
+
   #createTaskListFromTasks() {
     const taskList: TaskList = {};
 
@@ -324,4 +428,8 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
       throw error;
     }
   }
+
+  #logDebug(message: string, args?: any) {
+    logger.debug(`[worker][${this.#name}] ${message}`, args);
+  }
 }