Merge pull request #1804 from triggerdotdev/re2-reserve-concurrency

re2: New release concurrency system

Author: ericallam

.cursor/mcp.json

@@ -4,4 +4,4 @@
       "url": "http://localhost:3333/sse"
     }
   }
-}
+}

.vscode/launch.json

@@ -138,15 +138,15 @@
       "type": "node-terminal",
       "request": "launch",
       "name": "Debug RunEngine tests",
-      "command": "pnpm run test --filter @internal/run-engine",
-      "cwd": "${workspaceFolder}",
+      "command": "pnpm run test ./src/engine/tests/releaseConcurrencyQueue.test.ts -t 'Should manage token bucket and queue correctly'",
+      "cwd": "${workspaceFolder}/internal-packages/run-engine",
       "sourceMaps": true
     },
     {
       "type": "node-terminal",
       "request": "launch",
       "name": "Debug RunQueue tests",
-      "command": "pnpm run test ./src/engine/tests/waitpoints.test.ts",
+      "command": "pnpm run test ./src/run-queue/index.test.ts",
       "cwd": "${workspaceFolder}/internal-packages/run-engine",
       "sourceMaps": true
     }

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),

apps/webapp/app/routes/admin.api.v1.environments.$environmentId.ts

@@ -2,7 +2,7 @@ import { ActionFunctionArgs, json, LoaderFunctionArgs } from "@remix-run/server-
 import { z } from "zod";
 import { prisma } from "~/db.server";
 import { authenticateApiRequestWithPersonalAccessToken } from "~/services/personalAccessToken.server";
-import { marqs } from "~/v3/marqs/index.server";
+import { engine } from "~/v3/runEngine.server";
 import { updateEnvConcurrencyLimits } from "~/v3/runQueue.server";
 
 const ParamsSchema = z.object({
@@ -113,20 +113,15 @@ export async function loader({ request, params }: LoaderFunctionArgs) {
     Object.fromEntries(requestUrl.searchParams.entries())
   );
 
-  const concurrencyLimit = await marqs.getEnvConcurrencyLimit(environment);
-  const currentConcurrency = await marqs.currentConcurrencyOfEnvironment(environment);
-  const reserveConcurrency = await marqs.reserveConcurrencyOfEnvironment(environment);
+  const concurrencyLimit = await engine.runQueue.getEnvConcurrencyLimit(environment);
+  const currentConcurrency = await engine.runQueue.currentConcurrencyOfEnvironment(environment);
 
   if (searchParams.queue) {
-    const queueConcurrencyLimit = await marqs.getQueueConcurrencyLimit(
+    const queueConcurrencyLimit = await engine.runQueue.getQueueConcurrencyLimit(
       environment,
       searchParams.queue
     );
-    const queueCurrentConcurrency = await marqs.currentConcurrencyOfQueue(
-      environment,
-      searchParams.queue
-    );
-    const queueReserveConcurrency = await marqs.reserveConcurrencyOfQueue(
+    const queueCurrentConcurrency = await engine.runQueue.currentConcurrencyOfQueue(
       environment,
       searchParams.queue
     );
@@ -135,12 +130,10 @@ export async function loader({ request, params }: LoaderFunctionArgs) {
       id: environment.id,
       concurrencyLimit,
       currentConcurrency,
-      reserveConcurrency,
       queueConcurrencyLimit,
       queueCurrentConcurrency,
-      queueReserveConcurrency,
     });
   }
 
-  return json({ id: environment.id, concurrencyLimit, currentConcurrency, reserveConcurrency });
+  return json({ id: environment.id, concurrencyLimit, currentConcurrency });
 }

apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.wait.duration.ts

@@ -48,12 +48,9 @@ const { action } = createActionApiRoute(
       const waitResult = await engine.blockRunWithWaitpoint({
         runId: run.id,
         waitpoints: waitpoint.id,
-        environmentId: authentication.environment.id,
         projectId: authentication.environment.project.id,
         organizationId: authentication.environment.organization.id,
-        releaseConcurrency: {
-          releaseQueue: true,
-        },
+        releaseConcurrency: body.releaseConcurrency,
       });
 
       return json({

apps/webapp/app/routes/engine.v1.runs.$runFriendlyId.waitpoints.tokens.$waitpointFriendlyId.wait.ts

@@ -34,10 +34,10 @@ const { action } = createActionApiRoute(
         throw json({ error: "Waitpoint not found" }, { status: 404 });
       }
 
+      // TODO: Add releaseConcurrency from the body
       const result = await engine.blockRunWithWaitpoint({
         runId,
         waitpoints: [waitpointId],
-        environmentId: authentication.environment.id,
         projectId: authentication.environment.project.id,
         organizationId: authentication.environment.organization.id,
       });

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;

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

@@ -4,6 +4,9 @@ import {
   packetRequiresOffloading,
   QueueOptions,
   SemanticInternalAttributes,
+  TaskRunError,
+  taskRunErrorEnhancer,
+  taskRunErrorToString,
   TriggerTaskRequestBody,
 } from "@trigger.dev/core/v3";
 import {
@@ -164,10 +167,10 @@ export class TriggerTaskServiceV2 extends WithRunEngine {
                         index: options.batchIndex ?? 0,
                       }
                     : undefined,
-                  environmentId: environment.id,
                   projectId: environment.projectId,
                   organizationId: environment.organizationId,
                   tx: this._prisma,
+                  releaseConcurrency: body.options?.releaseConcurrency,
                 });
               }
             );
@@ -271,7 +274,7 @@ export class TriggerTaskServiceV2 extends WithRunEngine {
             immediate: true,
           },
           async (event, traceContext, traceparent) => {
-            const run = await autoIncrementCounter.incrementInTransaction(
+            const result = await autoIncrementCounter.incrementInTransaction(
               `v3-run:${environment.id}:${taskId}`,
               async (num, tx) => {
                 const lockedToBackgroundWorker = body.options?.lockToVersion
@@ -370,11 +373,18 @@ export class TriggerTaskServiceV2 extends WithRunEngine {
                       : undefined,
                     machine: body.options?.machine,
                     priorityMs: body.options?.priority ? body.options.priority * 1_000 : undefined,
+                    releaseConcurrency: body.options?.releaseConcurrency,
                   },
                   this._prisma
                 );
 
-                return { run: taskRun, isCached: false };
+                const error = taskRun.error ? TaskRunError.parse(taskRun.error) : undefined;
+
+                if (error) {
+                  event.failWithError(error);
+                }
+
+                return { run: taskRun, error, isCached: false };
               },
               async (_, tx) => {
                 const counter = await tx.taskRunNumberCounter.findFirst({
@@ -390,7 +400,13 @@ export class TriggerTaskServiceV2 extends WithRunEngine {
               this._prisma
             );
 
-            return run;
+            if (result?.error) {
+              throw new ServiceValidationError(
+                taskRunErrorToString(taskRunErrorEnhancer(result.error))
+              );
+            }
+
+            return result;
           }
         );
       } catch (error) {

internal-packages/database/prisma/migrations/20250314133612_add_queued_executing_status_to_snapshots/migration.sql

@@ -0,0 +1,4 @@
+-- AlterEnum
+ALTER TYPE "TaskRunExecutionStatus"
+ADD
+  VALUE 'QUEUED_EXECUTING';

internal-packages/database/prisma/migrations/20250318163201_add_previous_snapshot_id_to_task_run_execution_snapshot/migration.sql

@@ -0,0 +1,8 @@
+-- DropIndex
+DROP INDEX "SecretStore_key_idx";
+
+-- AlterTable
+ALTER TABLE "TaskRunExecutionSnapshot" ADD COLUMN     "previousSnapshotId" TEXT;
+
+-- CreateIndex
+CREATE INDEX "SecretStore_key_idx" ON "SecretStore"("key" text_pattern_ops);