Remove message from worker queue in certain circumstances when acking

Author: ericallam

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

@@ -1050,7 +1050,9 @@ export class RunAttemptSystem {
         }
 
         //remove it from the queue and release concurrency
-        await this.$.runQueue.acknowledgeMessage(run.runtimeEnvironment.organizationId, runId);
+        await this.$.runQueue.acknowledgeMessage(run.runtimeEnvironment.organizationId, runId, {
+          removeFromWorkerQueue: true,
+        });
 
         await this.releaseConcurrencySystem.refillTokensForSnapshot(latestSnapshot);
 
@@ -1233,7 +1235,9 @@ export class RunAttemptSystem {
         throw new ServiceValidationError("No associated waitpoint found", 400);
       }
 
-      await this.$.runQueue.acknowledgeMessage(run.runtimeEnvironment.organizationId, runId);
+      await this.$.runQueue.acknowledgeMessage(run.runtimeEnvironment.organizationId, runId, {
+        removeFromWorkerQueue: true,
+      });
 
       await this.waitpointSystem.completeWaitpoint({
         id: run.associatedWaitpoint.id,

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

@@ -106,7 +106,13 @@ export class TtlSystem {
         },
       });
 
-      await this.$.runQueue.acknowledgeMessage(updatedRun.runtimeEnvironment.organizationId, runId);
+      await this.$.runQueue.acknowledgeMessage(
+        updatedRun.runtimeEnvironment.organizationId,
+        runId,
+        {
+          removeFromWorkerQueue: true,
+        }
+      );
 
       if (!updatedRun.associatedWaitpoint) {
         throw new ServiceValidationError("No associated waitpoint found", 400);

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

@@ -482,7 +482,7 @@ describe("RunQueue", () => {
     }
   );
 
-  redisTest("Enqueue/Dequeue a 8 shards", async ({ redisContainer }) => {
+  redisTest("Enqueue/Dequeue with 8 shards", async ({ redisContainer }) => {
     const queue = new RunQueue({
       ...testOptions,
       shardCount: 8,
@@ -716,6 +716,98 @@ describe("RunQueue", () => {
     }
   });
 
+  redisTest(
+    "Ack after moving to workerQueue with removeFromWorkerQueue = undefined",
+    { timeout: 5_000 },
+    async ({ redisContainer }) => {
+      const queue = new RunQueue({
+        ...testOptions,
+        queueSelectionStrategy: new FairQueueSelectionStrategy({
+          redis: {
+            keyPrefix: "runqueue:test:",
+            host: redisContainer.getHost(),
+            port: redisContainer.getPort(),
+          },
+          keys: testOptions.keys,
+        }),
+        redis: {
+          keyPrefix: "runqueue:test:",
+          host: redisContainer.getHost(),
+          port: redisContainer.getPort(),
+        },
+      });
+
+      try {
+        await queue.enqueueMessage({
+          env: authenticatedEnvProd,
+          message: messageProd,
+          workerQueue: "main",
+        });
+
+        const queueLength = await queue.lengthOfQueue(authenticatedEnvProd, messageProd.queue);
+        expect(queueLength).toBe(1);
+        const envQueueLength = await queue.lengthOfEnvQueue(authenticatedEnvProd);
+        expect(envQueueLength).toBe(1);
+
+        await setTimeout(1000);
+
+        await queue.acknowledgeMessage(messageProd.orgId, messageProd.runId);
+
+        const messages = await queue.peekAllOnWorkerQueue("main");
+        expect(messages.length).toEqual(1);
+      } finally {
+        await queue.quit();
+      }
+    }
+  );
+
+  redisTest(
+    "Ack after moving to workerQueue with removeFromWorkerQueue = true",
+    { timeout: 5_000 },
+    async ({ redisContainer }) => {
+      const queue = new RunQueue({
+        ...testOptions,
+        queueSelectionStrategy: new FairQueueSelectionStrategy({
+          redis: {
+            keyPrefix: "runqueue:test:",
+            host: redisContainer.getHost(),
+            port: redisContainer.getPort(),
+          },
+          keys: testOptions.keys,
+        }),
+        redis: {
+          keyPrefix: "runqueue:test:",
+          host: redisContainer.getHost(),
+          port: redisContainer.getPort(),
+        },
+      });
+
+      try {
+        await queue.enqueueMessage({
+          env: authenticatedEnvProd,
+          message: messageProd,
+          workerQueue: "main",
+        });
+
+        const queueLength = await queue.lengthOfQueue(authenticatedEnvProd, messageProd.queue);
+        expect(queueLength).toBe(1);
+        const envQueueLength = await queue.lengthOfEnvQueue(authenticatedEnvProd);
+        expect(envQueueLength).toBe(1);
+
+        await setTimeout(1000);
+
+        await queue.acknowledgeMessage(messageProd.orgId, messageProd.runId, {
+          removeFromWorkerQueue: true,
+        });
+
+        const messages = await queue.peekAllOnWorkerQueue("main");
+        expect(messages.length).toEqual(0);
+      } finally {
+        await queue.quit();
+      }
+    }
+  );
+
   redisTest("Nacking", { timeout: 15_000 }, async ({ redisContainer, redisOptions }) => {
     const queue = new RunQueue({
       ...testOptions,

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

@@ -550,7 +550,7 @@ export class RunQueue {
   public async acknowledgeMessage(
     orgId: string,
     messageId: string,
-    options?: { skipDequeueProcessing?: boolean }
+    options?: { skipDequeueProcessing?: boolean; removeFromWorkerQueue?: boolean }
   ) {
     return this.#trace(
       "acknowledgeMessage",
@@ -586,6 +586,7 @@ export class RunQueue {
 
         await this.#callAcknowledgeMessage({
           message,
+          removeFromWorkerQueue: options?.removeFromWorkerQueue,
         });
       },
       {
@@ -836,6 +837,14 @@ export class RunQueue {
     await this.redis.quit();
   }
 
+  /**
+   * Peek all messages on a worker queue (useful for tests or debugging)
+   */
+  async peekAllOnWorkerQueue(workerQueue: string) {
+    const workerQueueKey = this.keys.workerQueueKey(workerQueue);
+    return await this.redis.lrange(workerQueueKey, 0, -1);
+  }
+
   private async handleRedriveMessage(channel: string, message: string) {
     try {
       const { runId, envId, projectId, orgId } = JSON.parse(message) as any;
@@ -1430,7 +1439,13 @@ export class RunQueue {
     };
   }
 
-  async #callAcknowledgeMessage({ message }: { message: OutputPayload }) {
+  async #callAcknowledgeMessage({
+    message,
+    removeFromWorkerQueue,
+  }: {
+    message: OutputPayload;
+    removeFromWorkerQueue?: boolean;
+  }) {
     const messageId = message.runId;
     const messageKey = this.keys.messageKey(message.orgId, messageId);
     const messageQueue = message.queue;
@@ -1441,6 +1456,9 @@ export class RunQueue {
       message.environmentId,
       this.shardCount
     );
+    const workerQueue = this.#getWorkerQueueFromMessage(message);
+    const workerQueueKey = this.keys.workerQueueKey(workerQueue);
+    const messageKeyValue = this.keys.messageKey(message.orgId, messageId);
 
     this.logger.debug("Calling acknowledgeMessage", {
       messageKey,
@@ -1450,6 +1468,10 @@ export class RunQueue {
       envQueueKey,
       messageId,
       masterQueueKey,
+      workerQueue,
+      workerQueueKey,
+      removeFromWorkerQueue,
+      messageKeyValue,
       service: this.name,
     });
 
@@ -1460,8 +1482,11 @@ export class RunQueue {
       queueCurrentConcurrencyKey,
       envCurrentConcurrencyKey,
       envQueueKey,
+      workerQueueKey,
       messageId,
-      messageQueue
+      messageQueue,
+      messageKeyValue,
+      removeFromWorkerQueue ? "1" : "0"
     );
   }
 
@@ -1767,7 +1792,7 @@ return results
     });
 
     this.redis.defineCommand("acknowledgeMessage", {
-      numberOfKeys: 6,
+      numberOfKeys: 7,
       lua: `
 -- Keys:
 local masterQueueKey = KEYS[1]
@@ -1776,10 +1801,13 @@ local messageQueueKey = KEYS[3]
 local queueCurrentConcurrencyKey = KEYS[4]
 local envCurrentConcurrencyKey = KEYS[5]
 local envQueueKey = KEYS[6]
+local workerQueueKey = KEYS[7]
 
 -- Args:
 local messageId = ARGV[1]
 local messageQueueName = ARGV[2]
+local messageKeyValue = ARGV[3]
+local removeFromWorkerQueue = ARGV[4]
 
 -- Remove the message from the message key
 redis.call('DEL', messageKey)
@@ -1799,6 +1827,11 @@ end
 -- Update the concurrency keys
 redis.call('SREM', queueCurrentConcurrencyKey, messageId)
 redis.call('SREM', envCurrentConcurrencyKey, messageId)
+
+-- Remove the message from the worker queue
+if removeFromWorkerQueue == '1' then
+  redis.call('LREM', workerQueueKey, 0, messageKeyValue)
+end
 `,
     });
 
@@ -2025,8 +2058,11 @@ declare module "@internal/redis" {
       concurrencyKey: string,
       envConcurrencyKey: string,
       envQueueKey: string,
+      workerQueueKey: string,
       messageId: string,
       messageQueueName: string,
+      messageKeyValue: string,
+      removeFromWorkerQueue: string,
       callback?: Callback<void>
     ): Result<void, Context>;