Implement reserve concurrency clearing when the child run is acked

ericallam · ericallam · commit 8de78cb93ab0 · 2025-03-19T16:29:18.000Z
diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts
@@ -1540,6 +1540,7 @@ export class RunEngine {
             createdAt: true,
             completedAt: true,
             taskEventStore: true,
+            parentTaskRunId: true,
             runtimeEnvironment: {
               select: {
                 organizationId: true,
@@ -1559,7 +1560,9 @@ export class RunEngine {
         });
 
         //remove it from the queue and release concurrency
-        await this.runQueue.acknowledgeMessage(run.runtimeEnvironment.organizationId, runId);
+        await this.runQueue.acknowledgeMessage(run.runtimeEnvironment.organizationId, runId, {
+          messageId: run.parentTaskRunId ?? undefined,
+        });
 
         //if executing, we need to message the worker to cancel the run and put it into `PENDING_CANCEL` status
         if (isExecuting(latestSnapshot.executionStatus)) {
@@ -2778,10 +2781,13 @@ export class RunEngine {
           createdAt: true,
           completedAt: true,
           taskEventStore: true,
+          parentTaskRunId: true,
         },
       });
 
-      await this.runQueue.acknowledgeMessage(updatedRun.runtimeEnvironment.organizationId, runId);
+      await this.runQueue.acknowledgeMessage(updatedRun.runtimeEnvironment.organizationId, runId, {
+        messageId: updatedRun.parentTaskRunId ?? undefined,
+      });
 
       if (!updatedRun.associatedWaitpoint) {
         throw new ServiceValidationError("No associated waitpoint found", 400);
@@ -2919,10 +2925,14 @@ export class RunEngine {
             createdAt: true,
             completedAt: true,
             taskEventStore: true,
+            parentTaskRunId: true,
           },
         });
         const newSnapshot = await getLatestExecutionSnapshot(prisma, runId);
-        await this.runQueue.acknowledgeMessage(run.project.organizationId, runId);
+
+        await this.runQueue.acknowledgeMessage(run.project.organizationId, runId, {
+          messageId: run.parentTaskRunId ?? undefined,
+        });
 
         // We need to manually emit this as we created the final snapshot as part of the task run update
         this.eventBus.emit("executionSnapshotCreated", {
@@ -3236,6 +3246,7 @@ export class RunEngine {
           attemptNumber: true,
           spanId: true,
           batchId: true,
+          parentTaskRunId: true,
           associatedWaitpoint: {
             select: {
               id: true,
@@ -3270,13 +3281,15 @@ export class RunEngine {
         throw new ServiceValidationError("No associated waitpoint found", 400);
       }
 
+      await this.runQueue.acknowledgeMessage(run.runtimeEnvironment.organizationId, runId, {
+        messageId: run.parentTaskRunId ?? undefined,
+      });
+
       await this.completeWaitpoint({
         id: run.associatedWaitpoint.id,
         output: { value: JSON.stringify(error), isError: true },
       });
 
-      await this.runQueue.acknowledgeMessage(run.runtimeEnvironment.organizationId, runId);
-
       this.eventBus.emit("runFailed", {
         time: failedAt,
         run: {
diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts
@@ -411,18 +411,20 @@ export class RunQueue {
    * This is done when the run is in a final state.
    * @param messageId
    */
-  public async acknowledgeMessage(orgId: string, messageId: string) {
+  public async acknowledgeMessage(
+    orgId: string,
+    messageId: string,
+    reserveConcurrency?: {
+      messageId?: string;
+    }
+  ) {
     return this.#trace(
       "acknowledgeMessage",
       async (span) => {
         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({
@@ -434,6 +436,7 @@ export class RunQueue {
 
         await this.#callAcknowledgeMessage({
           message,
+          reserveConcurrency,
         });
       },
       {
@@ -936,7 +939,15 @@ export class RunQueue {
     };
   }
 
-  async #callAcknowledgeMessage({ message }: { message: OutputPayload }) {
+  async #callAcknowledgeMessage({
+    message,
+    reserveConcurrency,
+  }: {
+    message: OutputPayload;
+    reserveConcurrency?: {
+      messageId?: string;
+    };
+  }) {
     const messageId = message.runId;
     const messageKey = this.keys.messageKey(message.orgId, messageId);
     const messageQueue = message.queue;
@@ -956,22 +967,37 @@ export class RunQueue {
       service: this.name,
     });
 
-    const queueReserveConcurrencyKey = this.keys.reserveConcurrencyKeyFromQueue(messageQueue);
-    const envReserveConcurrencyKey = this.keys.envReserveConcurrencyKeyFromQueue(messageQueue);
+    if (!reserveConcurrency?.messageId) {
+      return this.redis.acknowledgeMessage(
+        messageKey,
+        messageQueue,
+        queueCurrentConcurrencyKey,
+        envCurrentConcurrencyKey,
+        envQueueKey,
+        messageId,
+        messageQueue,
+        JSON.stringify(masterQueues),
+        this.options.redis.keyPrefix ?? ""
+      );
+    } else {
+      const queueReserveConcurrencyKey = this.keys.reserveConcurrencyKeyFromQueue(messageQueue);
+      const envReserveConcurrencyKey = this.keys.envReserveConcurrencyKeyFromQueue(messageQueue);
 
-    return this.redis.acknowledgeMessage(
-      messageKey,
-      messageQueue,
-      queueCurrentConcurrencyKey,
-      envCurrentConcurrencyKey,
-      envQueueKey,
-      queueReserveConcurrencyKey,
-      envReserveConcurrencyKey,
-      messageId,
-      messageQueue,
-      JSON.stringify(masterQueues),
-      this.options.redis.keyPrefix ?? ""
-    );
+      return this.redis.acknowledgeMessageWithReserveConcurrency(
+        messageKey,
+        messageQueue,
+        queueCurrentConcurrencyKey,
+        envCurrentConcurrencyKey,
+        envQueueKey,
+        envReserveConcurrencyKey,
+        queueReserveConcurrencyKey,
+        messageId,
+        messageQueue,
+        JSON.stringify(masterQueues),
+        this.options.redis.keyPrefix ?? "",
+        reserveConcurrency.messageId
+      );
+    }
   }
 
   async #callNackMessage({ message, retryAt }: { message: OutputPayload; retryAt?: number }) {
@@ -1323,12 +1349,52 @@ return {messageId, messageScore, messagePayload} -- Return message details
     });
 
     this.redis.defineCommand("acknowledgeMessage", {
+      numberOfKeys: 5,
+      lua: `
+-- Keys:
+local messageKey = KEYS[1]
+local messageQueueKey = KEYS[2]
+local queueCurrentConcurrencyKey = KEYS[3]
+local envCurrentConcurrencyKey = KEYS[4]
+local envQueueKey = KEYS[5]
+
+-- Args:
+local messageId = ARGV[1]
+local messageQueueName = ARGV[2]
+local parentQueues = cjson.decode(ARGV[3])
+local keyPrefix = ARGV[4]
+
+-- Remove the message from the message key
+redis.call('DEL', messageKey)
+
+-- Remove the message from the queue
+redis.call('ZREM', messageQueueKey, messageId)
+redis.call('ZREM', envQueueKey, messageId)
+
+-- Rebalance the parent queues
+local earliestMessage = redis.call('ZRANGE', messageQueueKey, 0, 0, 'WITHSCORES')
+for _, parentQueue in ipairs(parentQueues) do
+  local prefixedParentQueue = keyPrefix .. parentQueue
+    if #earliestMessage == 0 then
+        redis.call('ZREM', prefixedParentQueue, messageQueueName)
+    else
+        redis.call('ZADD', prefixedParentQueue, earliestMessage[2], messageQueueName)
+    end
+end
+
+-- Update the concurrency keys
+redis.call('SREM', queueCurrentConcurrencyKey, messageId)
+redis.call('SREM', envCurrentConcurrencyKey, messageId)
+`,
+    });
+
+    this.redis.defineCommand("acknowledgeMessageWithReserveConcurrency", {
       numberOfKeys: 7,
       lua: `
 -- Keys:
 local messageKey = KEYS[1]
-local messageQueue = KEYS[2]
-local concurrencyKey = KEYS[3]
+local messageQueueKey = KEYS[2]
+local queueCurrentConcurrencyKey = KEYS[3]
 local envCurrentConcurrencyKey = KEYS[4]
 local envQueueKey = KEYS[5]
 local queueReserveConcurrencyKey = KEYS[6]
@@ -1339,16 +1405,17 @@ local messageId = ARGV[1]
 local messageQueueName = ARGV[2]
 local parentQueues = cjson.decode(ARGV[3])
 local keyPrefix = ARGV[4]
+local reserveMessageId = ARGV[5]
 
 -- Remove the message from the message key
 redis.call('DEL', messageKey)
 
 -- Remove the message from the queue
-redis.call('ZREM', messageQueue, messageId)
+redis.call('ZREM', messageQueueKey, messageId)
 redis.call('ZREM', envQueueKey, messageId)
 
 -- Rebalance the parent queues
-local earliestMessage = redis.call('ZRANGE', messageQueue, 0, 0, 'WITHSCORES')
+local earliestMessage = redis.call('ZRANGE', messageQueueKey, 0, 0, 'WITHSCORES')
 for _, parentQueue in ipairs(parentQueues) do
   local prefixedParentQueue = keyPrefix .. parentQueue
     if #earliestMessage == 0 then
@@ -1359,12 +1426,12 @@ for _, parentQueue in ipairs(parentQueues) do
 end
 
 -- Update the concurrency keys
-redis.call('SREM', concurrencyKey, messageId)
+redis.call('SREM', queueCurrentConcurrencyKey, messageId)
 redis.call('SREM', envCurrentConcurrencyKey, messageId)
 
 -- Clear reserve concurrency
-redis.call('SREM', queueReserveConcurrencyKey, messageId)
-redis.call('SREM', envReserveConcurrencyKey, messageId)
+redis.call('SREM', queueReserveConcurrencyKey, reserveMessageId)
+redis.call('SREM', envReserveConcurrencyKey, reserveMessageId)
 `,
     });
 
@@ -1515,12 +1582,6 @@ end
 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
 `,
     });
@@ -1631,12 +1692,26 @@ declare module "@internal/redis" {
       concurrencyKey: string,
       envConcurrencyKey: string,
       envQueueKey: string,
-      queueReserveConcurrencyKey: string,
+      messageId: string,
+      messageQueueName: string,
+      masterQueues: string,
+      keyPrefix: string,
+      callback?: Callback<void>
+    ): Result<void, Context>;
+
+    acknowledgeMessageWithReserveConcurrency(
+      messageKey: string,
+      messageQueue: string,
+      concurrencyKey: string,
+      envConcurrencyKey: string,
+      envQueueKey: string,
       envReserveConcurrencyKey: string,
+      queueReserveConcurrencyKey: string,
       messageId: string,
       messageQueueName: string,
       masterQueues: string,
       keyPrefix: string,
+      reserveMessageId: string,
       callback?: Callback<void>
     ): Result<void, Context>;
 
diff --git a/internal-packages/run-engine/src/run-queue/tests/ack.test.ts b/internal-packages/run-engine/src/run-queue/tests/ack.test.ts
@@ -274,7 +274,7 @@ describe("RunQueue.acknowledgeMessage", () => {
   });
 
   redisTest(
-    "acknowledging a message clears reserve concurrency sets even when not dequeued",
+    "acknowledging a message clears env reserve concurrency when recursive queue is false",
     async ({ redisContainer }) => {
       const queue = new RunQueue({
         ...testOptions,
@@ -302,8 +302,8 @@ describe("RunQueue.acknowledgeMessage", () => {
           message: messageDev,
           masterQueues: ["main", envMasterQueue],
           reserveConcurrency: {
-            messageId: messageDev.runId,
-            recursiveQueue: true,
+            messageId: "r1235",
+            recursiveQueue: false,
           },
         });
 
@@ -317,7 +317,7 @@ describe("RunQueue.acknowledgeMessage", () => {
           authenticatedEnvDev,
           messageDev.queue
         );
-        expect(queueReserveConcurrency).toBe(1);
+        expect(queueReserveConcurrency).toBe(0);
 
         // Verify message is in queue before acknowledging
         const queueLengthBefore = await queue.lengthOfQueue(authenticatedEnvDev, messageDev.queue);
@@ -327,7 +327,9 @@ describe("RunQueue.acknowledgeMessage", () => {
         expect(envQueueLengthBefore).toBe(1);
 
         // Acknowledge the message before dequeuing
-        await queue.acknowledgeMessage(messageDev.orgId, messageDev.runId);
+        await queue.acknowledgeMessage(messageDev.orgId, messageDev.runId, {
+          messageId: "r1235",
+        });
 
         // Verify reserve concurrency is cleared
         const envReserveConcurrencyAfter = await queue.reserveConcurrencyOfEnvironment(
