WIP new release concurrency system

Author: ericallam

internal-packages/database/prisma/migrations/20250715124534_add_concurrency_limit_burst_factor_to_runtime_environment/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "RuntimeEnvironment" ADD COLUMN     "concurrencyLimitBurstFactor" DECIMAL(4,2) NOT NULL DEFAULT 2.00;

internal-packages/database/prisma/schema.prisma

@@ -247,8 +247,9 @@ model RuntimeEnvironment {
   ///A memorable code for the environment
   shortcode String
 
-  maximumConcurrencyLimit Int     @default(5)
-  paused                  Boolean @default(false)
+  maximumConcurrencyLimit     Int     @default(5)
+  concurrencyLimitBurstFactor Decimal @default("2.00") @db.Decimal(4, 2)
+  paused                      Boolean @default(false)
 
   autoEnableInternalSources Boolean @default(true)
 

internal-packages/database/src/transaction.ts

@@ -9,6 +9,8 @@ export type PrismaClientOrTransaction = PrismaClient | PrismaTransactionClient;
 
 export type PrismaReplicaClient = Omit<PrismaClient, "$transaction">;
 
+export const Decimal = Prisma.Decimal;
+
 function isTransactionClient(prisma: PrismaClientOrTransaction): prisma is PrismaTransactionClient {
   return !("$transaction" in prisma);
 }

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

@@ -275,7 +275,6 @@ export class RunEngine {
       resources,
       executionSnapshotSystem: this.executionSnapshotSystem,
       enqueueSystem: this.enqueueSystem,
-      releaseConcurrencySystem: this.releaseConcurrencySystem,
     });
 
     this.ttlSystem = new TtlSystem({

internal-packages/run-engine/src/engine/systems/checkpointSystem.ts

@@ -195,10 +195,13 @@ export class CheckpointSystem {
           checkpointId: taskRunCheckpoint.id,
         });
 
-        this.$.logger.debug("Refilling token bucket for release concurrency queue", {
+        this.$.logger.debug("Releasing concurrency for run because it was checkpointed", {
           snapshot,
+          newSnapshot,
         });
 
+        await this.releaseConcurrencySystem.releaseConcurrency(run);
+
         return {
           ok: true as const,
           ...executionResultFromSnapshot(newSnapshot),
@@ -228,6 +231,13 @@ export class CheckpointSystem {
           runnerId,
         });
 
+        this.$.logger.debug("Releasing concurrency for run because it was checkpointed", {
+          snapshot,
+          newSnapshot,
+        });
+
+        await this.releaseConcurrencySystem.releaseConcurrency(run);
+
         return {
           ok: true as const,
           ...executionResultFromSnapshot(newSnapshot),

internal-packages/run-engine/src/engine/systems/waitpointSystem.ts

@@ -20,20 +20,17 @@ export type WaitpointSystemOptions = {
   resources: SystemResources;
   executionSnapshotSystem: ExecutionSnapshotSystem;
   enqueueSystem: EnqueueSystem;
-  releaseConcurrencySystem: ReleaseConcurrencySystem;
 };
 
 export class WaitpointSystem {
   private readonly $: SystemResources;
   private readonly executionSnapshotSystem: ExecutionSnapshotSystem;
-  private readonly releaseConcurrencySystem: ReleaseConcurrencySystem;
   private readonly enqueueSystem: EnqueueSystem;
 
   constructor(private readonly options: WaitpointSystemOptions) {
     this.$ = options.resources;
     this.executionSnapshotSystem = options.executionSnapshotSystem;
     this.enqueueSystem = options.enqueueSystem;
-    this.releaseConcurrencySystem = options.releaseConcurrencySystem;
   }
 
   shouldReleaseConcurrencyOnWaitpointForQueue(queue: TaskQueue) {
@@ -456,29 +453,6 @@ export class WaitpointSystem {
 
         // Let the worker know immediately, so it can suspend the run
         await sendNotificationToWorker({ runId, snapshot, eventBus: this.$.eventBus });
-
-        if (isRunBlocked) {
-          //release concurrency
-          const run = await this.$.prisma.taskRun.findFirst({
-            where: { id: runId },
-            select: {
-              id: true,
-              organizationId: true,
-              lockedQueueReleaseConcurrencyOnWaitpoint: true,
-            },
-          });
-
-          if (!run) {
-            this.$.logger.error(
-              "WaitpointSystem.blockRunWithWaitpoint(): Run not found, cannot release concurrency",
-              {
-                runId,
-              }
-            );
-          } else {
-            await this.releaseConcurrencySystem.releaseConcurrency(run, releaseConcurrency);
-          }
-        }
       }
 
       if (timeout) {
@@ -557,6 +531,7 @@ export class WaitpointSystem {
               id: true,
               type: true,
               maximumConcurrencyLimit: true,
+              concurrencyLimitBurstFactor: true,
               project: { select: { id: true } },
               organization: { select: { id: true } },
             },