WIP new reacquire concurrency system

Author: ericallam

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/schema.prisma

@@ -2011,6 +2011,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

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

@@ -3908,6 +3908,9 @@ export class RunEngine {
         case "QUEUED": {
           throw new NotImplementedError("There shouldn't be a heartbeat for QUEUED");
         }
+        case "QUEUED_EXECUTING": {
+          throw new NotImplementedError("There shouldn't be a heartbeat for QUEUED_EXECUTING");
+        }
         case "PENDING_EXECUTING": {
           //the run didn't start executing, we need to requeue it
           const run = await prisma.taskRun.findFirst({

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

@@ -1,7 +1,7 @@
 import { TaskRunExecutionStatus, TaskRunStatus } from "@trigger.dev/database";
 
 export function isDequeueableExecutionStatus(status: TaskRunExecutionStatus): boolean {
-  const dequeuableExecutionStatuses: TaskRunExecutionStatus[] = ["QUEUED"];
+  const dequeuableExecutionStatuses: TaskRunExecutionStatus[] = ["QUEUED", "QUEUED_EXECUTING"];
   return dequeuableExecutionStatuses.includes(status);
 }
 

…es/run-engine/src/engine/locking.test.ts …-engine/src/engine/tests/locking.test.tsinternal-packages/run-engine/src/engine/locking.test.ts renamed to internal-packages/run-engine/src/engine/tests/locking.test.ts internal-packages/run-engine/src/engine/locking.test.ts renamed to internal-packages/run-engine/src/engine/tests/locking.test.ts

@@ -0,0 +1,5 @@
+export class MessageNotFoundError extends Error {
+  constructor(messageId: string) {
+    super(`Message not found: ${messageId}`);
+  }
+}

internal-packages/run-engine/src/run-queue/errors.ts

@@ -29,6 +29,7 @@ import {
   type RedisOptions,
   type Result,
 } from "@internal/redis";
+import { MessageNotFoundError } from "./errors.js";
 
 const SemanticAttributes = {
   QUEUE: "runqueue.queue",
@@ -530,12 +531,9 @@ export class RunQueue {
         });
 
         return this.redis.releaseConcurrency(
-          this.keys.messageKey(orgId, messageId),
-          message.queue,
           this.keys.currentConcurrencyKeyFromQueue(message.queue),
           this.keys.envCurrentConcurrencyKeyFromQueue(message.queue),
-          messageId,
-          JSON.stringify(message.masterQueues)
+          messageId
         );
       },
       {
@@ -556,11 +554,7 @@ export class RunQueue {
         const message = await this.readMessage(orgId, messageId);
 
         if (!message) {
-          this.logger.log(`[${this.name}].acknowledgeMessage() message not found`, {
-            messageId,
-            service: this.name,
-          });
-          return;
+          throw new MessageNotFoundError(messageId);
         }
 
         span.setAttributes({
@@ -570,14 +564,25 @@ export class RunQueue {
           [SemanticAttributes.CONCURRENCY_KEY]: message.concurrencyKey,
         });
 
-        return this.redis.reacquireConcurrency(
-          this.keys.messageKey(orgId, messageId),
-          message.queue,
-          this.keys.currentConcurrencyKeyFromQueue(message.queue),
-          this.keys.envCurrentConcurrencyKeyFromQueue(message.queue),
+        const queueCurrentConcurrencyKey = this.keys.currentConcurrencyKeyFromQueue(message.queue);
+        const envCurrentConcurrencyKey = this.keys.envCurrentConcurrencyKeyFromQueue(message.queue);
+        const queueReserveConcurrencyKey = this.keys.reserveConcurrencyKeyFromQueue(message.queue);
+        const envReserveConcurrencyKey = this.keys.envReserveConcurrencyKeyFromQueue(message.queue);
+        const queueConcurrencyLimitKey = this.keys.concurrencyLimitKeyFromQueue(message.queue);
+        const envConcurrencyLimitKey = this.keys.envConcurrencyLimitKeyFromQueue(message.queue);
+
+        const result = await this.redis.reacquireConcurrency(
+          queueCurrentConcurrencyKey,
+          envCurrentConcurrencyKey,
+          queueReserveConcurrencyKey,
+          envReserveConcurrencyKey,
+          queueConcurrencyLimitKey,
+          envConcurrencyLimitKey,
           messageId,
-          JSON.stringify(message.masterQueues)
+          String(this.options.defaultEnvConcurrency)
         );
+
+        return !!result;
       },
       {
         kind: SpanKind.CONSUMER,
@@ -1447,41 +1452,76 @@ redis.call('SREM', envCurrentConcurrencyKey, messageId)
     });
 
     this.redis.defineCommand("releaseConcurrency", {
-      numberOfKeys: 5,
+      numberOfKeys: 2,
       lua: `
 -- Keys:
-local messageKey = KEYS[1]
-local messageQueue = KEYS[2]
-local concurrencyKey = KEYS[3]
-local envCurrentConcurrencyKey = KEYS[4]
-local envQueueKey = KEYS[5]
+local queueCurrentConcurrencyKey = KEYS[1]
+local envCurrentConcurrencyKey = KEYS[2]
 
 -- Args:
 local messageId = ARGV[1]
 
 -- Update the concurrency keys
-if concurrencyKey ~= "" then
-  redis.call('SREM', concurrencyKey, messageId)
-end
+redis.call('SREM', queueCurrentConcurrencyKey, messageId)
 redis.call('SREM', envCurrentConcurrencyKey, messageId)
 `,
     });
 
     this.redis.defineCommand("reacquireConcurrency", {
-      numberOfKeys: 4,
+      numberOfKeys: 6,
       lua: `
 -- Keys:
-local messageKey = KEYS[1]
-local messageQueue = KEYS[2]
-local concurrencyKey = KEYS[3]
-local envCurrentConcurrencyKey = KEYS[4]
+local queueCurrentConcurrencyKey = KEYS[1]
+local envCurrentConcurrencyKey = KEYS[2]
+local queueReserveConcurrencyKey = KEYS[3]
+local envReserveConcurrencyKey = KEYS[4]
+local queueConcurrencyLimitKey = KEYS[5]
+local envConcurrencyLimitKey = KEYS[6]
 
 -- Args:
 local messageId = ARGV[1]
+local defaultEnvConcurrencyLimit = ARGV[2]
+
+-- Check if the message is already in either current concurrency set
+local isInQueueConcurrency = redis.call('SISMEMBER', queueCurrentConcurrencyKey, messageId) == 1
+local isInEnvConcurrency = redis.call('SISMEMBER', envCurrentConcurrencyKey, messageId) == 1
+
+-- If it's already in both sets, we're done
+if isInQueueConcurrency and isInEnvConcurrency then
+    return true
+end
+
+-- Check current env concurrency against the limit
+local envCurrentConcurrency = tonumber(redis.call('SCARD', envCurrentConcurrencyKey) or '0')
+local envConcurrencyLimit = tonumber(redis.call('GET', envConcurrencyLimitKey) or defaultEnvConcurrencyLimit)
+local envReserveConcurrency = tonumber(redis.call('SCARD', envReserveConcurrencyKey) or '0')
+local totalEnvConcurrencyLimit = envConcurrencyLimit + envReserveConcurrency
+
+if envCurrentConcurrency >= totalEnvConcurrencyLimit then
+    return false
+end
+
+-- Check current queue concurrency against the limit
+local queueCurrentConcurrency = tonumber(redis.call('SCARD', queueCurrentConcurrencyKey) or '0')
+local queueConcurrencyLimit = math.min(tonumber(redis.call('GET', queueConcurrencyLimitKey) or '1000000'), envConcurrencyLimit)
+local queueReserveConcurrency = tonumber(redis.call('SCARD', queueReserveConcurrencyKey) or '0')
+local totalQueueConcurrencyLimit = queueConcurrencyLimit + queueReserveConcurrency
+
+if queueCurrentConcurrency >= totalQueueConcurrencyLimit then
+    return false
+end
 
 -- Update the concurrency keys
-redis.call('SADD', concurrencyKey, messageId)
+redis.call('SADD', queueCurrentConcurrencyKey, messageId)
 redis.call('SADD', envCurrentConcurrencyKey, messageId)
+
+-- Remove the message from the queue reserve concurrency set
+redis.call('SREM', queueReserveConcurrencyKey, messageId)
+
+-- Remove the message from the env reserve concurrency set
+redis.call('SREM', envReserveConcurrencyKey, messageId)
+
+return true
 `,
     });
 
@@ -1630,24 +1670,23 @@ declare module "@internal/redis" {
     ): Result<void, Context>;
 
     releaseConcurrency(
-      messageKey: string,
-      messageQueue: string,
-      concurrencyKey: string,
-      envConcurrencyKey: string,
+      queueCurrentConcurrencyKey: string,
+      envCurrentConcurrencyKey: string,
       messageId: string,
-      masterQueues: string,
       callback?: Callback<void>
     ): Result<void, Context>;
 
     reacquireConcurrency(
-      messageKey: string,
-      messageQueue: string,
-      concurrencyKey: string,
-      envConcurrencyKey: string,
+      queueCurrentConcurrencyKey: string,
+      envCurrentConcurrencyKey: string,
+      queueReserveConcurrencyKey: string,
+      envReserveConcurrencyKey: string,
+      queueConcurrencyLimitKey: string,
+      envConcurrencyLimitKey: string,
       messageId: string,
-      masterQueues: string,
-      callback?: Callback<void>
-    ): Result<void, Context>;
+      defaultEnvConcurrencyLimit: string,
+      callback?: Callback<string>
+    ): Result<string, Context>;
 
     updateGlobalConcurrencyLimits(
       envConcurrencyLimitKey: string,