move the release concurrency queue into the release concurrency system, make it disabled by default, configure the run engine in the webapp with env vars

ericallam · ericallam · commit 6eaba5fa5614 · 2025-03-19T16:51:04.000Z
diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts
@@ -565,6 +565,13 @@ const EnvironmentSchema = z.object({
   RUN_ENGINE_RATE_LIMIT_REJECTION_LOGS_ENABLED: z.string().default("1"),
   RUN_ENGINE_RATE_LIMIT_LIMITER_LOGS_ENABLED: z.string().default("0"),
 
+  RUN_ENGINE_RELEASE_CONCURRENCY_ENABLED: z.string().default("0"),
+  RUN_ENGINE_RELEASE_CONCURRENCY_MAX_TOKENS_RATIO: z.coerce.number().default(1),
+  RUN_ENGINE_RELEASE_CONCURRENCY_MAX_RETRIES: z.coerce.number().int().default(3),
+  RUN_ENGINE_RELEASE_CONCURRENCY_CONSUMERS_COUNT: z.coerce.number().int().default(1),
+  RUN_ENGINE_RELEASE_CONCURRENCY_POLL_INTERVAL: z.coerce.number().int().default(500),
+  RUN_ENGINE_RELEASE_CONCURRENCY_BATCH_SIZE: z.coerce.number().int().default(10),
+
   /** How long should the presence ttl last */
   DEV_PRESENCE_TTL_MS: z.coerce.number().int().default(30_000),
   DEV_PRESENCE_POLL_INTERVAL_MS: z.coerce.number().int().default(5_000),
diff --git a/apps/webapp/app/v3/runEngine.server.ts b/apps/webapp/app/v3/runEngine.server.ts
@@ -1,10 +1,10 @@
 import { RunEngine } from "@internal/run-engine";
+import { defaultMachine } from "@trigger.dev/platform/v3";
 import { prisma } from "~/db.server";
 import { env } from "~/env.server";
-import { tracer } from "./tracer.server";
 import { singleton } from "~/utils/singleton";
-import { defaultMachine, machines } from "@trigger.dev/platform/v3";
 import { allMachines } from "./machinePresets.server";
+import { tracer } from "./tracer.server";
 
 export const engine = singleton("RunEngine", createRunEngine);
 
@@ -73,6 +73,23 @@ function createRunEngine() {
       EXECUTING: env.RUN_ENGINE_TIMEOUT_EXECUTING,
       EXECUTING_WITH_WAITPOINTS: env.RUN_ENGINE_TIMEOUT_EXECUTING_WITH_WAITPOINTS,
     },
+    releaseConcurrency: {
+      disabled: env.RUN_ENGINE_RELEASE_CONCURRENCY_ENABLED === "0",
+      maxTokensRatio: env.RUN_ENGINE_RELEASE_CONCURRENCY_MAX_TOKENS_RATIO,
+      maxRetries: env.RUN_ENGINE_RELEASE_CONCURRENCY_MAX_RETRIES,
+      consumersCount: env.RUN_ENGINE_RELEASE_CONCURRENCY_CONSUMERS_COUNT,
+      pollInterval: env.RUN_ENGINE_RELEASE_CONCURRENCY_POLL_INTERVAL,
+      batchSize: env.RUN_ENGINE_RELEASE_CONCURRENCY_BATCH_SIZE,
+      redis: {
+        keyPrefix: "engine:",
+        port: env.RUN_ENGINE_RUN_QUEUE_REDIS_PORT ?? undefined,
+        host: env.RUN_ENGINE_RUN_QUEUE_REDIS_HOST ?? undefined,
+        username: env.RUN_ENGINE_RUN_QUEUE_REDIS_USERNAME ?? undefined,
+        password: env.RUN_ENGINE_RUN_QUEUE_REDIS_PASSWORD ?? undefined,
+        enableAutoPipelining: true,
+        ...(env.RUN_ENGINE_RUN_QUEUE_REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
+      },
+    },
   });
 
   return engine;
diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts
@@ -35,7 +35,6 @@ import {
 } from "./errors.js";
 import { EventBus, EventBusEvents } from "./eventBus.js";
 import { RunLocker } from "./locking.js";
-import { ReleaseConcurrencyTokenBucketQueue } from "./releaseConcurrencyTokenBucketQueue.js";
 import { BatchSystem } from "./systems/batchSystem.js";
 import { CheckpointSystem } from "./systems/checkpointSystem.js";
 import { DelayedRunSystem } from "./systems/delayedRunSystem.js";
@@ -63,11 +62,6 @@ export class RunEngine {
   private logger = new Logger("RunEngine", "debug");
   private tracer: Tracer;
   private heartbeatTimeouts: HeartbeatTimeouts;
-  releaseConcurrencyQueue: ReleaseConcurrencyTokenBucketQueue<{
-    orgId: string;
-    projectId: string;
-    envId: string;
-  }>;
   eventBus: EventBus = new EventEmitter<EventBusEvents>();
   executionSnapshotSystem: ExecutionSnapshotSystem;
   runAttemptSystem: RunAttemptSystem;
@@ -184,51 +178,6 @@ export class RunEngine {
       ...(options.heartbeatTimeoutsMs ?? {}),
     };
 
-    // Initialize the ReleaseConcurrencyQueue
-    this.releaseConcurrencyQueue = new ReleaseConcurrencyTokenBucketQueue({
-      redis: {
-        ...options.queue.redis, // Use base queue redis options
-        ...options.releaseConcurrency?.redis, // Allow overrides
-        keyPrefix: `${options.queue.redis.keyPrefix ?? ""}release-concurrency:`,
-      },
-      retry: {
-        maxRetries: options.releaseConcurrency?.maxRetries ?? 5,
-        backoff: {
-          minDelay: options.releaseConcurrency?.backoff?.minDelay ?? 1000,
-          maxDelay: options.releaseConcurrency?.backoff?.maxDelay ?? 10000,
-          factor: options.releaseConcurrency?.backoff?.factor ?? 2,
-        },
-      },
-      consumersCount: options.releaseConcurrency?.consumersCount ?? 1,
-      pollInterval: options.releaseConcurrency?.pollInterval ?? 1000,
-      batchSize: options.releaseConcurrency?.batchSize ?? 10,
-      executor: async (descriptor, snapshotId) => {
-        await this.releaseConcurrencySystem.executeReleaseConcurrencyForSnapshot(snapshotId);
-      },
-      maxTokens: async (descriptor) => {
-        const environment = await this.prisma.runtimeEnvironment.findFirstOrThrow({
-          where: { id: descriptor.envId },
-          select: {
-            maximumConcurrencyLimit: true,
-          },
-        });
-
-        return (
-          environment.maximumConcurrencyLimit * (options.releaseConcurrency?.maxTokensRatio ?? 1.0)
-        );
-      },
-      keys: {
-        fromDescriptor: (descriptor) =>
-          `org:${descriptor.orgId}:proj:${descriptor.projectId}:env:${descriptor.envId}`,
-        toDescriptor: (name) => ({
-          orgId: name.split(":")[1],
-          projectId: name.split(":")[3],
-          envId: name.split(":")[5],
-        }),
-      },
-      tracer: this.tracer,
-    });
-
     const resources: SystemResources = {
       prisma: this.prisma,
       worker: this.worker,
@@ -237,11 +186,60 @@ export class RunEngine {
       tracer: this.tracer,
       runLock: this.runLock,
       runQueue: this.runQueue,
-      releaseConcurrencyQueue: this.releaseConcurrencyQueue,
     };
 
     this.releaseConcurrencySystem = new ReleaseConcurrencySystem({
       resources,
+      queueOptions:
+        typeof options.releaseConcurrency?.disabled === "boolean" &&
+        options.releaseConcurrency.disabled
+          ? undefined
+          : {
+              redis: {
+                ...options.queue.redis, // Use base queue redis options
+                ...options.releaseConcurrency?.redis, // Allow overrides
+                keyPrefix: `${options.queue.redis.keyPrefix ?? ""}release-concurrency:`,
+              },
+              retry: {
+                maxRetries: options.releaseConcurrency?.maxRetries ?? 5,
+                backoff: {
+                  minDelay: options.releaseConcurrency?.backoff?.minDelay ?? 1000,
+                  maxDelay: options.releaseConcurrency?.backoff?.maxDelay ?? 10000,
+                  factor: options.releaseConcurrency?.backoff?.factor ?? 2,
+                },
+              },
+              consumersCount: options.releaseConcurrency?.consumersCount ?? 1,
+              pollInterval: options.releaseConcurrency?.pollInterval ?? 1000,
+              batchSize: options.releaseConcurrency?.batchSize ?? 10,
+              executor: async (descriptor, snapshotId) => {
+                await this.releaseConcurrencySystem.executeReleaseConcurrencyForSnapshot(
+                  snapshotId
+                );
+              },
+              maxTokens: async (descriptor) => {
+                const environment = await this.prisma.runtimeEnvironment.findFirstOrThrow({
+                  where: { id: descriptor.envId },
+                  select: {
+                    maximumConcurrencyLimit: true,
+                  },
+                });
+
+                return (
+                  environment.maximumConcurrencyLimit *
+                  (options.releaseConcurrency?.maxTokensRatio ?? 1.0)
+                );
+              },
+              keys: {
+                fromDescriptor: (descriptor) =>
+                  `org:${descriptor.orgId}:proj:${descriptor.projectId}:env:${descriptor.envId}`,
+                toDescriptor: (name) => ({
+                  orgId: name.split(":")[1],
+                  projectId: name.split(":")[3],
+                  envId: name.split(":")[5],
+                }),
+              },
+              tracer: this.tracer,
+            },
     });
 
     this.executionSnapshotSystem = new ExecutionSnapshotSystem({
@@ -1201,7 +1199,7 @@ export class RunEngine {
   async quit() {
     try {
       //stop the run queue
-      await this.releaseConcurrencyQueue.quit();
+      await this.releaseConcurrencySystem.quit();
       await this.runQueue.quit();
       await this.worker.stop();
       await this.runLock.quit();
diff --git a/internal-packages/run-engine/src/engine/systems/releaseConcurrencySystem.ts b/internal-packages/run-engine/src/engine/systems/releaseConcurrencySystem.ts
@@ -3,6 +3,10 @@ import { SystemResources } from "./systems.js";
 import { getLatestExecutionSnapshot } from "./executionSnapshotSystem.js";
 import { canReleaseConcurrency } from "../statuses.js";
 import { z } from "zod";
+import {
+  ReleaseConcurrencyQueueOptions,
+  ReleaseConcurrencyTokenBucketQueue,
+} from "../releaseConcurrencyTokenBucketQueue.js";
 
 const ReleaseConcurrencyMetadata = z.object({
   releaseConcurrency: z.boolean().optional(),
@@ -12,17 +16,65 @@ type ReleaseConcurrencyMetadata = z.infer<typeof ReleaseConcurrencyMetadata>;
 
 export type ReleaseConcurrencySystemOptions = {
   resources: SystemResources;
+  queueOptions?: ReleaseConcurrencyQueueOptions<{
+    orgId: string;
+    projectId: string;
+    envId: string;
+  }>;
 };
 
 export class ReleaseConcurrencySystem {
   private readonly $: SystemResources;
+  releaseConcurrencyQueue?: ReleaseConcurrencyTokenBucketQueue<{
+    orgId: string;
+    projectId: string;
+    envId: string;
+  }>;
 
   constructor(private readonly options: ReleaseConcurrencySystemOptions) {
     this.$ = options.resources;
+
+    if (options.queueOptions) {
+      this.releaseConcurrencyQueue = new ReleaseConcurrencyTokenBucketQueue(options.queueOptions);
+    }
+  }
+
+  public async consumeToken(
+    descriptor: { orgId: string; projectId: string; envId: string },
+    releaserId: string
+  ) {
+    if (!this.releaseConcurrencyQueue) {
+      return;
+    }
+
+    await this.releaseConcurrencyQueue.consumeToken(descriptor, releaserId);
+  }
+
+  public async returnToken(
+    descriptor: { orgId: string; projectId: string; envId: string },
+    releaserId: string
+  ) {
+    if (!this.releaseConcurrencyQueue) {
+      return;
+    }
+
+    await this.releaseConcurrencyQueue.returnToken(descriptor, releaserId);
+  }
+
+  public async quit() {
+    if (!this.releaseConcurrencyQueue) {
+      return;
+    }
+
+    await this.releaseConcurrencyQueue.quit();
   }
 
   public async checkpointCreatedOnEnvironment(environment: RuntimeEnvironment) {
-    await this.$.releaseConcurrencyQueue.refillTokens(
+    if (!this.releaseConcurrencyQueue) {
+      return;
+    }
+
+    await this.releaseConcurrencyQueue.refillTokens(
       {
         orgId: environment.organizationId,
         projectId: environment.projectId,
@@ -33,12 +85,16 @@ export class ReleaseConcurrencySystem {
   }
 
   public async releaseConcurrencyForSnapshot(snapshot: TaskRunExecutionSnapshot) {
+    if (!this.releaseConcurrencyQueue) {
+      return;
+    }
+
     // Go ahead and release concurrency immediately if the run is in a development environment
     if (snapshot.environmentType === "DEVELOPMENT") {
       return await this.executeReleaseConcurrencyForSnapshot(snapshot.id);
     }
 
-    await this.$.releaseConcurrencyQueue.attemptToRelease(
+    await this.releaseConcurrencyQueue.attemptToRelease(
       {
         orgId: snapshot.organizationId,
         projectId: snapshot.projectId,
@@ -49,6 +105,10 @@ export class ReleaseConcurrencySystem {
   }
 
   public async executeReleaseConcurrencyForSnapshot(snapshotId: string) {
+    if (!this.releaseConcurrencyQueue) {
+      return;
+    }
+
     this.$.logger.debug("Executing released concurrency", {
       snapshotId,
     });
diff --git a/internal-packages/run-engine/src/engine/systems/systems.ts b/internal-packages/run-engine/src/engine/systems/systems.ts
@@ -5,7 +5,6 @@ import { RunQueue } from "../../run-queue/index.js";
 import { EventBus } from "../eventBus.js";
 import { RunLocker } from "../locking.js";
 import { EngineWorker } from "../types.js";
-import { ReleaseConcurrencyTokenBucketQueue } from "../releaseConcurrencyTokenBucketQueue.js";
 
 export type SystemResources = {
   prisma: PrismaClient;
@@ -15,9 +14,4 @@ export type SystemResources = {
   tracer: Tracer;
   runLock: RunLocker;
   runQueue: RunQueue;
-  releaseConcurrencyQueue: ReleaseConcurrencyTokenBucketQueue<{
-    orgId: string;
-    projectId: string;
-    envId: string;
-  }>;
 };
diff --git a/internal-packages/run-engine/src/engine/tests/releaseConcurrency.test.ts b/internal-packages/run-engine/src/engine/tests/releaseConcurrency.test.ts
@@ -673,7 +673,7 @@ describe("RunEngine Releasing Concurrency", () => {
         projectId: authenticatedEnvironment.projectId,
       });
 
-      await engine.releaseConcurrencyQueue.consumeToken(
+      await engine.releaseConcurrencySystem.consumeToken(
         {
           orgId: authenticatedEnvironment.organizationId,
           projectId: authenticatedEnvironment.projectId,
@@ -708,7 +708,7 @@ describe("RunEngine Releasing Concurrency", () => {
       expect(envConcurrencyAfter).toBe(1);
 
       // Now we return the token so the concurrency can be released
-      await engine.releaseConcurrencyQueue.returnToken(
+      await engine.releaseConcurrencySystem.returnToken(
         {
           orgId: authenticatedEnvironment.organizationId,
           projectId: authenticatedEnvironment.projectId,
@@ -835,7 +835,7 @@ describe("RunEngine Releasing Concurrency", () => {
         projectId: authenticatedEnvironment.projectId,
       });
 
-      await engine.releaseConcurrencyQueue.consumeToken(
+      await engine.releaseConcurrencySystem.consumeToken(
         {
           orgId: authenticatedEnvironment.organizationId,
           projectId: authenticatedEnvironment.projectId,
@@ -890,7 +890,7 @@ describe("RunEngine Releasing Concurrency", () => {
       expect(snapshot.executionStatus).toBe("SUSPENDED");
 
       // Now we return the token so the concurrency can be released
-      await engine.releaseConcurrencyQueue.returnToken(
+      await engine.releaseConcurrencySystem.returnToken(
         {
           orgId: authenticatedEnvironment.organizationId,
           projectId: authenticatedEnvironment.projectId,
@@ -1017,7 +1017,7 @@ describe("RunEngine Releasing Concurrency", () => {
         projectId: authenticatedEnvironment.projectId,
       });
 
-      await engine.releaseConcurrencyQueue.consumeToken(
+      await engine.releaseConcurrencySystem.consumeToken(
         {
           orgId: authenticatedEnvironment.organizationId,
           projectId: authenticatedEnvironment.projectId,
@@ -1063,7 +1063,7 @@ describe("RunEngine Releasing Concurrency", () => {
       expect(executionDataAfter?.snapshot.executionStatus).toBe("EXECUTING");
 
       // Now we return the token so the concurrency can be released
-      await engine.releaseConcurrencyQueue.returnToken(
+      await engine.releaseConcurrencySystem.returnToken(
         {
           orgId: authenticatedEnvironment.organizationId,
           projectId: authenticatedEnvironment.projectId,
diff --git a/internal-packages/run-engine/src/engine/types.ts b/internal-packages/run-engine/src/engine/types.ts
@@ -37,6 +37,7 @@ export type RunEngineOptions = {
   queueRunsWaitingForWorkerBatchSize?: number;
   tracer: Tracer;
   releaseConcurrency?: {
+    disabled?: boolean;
     maxTokensRatio?: number;
     redis?: Partial<RedisOptions>;
     maxRetries?: number;
