re2: New release concurrency system

.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);

internal-packages/database/prisma/migrations/20250319103257_add_release_concurrency_on_waitpoint_to_task_queue/migration.sql

@@ -0,0 +1,5 @@
+-- AlterTable
+ALTER TABLE
+  "TaskQueue"
+ADD
+  COLUMN "releaseConcurrencyOnWaitpoint" BOOLEAN NOT NULL DEFAULT false;

internal-packages/database/prisma/migrations/20250319110754_add_org_and_project_to_execution_snapshots/migration.sql

@@ -0,0 +1,26 @@
+/*
+ Warnings:
+
+ - Added the required column `organizationId` to the `TaskRunExecutionSnapshot` table without a default value. This is not possible if the table is not empty.
+ - Added the required column `projectId` to the `TaskRunExecutionSnapshot` table without a default value. This is not possible if the table is not empty.
+
+ */
+-- AlterTable
+ALTER TABLE
+  "TaskRunExecutionSnapshot"
+ADD
+  COLUMN "organizationId" TEXT NOT NULL,
+ADD
+  COLUMN "projectId" TEXT NOT NULL;
+
+-- AddForeignKey
+ALTER TABLE
+  "TaskRunExecutionSnapshot"
+ADD
+  CONSTRAINT "TaskRunExecutionSnapshot_projectId_fkey" FOREIGN KEY ("projectId") REFERENCES "Project"("id") ON DELETE RESTRICT ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE
+  "TaskRunExecutionSnapshot"
+ADD
+  CONSTRAINT "TaskRunExecutionSnapshot_organizationId_fkey" FOREIGN KEY ("organizationId") REFERENCES "Organization"("id") ON DELETE RESTRICT ON UPDATE CASCADE;

internal-packages/database/prisma/migrations/20250319114436_add_metadata_to_task_run_execution_snapshots/migration.sql

@@ -0,0 +1,5 @@
+-- AlterTable
+ALTER TABLE
+  "TaskRunExecutionSnapshot"
+ADD
+  COLUMN "metadata" JSONB;

internal-packages/database/prisma/migrations/20250319131807_add_locked_queue_id_to_task_run/migration.sql

@@ -0,0 +1,5 @@
+-- AlterTable
+ALTER TABLE
+  "TaskRun"
+ADD
+  COLUMN "lockedQueueId" TEXT;

internal-packages/database/prisma/schema.prisma

@@ -164,6 +164,7 @@ model Organization {
   organizationIntegrations OrganizationIntegration[]
   workerGroups             WorkerInstanceGroup[]
   workerInstances          WorkerInstance[]
+  executionSnapshots       TaskRunExecutionSnapshot[]
 }
 
 model ExternalAccount {
@@ -504,6 +505,7 @@ model Project {
   waitpoints             Waitpoint[]
   taskRunWaitpoints      TaskRunWaitpoint[]
   taskRunCheckpoints     TaskRunCheckpoint[]
+  executionSnapshots     TaskRunExecutionSnapshot[]
 }
 
 enum ProjectVersion {
@@ -1724,7 +1726,9 @@ model TaskRun {
   projectId String
 
   // The specific queue this run is in
-  queue String
+  queue         String
+  // The queueId is set when the run is locked to a specific queue
+  lockedQueueId String?
 
   /// The main queue that this run is part of
   masterQueue          String  @default("main")
@@ -1965,6 +1969,9 @@ model TaskRunExecutionSnapshot {
   isValid Boolean @default(true)
   error   String?
 
+  /// The previous snapshot ID
+  previousSnapshotId String?
+
   /// Run
   runId     String
   run       TaskRun       @relation(fields: [runId], references: [id])
@@ -1982,6 +1989,12 @@ model TaskRunExecutionSnapshot {
   environment     RuntimeEnvironment     @relation(fields: [environmentId], references: [id])
   environmentType RuntimeEnvironmentType
 
+  projectId String
+  project   Project @relation(fields: [projectId], references: [id])
+
+  organizationId String
+  organization   Organization @relation(fields: [organizationId], references: [id])
+
   /// Waitpoints that have been completed for this execution
   completedWaitpoints Waitpoint[] @relation("completedWaitpoints")
 
@@ -2003,6 +2016,9 @@ model TaskRunExecutionSnapshot {
 
   lastHeartbeatAt DateTime?
 
+  /// Metadata used by various systems in the run engine
+  metadata Json?
+
   /// Used to get the latest valid snapshot quickly
   @@index([runId, isValid, createdAt(sort: Desc)])
 }
@@ -2012,6 +2028,8 @@ enum TaskRunExecutionStatus {
   RUN_CREATED
   /// Run is in the RunQueue
   QUEUED
+  /// Run is in the RunQueue, and is also executing. This happens when a run is continued cannot reacquire concurrency
+  QUEUED_EXECUTING
   /// Run has been pulled from the queue, but isn't executing yet
   PENDING_EXECUTING
   /// Run is executing on a worker
@@ -2526,6 +2544,9 @@ model TaskQueue {
 
   paused Boolean @default(false)
 
+  /// If true, when a run is paused and waiting for waitpoints to be completed, the run will release the concurrency capacity.
+  releaseConcurrencyOnWaitpoint Boolean @default(false)
+
   createdAt DateTime @default(now())
   updatedAt DateTime @updatedAt