fix: refill release concurrency token bucket queue when runs resume before checkpoints are created (#1933)

* fix: refill release concurrency token bucket queue when runs resume before checkpoints are created

* Fix the engine package.json tests

* Remove log

Author: ericallam

Diff for: internal-packages/run-engine/package.json

@@ -41,4 +41,4 @@
     "build": "pnpm run clean && tsc -p tsconfig.build.json",
     "dev": "tsc --watch  -p tsconfig.build.json"
   }
-}
+}

Diff for: internal-packages/run-engine/src/engine/index.ts

@@ -305,6 +305,7 @@ export class RunEngine {
       executionSnapshotSystem: this.executionSnapshotSystem,
       runAttemptSystem: this.runAttemptSystem,
       machines: this.options.machines,
+      releaseConcurrencySystem: this.releaseConcurrencySystem,
     });
   }
 

Diff for: internal-packages/run-engine/src/engine/releaseConcurrencyTokenBucketQueue.ts

@@ -259,6 +259,59 @@ export class ReleaseConcurrencyTokenBucketQueue<T> {
     });
   }
 
+  public async getReleaseQueueMetrics(releaseQueueDescriptor: T) {
+    const releaseQueue = this.keys.fromDescriptor(releaseQueueDescriptor);
+    const currentTokensRaw = await this.redis.get(this.#bucketKey(releaseQueue));
+    const queueLength = await this.redis.zcard(this.#queueKey(releaseQueue));
+
+    const currentTokens = currentTokensRaw ? Number(currentTokensRaw) : undefined;
+
+    return { currentTokens, queueLength };
+  }
+
+  /**
+   * Refill a token only if the releaserId is not in the release queue.
+   * Returns true if the token was refilled, false if the releaserId was found in the queue.
+   */
+  public async refillTokenIfNotInQueue(
+    releaseQueueDescriptor: T,
+    releaserId: string
+  ): Promise<boolean> {
+    const maxTokens = await this.#callMaxTokens(releaseQueueDescriptor);
+    const releaseQueue = this.keys.fromDescriptor(releaseQueueDescriptor);
+
+    if (maxTokens === 0) {
+      this.logger.debug("No tokens available, skipping refill", {
+        releaseQueueDescriptor,
+        releaserId,
+        maxTokens,
+        releaseQueue,
+      });
+
+      return false;
+    }
+
+    const result = await this.redis.refillTokenIfNotInQueue(
+      this.masterQueuesKey,
+      this.#bucketKey(releaseQueue),
+      this.#queueKey(releaseQueue),
+      this.#metadataKey(releaseQueue),
+      releaseQueue,
+      releaserId,
+      String(maxTokens)
+    );
+
+    this.logger.debug("Attempted to refill token if not in queue", {
+      releaseQueueDescriptor,
+      releaserId,
+      maxTokens,
+      releaseQueue,
+      result,
+    });
+
+    return result === "true";
+  }
+
   /**
    * Get the next queue that has available capacity and process one item from it
    * Returns true if an item was processed, false if no items were available
@@ -783,6 +836,51 @@ end
 return true
       `,
     });
+
+    this.redis.defineCommand("refillTokenIfNotInQueue", {
+      numberOfKeys: 4,
+      lua: `
+local masterQueuesKey = KEYS[1]
+local bucketKey = KEYS[2]
+local queueKey = KEYS[3]
+local metadataKey = KEYS[4]
+
+local releaseQueue = ARGV[1]
+local releaserId = ARGV[2]
+local maxTokens = tonumber(ARGV[3])
+
+-- Check if the releaserId is in the queue
+local score = redis.call("ZSCORE", queueKey, releaserId)
+if score then
+  -- Item is in queue, don't refill token
+  return redis.status_reply("false")
+end
+
+-- Return the token to the bucket
+local currentTokens = tonumber(redis.call("GET", bucketKey) or maxTokens)
+local remainingTokens = currentTokens + 1
+
+-- Don't exceed maxTokens
+if remainingTokens > maxTokens then
+  remainingTokens = maxTokens
+end
+
+redis.call("SET", bucketKey, remainingTokens)
+
+-- Clean up any metadata just in case
+redis.call("HDEL", metadataKey, releaserId)
+
+-- Update the master queue based on remaining queue length
+local queueLength = redis.call("ZCARD", queueKey)
+if queueLength > 0 then
+  redis.call("ZADD", masterQueuesKey, remainingTokens, releaseQueue)
+else
+  redis.call("ZREM", masterQueuesKey, releaseQueue)
+end
+
+return redis.status_reply("true")
+      `,
+    });
   }
 }
 
@@ -839,6 +937,17 @@ declare module "@internal/redis" {
       releaserId: string,
       callback?: Callback<void>
     ): Result<void, Context>;
+
+    refillTokenIfNotInQueue(
+      masterQueuesKey: string,
+      bucketKey: string,
+      queueKey: string,
+      metadataKey: string,
+      releaseQueue: string,
+      releaserId: string,
+      maxTokens: string,
+      callback?: Callback<string>
+    ): Result<string, Context>;
   }
 }
 

Diff for: internal-packages/run-engine/src/engine/systems/dequeueSystem.ts

@@ -11,23 +11,27 @@ import { ExecutionSnapshotSystem, getLatestExecutionSnapshot } from "./execution
 import { RunAttemptSystem } from "./runAttemptSystem.js";
 import { SystemResources } from "./systems.js";
 import { sendNotificationToWorker } from "../eventBus.js";
+import { ReleaseConcurrencySystem } from "./releaseConcurrencySystem.js";
 
 export type DequeueSystemOptions = {
   resources: SystemResources;
   machines: RunEngineOptions["machines"];
   executionSnapshotSystem: ExecutionSnapshotSystem;
   runAttemptSystem: RunAttemptSystem;
+  releaseConcurrencySystem: ReleaseConcurrencySystem;
 };
 
 export class DequeueSystem {
   private readonly $: SystemResources;
   private readonly executionSnapshotSystem: ExecutionSnapshotSystem;
   private readonly runAttemptSystem: RunAttemptSystem;
+  private readonly releaseConcurrencySystem: ReleaseConcurrencySystem;
 
   constructor(private readonly options: DequeueSystemOptions) {
     this.$ = options.resources;
     this.executionSnapshotSystem = options.executionSnapshotSystem;
     this.runAttemptSystem = options.runAttemptSystem;
+    this.releaseConcurrencySystem = options.releaseConcurrencySystem;
   }
 
   /**
@@ -158,6 +162,12 @@ export class DequeueSystem {
                   }
                 );
 
+                if (snapshot.previousSnapshotId) {
+                  await this.releaseConcurrencySystem.refillTokensForSnapshot(
+                    snapshot.previousSnapshotId
+                  );
+                }
+
                 await sendNotificationToWorker({
                   runId,
                   snapshot: newSnapshot,

Diff for: internal-packages/run-engine/src/engine/systems/releaseConcurrencySystem.ts

@@ -69,6 +69,52 @@ export class ReleaseConcurrencySystem {
     await this.releaseConcurrencyQueue.quit();
   }
 
+  public async refillTokensForSnapshot(snapshotId: string | undefined): Promise<void>;
+  public async refillTokensForSnapshot(snapshot: TaskRunExecutionSnapshot): Promise<void>;
+  public async refillTokensForSnapshot(
+    snapshotOrId: TaskRunExecutionSnapshot | string | undefined
+  ) {
+    if (!this.releaseConcurrencyQueue) {
+      return;
+    }
+
+    if (typeof snapshotOrId === "undefined") {
+      return;
+    }
+
+    const snapshot =
+      typeof snapshotOrId === "string"
+        ? await this.$.prisma.taskRunExecutionSnapshot.findFirst({
+            where: { id: snapshotOrId },
+          })
+        : snapshotOrId;
+
+    if (!snapshot) {
+      this.$.logger.error("Snapshot not found", {
+        snapshotId: snapshotOrId,
+      });
+
+      return;
+    }
+
+    if (snapshot.executionStatus !== "EXECUTING_WITH_WAITPOINTS") {
+      this.$.logger.debug("Snapshot is not in a valid state to refill tokens", {
+        snapshot,
+      });
+
+      return;
+    }
+
+    await this.releaseConcurrencyQueue.refillTokenIfNotInQueue(
+      {
+        orgId: snapshot.organizationId,
+        projectId: snapshot.projectId,
+        envId: snapshot.environmentId,
+      },
+      snapshot.id
+    );
+  }
+
   public async checkpointCreatedOnEnvironment(environment: RuntimeEnvironment) {
     if (!this.releaseConcurrencyQueue) {
       return;
@@ -86,11 +132,19 @@ export class ReleaseConcurrencySystem {
 
   public async releaseConcurrencyForSnapshot(snapshot: TaskRunExecutionSnapshot) {
     if (!this.releaseConcurrencyQueue) {
+      this.$.logger.debug("Release concurrency queue not enabled, skipping release", {
+        snapshotId: snapshot.id,
+      });
+
       return;
     }
 
     // Go ahead and release concurrency immediately if the run is in a development environment
     if (snapshot.environmentType === "DEVELOPMENT") {
+      this.$.logger.debug("Immediate release of concurrency for development environment", {
+        snapshotId: snapshot.id,
+      });
+
       return await this.executeReleaseConcurrencyForSnapshot(snapshot.id);
     }
 

Diff for: internal-packages/run-engine/src/engine/systems/waitpointSystem.ts

@@ -543,6 +543,8 @@ export class WaitpointSystem {
             }
           );
 
+          await this.releaseConcurrencySystem.refillTokensForSnapshot(snapshot);
+
           await sendNotificationToWorker({
             runId,
             snapshot: newSnapshot,