Implement release concurrency system

Author: ericallam

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

@@ -48,10 +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: true,
+        releaseConcurrency: body.releaseConcurrency,
       });
 
       return json({

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

@@ -34,13 +34,12 @@ 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,
-        releaseConcurrency: true,
       });
 
       return json<WaitForWaitpointTokenResponseBody>(

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

@@ -167,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,
                 });
               }
             );
@@ -373,6 +373,7 @@ 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
                 );

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 {
@@ -503,6 +504,7 @@ model Project {
   waitpoints             Waitpoint[]
   taskRunWaitpoints      TaskRunWaitpoint[]
   taskRunCheckpoints     TaskRunCheckpoint[]
+  executionSnapshots     TaskRunExecutionSnapshot[]
 }
 
 enum ProjectVersion {
@@ -1723,7 +1725,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")
@@ -1984,6 +1988,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")
 
@@ -2005,6 +2015,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)])
 }
@@ -2528,6 +2541,9 @@ model TaskQueue {
   concurrencyLimit Int?
   rateLimit        Json?
 
+  /// 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
 

internal-packages/run-engine/src/engine/index.ts

@@ -36,7 +36,6 @@ import {
 import { EventBus, EventBusEvents } from "./eventBus.js";
 import { RunLocker } from "./locking.js";
 import { ReleaseConcurrencyTokenBucketQueue } from "./releaseConcurrencyTokenBucketQueue.js";
-import { canReleaseConcurrency } from "./statuses.js";
 import { BatchSystem } from "./systems/batchSystem.js";
 import { CheckpointSystem } from "./systems/checkpointSystem.js";
 import { DelayedRunSystem } from "./systems/delayedRunSystem.js";
@@ -46,13 +45,14 @@ import {
   ExecutionSnapshotSystem,
   getLatestExecutionSnapshot,
 } from "./systems/executionSnapshotSystem.js";
+import { ReleaseConcurrencySystem } from "./systems/releaseConcurrencySystem.js";
 import { RunAttemptSystem } from "./systems/runAttemptSystem.js";
 import { SystemResources } from "./systems/systems.js";
 import { TtlSystem } from "./systems/ttlSystem.js";
+import { WaitingForWorkerSystem } from "./systems/waitingForWorkerSystem.js";
 import { WaitpointSystem } from "./systems/waitpointSystem.js";
 import { EngineWorker, HeartbeatTimeouts, RunEngineOptions, TriggerParams } from "./types.js";
 import { workerCatalog } from "./workerCatalog.js";
-import { WaitingForWorkerSystem } from "./systems/waitingForWorkerSystem.js";
 
 export class RunEngine {
   private runLockRedis: Redis;
@@ -63,7 +63,7 @@ export class RunEngine {
   private logger = new Logger("RunEngine", "debug");
   private tracer: Tracer;
   private heartbeatTimeouts: HeartbeatTimeouts;
-  private releaseConcurrencyQueue: ReleaseConcurrencyTokenBucketQueue<{
+  releaseConcurrencyQueue: ReleaseConcurrencyTokenBucketQueue<{
     orgId: string;
     projectId: string;
     envId: string;
@@ -79,6 +79,7 @@ export class RunEngine {
   delayedRunSystem: DelayedRunSystem;
   ttlSystem: TtlSystem;
   waitingForWorkerSystem: WaitingForWorkerSystem;
+  releaseConcurrencySystem: ReleaseConcurrencySystem;
 
   constructor(private readonly options: RunEngineOptions) {
     this.prisma = options.prisma;
@@ -188,7 +189,7 @@ export class RunEngine {
       redis: {
         ...options.queue.redis, // Use base queue redis options
         ...options.releaseConcurrency?.redis, // Allow overrides
-        keyPrefix: `${options.queue.redis.keyPrefix}release-concurrency:`,
+        keyPrefix: `${options.queue.redis.keyPrefix ?? ""}release-concurrency:`,
       },
       retry: {
         maxRetries: options.releaseConcurrency?.maxRetries ?? 5,
@@ -201,8 +202,8 @@ export class RunEngine {
       consumersCount: options.releaseConcurrency?.consumersCount ?? 1,
       pollInterval: options.releaseConcurrency?.pollInterval ?? 1000,
       batchSize: options.releaseConcurrency?.batchSize ?? 10,
-      executor: async (descriptor, runId) => {
-        await this.#executeReleasedConcurrencyFromQueue(descriptor, runId);
+      executor: async (descriptor, snapshotId) => {
+        await this.releaseConcurrencySystem.executeReleaseConcurrencyForSnapshot(snapshotId);
       },
       maxTokens: async (descriptor) => {
         const environment = await this.prisma.runtimeEnvironment.findFirstOrThrow({
@@ -239,6 +240,10 @@ export class RunEngine {
       releaseConcurrencyQueue: this.releaseConcurrencyQueue,
     };
 
+    this.releaseConcurrencySystem = new ReleaseConcurrencySystem({
+      resources,
+    });
+
     this.executionSnapshotSystem = new ExecutionSnapshotSystem({
       resources,
       heartbeatTimeouts: this.heartbeatTimeouts,
@@ -251,6 +256,7 @@ export class RunEngine {
 
     this.checkpointSystem = new CheckpointSystem({
       resources,
+      releaseConcurrencySystem: this.releaseConcurrencySystem,
       executionSnapshotSystem: this.executionSnapshotSystem,
       enqueueSystem: this.enqueueSystem,
     });
@@ -269,6 +275,7 @@ export class RunEngine {
       resources,
       executionSnapshotSystem: this.executionSnapshotSystem,
       enqueueSystem: this.enqueueSystem,
+      releaseConcurrencySystem: this.releaseConcurrencySystem,
     });
 
     this.ttlSystem = new TtlSystem({
@@ -344,6 +351,7 @@ export class RunEngine {
       machine,
       workerId,
       runnerId,
+      releaseConcurrency,
     }: TriggerParams,
     tx?: PrismaClientOrTransaction
   ): Promise<TaskRun> {
@@ -435,6 +443,8 @@ export class RunEngine {
                   runStatus: status,
                   environmentId: environment.id,
                   environmentType: environment.type,
+                  projectId: environment.project.id,
+                  organizationId: environment.organization.id,
                   workerId,
                   runnerId,
                 },
@@ -490,12 +500,11 @@ export class RunEngine {
               runId: parentTaskRunId,
               waitpoints: associatedWaitpoint.id,
               projectId: associatedWaitpoint.projectId,
-              organizationId: environment.organization.id,
               batch,
               workerId,
               runnerId,
               tx: prisma,
-              releaseConcurrency: true, // TODO: This needs to use the release concurrency system
+              releaseConcurrency,
             });
           }
 
@@ -1003,7 +1012,6 @@ export class RunEngine {
     runId,
     waitpoints,
     projectId,
-    organizationId,
     releaseConcurrency,
     timeout,
     spanIdToComplete,
@@ -1028,7 +1036,6 @@ export class RunEngine {
       runId,
       waitpoints,
       projectId,
-      organizationId,
       releaseConcurrency,
       timeout,
       spanIdToComplete,
@@ -1039,35 +1046,6 @@ export class RunEngine {
     });
   }
 
-  async #executeReleasedConcurrencyFromQueue(
-    descriptor: { orgId: string; projectId: string; envId: string },
-    runId: string
-  ) {
-    this.logger.debug("Executing released concurrency", {
-      descriptor,
-      runId,
-    });
-
-    // - Runlock the run
-    // - Get latest snapshot
-    // - If the run is non suspended or going to be, then bail
-    // - If the run is suspended or going to be, then release the concurrency
-    await this.runLock.lock([runId], 5_000, async () => {
-      const snapshot = await getLatestExecutionSnapshot(this.prisma, runId);
-
-      if (!canReleaseConcurrency(snapshot.executionStatus)) {
-        this.logger.debug("Run is not in a state to release concurrency", {
-          runId,
-          snapshot,
-        });
-
-        return;
-      }
-
-      return await this.runQueue.releaseConcurrency(descriptor.orgId, snapshot.runId);
-    });
-  }
-
   /** This completes a waitpoint and updates all entries so the run isn't blocked,
    * if they're no longer blocked. This doesn't suffer from race conditions. */
   async completeWaitpoint({
@@ -1328,7 +1306,8 @@ export class RunEngine {
               id: latestSnapshot.environmentId,
               type: latestSnapshot.environmentType,
             },
-            orgId: run.runtimeEnvironment.organizationId,
+            orgId: latestSnapshot.organizationId,
+            projectId: latestSnapshot.projectId,
             error: {
               type: "INTERNAL_ERROR",
               code: "TASK_RUN_DEQUEUED_MAX_RETRIES",