Using triggerAndWait or batchTriggerAndWait frees up concurrency

Normally it frees up just env and org concurrency. If it’s a recursive task then it will free up the run concurrency too (e.g. a task calling itself).

Author: matt-aitken

apps/webapp/app/v3/marqs/index.server.ts

@@ -542,6 +542,60 @@ export class MarQS {
     );
   }
 
+  public async releaseConcurrency(messageId: string, releaseForRun: boolean = false) {
+    return this.#trace(
+      "releaseConcurrency",
+      async (span) => {
+        span.setAttributes({
+          [SemanticAttributes.MESSAGE_ID]: messageId,
+        });
+
+        const message = await this.readMessage(messageId);
+
+        if (!message) {
+          return;
+        }
+
+        span.setAttributes({
+          [SemanticAttributes.QUEUE]: message.queue,
+          [SemanticAttributes.MESSAGE_ID]: message.messageId,
+          [SemanticAttributes.CONCURRENCY_KEY]: message.concurrencyKey,
+          [SemanticAttributes.PARENT_QUEUE]: message.parentQueue,
+        });
+
+        const concurrencyKey = this.keys.currentConcurrencyKeyFromQueue(message.queue);
+        const envConcurrencyKey = this.keys.envCurrentConcurrencyKeyFromQueue(message.queue);
+        const orgConcurrencyKey = this.keys.orgCurrentConcurrencyKeyFromQueue(message.queue);
+
+        logger.debug("Calling releaseConcurrency", {
+          messageId,
+          queue: message.queue,
+          concurrencyKey,
+          envConcurrencyKey,
+          orgConcurrencyKey,
+          service: this.name,
+          releaseForRun,
+        });
+
+        return this.redis.releaseConcurrency(
+          //don't release the for the run, it breaks concurrencyLimits
+          releaseForRun ? concurrencyKey : "",
+          envConcurrencyKey,
+          orgConcurrencyKey,
+          message.messageId
+        );
+      },
+      {
+        kind: SpanKind.CONSUMER,
+        attributes: {
+          [SEMATTRS_MESSAGING_OPERATION]: "releaseConcurrency",
+          [SEMATTRS_MESSAGE_ID]: messageId,
+          [SEMATTRS_MESSAGING_SYSTEM]: "marqs",
+        },
+      }
+    );
+  }
+
   async #trace<T>(
     name: string,
     fn: (span: Span) => Promise<T>,
@@ -1488,6 +1542,24 @@ end
 `,
     });
 
+    this.redis.defineCommand("releaseConcurrency", {
+      numberOfKeys: 3,
+      lua: `
+local concurrencyKey = KEYS[1]
+local envCurrentConcurrencyKey = KEYS[2]
+local orgCurrentConcurrencyKey = KEYS[3]
+
+local messageId = ARGV[1]
+
+-- Update the concurrency keys
+if concurrencyKey ~= "" then
+  redis.call('SREM', concurrencyKey, messageId)
+end
+redis.call('SREM', envCurrentConcurrencyKey, messageId)
+redis.call('SREM', orgCurrentConcurrencyKey, messageId)
+`,
+    });
+
     this.redis.defineCommand("heartbeatMessage", {
       numberOfKeys: 1,
       lua: `
@@ -1699,6 +1771,14 @@ declare module "ioredis" {
       callback?: Callback<void>
     ): Result<void, Context>;
 
+    releaseConcurrency(
+      concurrencyKey: string,
+      envConcurrencyKey: string,
+      orgConcurrencyKey: string,
+      messageId: string,
+      callback?: Callback<void>
+    ): Result<void, Context>;
+
     heartbeatMessage(
       visibilityQueue: string,
       messageId: string,

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

@@ -107,6 +107,7 @@ export class TriggerTaskService extends BaseService {
                 select: {
                   id: true,
                   status: true,
+                  taskIdentifier: true,
                 },
               },
             },
@@ -143,6 +144,7 @@ export class TriggerTaskService extends BaseService {
                     select: {
                       id: true,
                       status: true,
+                      taskIdentifier: true,
                     },
                   },
                 },
@@ -391,6 +393,20 @@ export class TriggerTaskService extends BaseService {
             this._prisma
           );
 
+          //release the concurrency for the env and org, if part of a (batch)triggerAndWait
+          if (dependentAttempt) {
+            const isSameTask = dependentAttempt.taskRun.taskIdentifier === taskId;
+            await marqs?.releaseConcurrency(dependentAttempt.taskRun.id, isSameTask);
+          }
+          if (dependentBatchRun?.dependentTaskAttempt) {
+            const isSameTask =
+              dependentBatchRun.dependentTaskAttempt.taskRun.taskIdentifier === taskId;
+            await marqs?.releaseConcurrency(
+              dependentBatchRun.dependentTaskAttempt.taskRun.id,
+              isSameTask
+            );
+          }
+
           if (!run) {
             return;
           }