Fixes an issue with the shared queue consumer relying on the queue name to get the env id

Author: ericallam

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

@@ -251,35 +251,12 @@ export class SharedQueueConsumer {
 
     logger.log("dequeueMessageInSharedQueue()", { queueMessage: message });
 
-    const envId = this.#envIdFromQueue(message.queue);
-
-    const environment = await prisma.runtimeEnvironment.findUnique({
-      include: {
-        organization: true,
-        project: true,
-      },
-      where: {
-        id: envId,
-      },
-    });
-
-    if (!environment) {
-      logger.error("Environment not found", {
-        queueMessage: message.data,
-        envId,
-      });
-
-      this.#ackAndDoMoreWork(message.messageId);
-      return;
-    }
-
     const messageBody = MessageBody.safeParse(message.data);
 
     if (!messageBody.success) {
       logger.error("Failed to parse message", {
         queueMessage: message.data,
         error: messageBody.error,
-        env: environment,
       });
 
       this.#ackAndDoMoreWork(message.messageId);
@@ -381,6 +358,7 @@ export class SharedQueueConsumer {
             lockedById: backgroundTask.id,
           },
           include: {
+            runtimeEnvironment: true,
             attempts: {
               take: 1,
               orderBy: { number: "desc" },
@@ -411,7 +389,7 @@ export class SharedQueueConsumer {
         const queue = await prisma.taskQueue.findUnique({
           where: {
             runtimeEnvironmentId_name: {
-              runtimeEnvironmentId: environment.id,
+              runtimeEnvironmentId: lockedTaskRun.runtimeEnvironmentId,
               name: lockedTaskRun.queue,
             },
           },
@@ -437,7 +415,7 @@ export class SharedQueueConsumer {
             backgroundWorkerTaskId: backgroundTask.id,
             status: "PENDING" as const,
             queueId: queue.id,
-            runtimeEnvironmentId: environment.id,
+            runtimeEnvironmentId: lockedTaskRun.runtimeEnvironmentId,
           },
           include: {
             backgroundWorkerTask: true,
@@ -493,10 +471,10 @@ export class SharedQueueConsumer {
                 machine: machine.data,
                 // identifiers
                 id: taskRunAttempt.id,
-                envId: environment.id,
-                envType: environment.type,
-                orgId: environment.organizationId,
-                projectId: environment.projectId,
+                envId: lockedTaskRun.runtimeEnvironment.id,
+                envType: lockedTaskRun.runtimeEnvironment.type,
+                orgId: lockedTaskRun.runtimeEnvironment.organizationId,
+                projectId: lockedTaskRun.runtimeEnvironment.projectId,
                 runId: taskRunAttempt.taskRunId,
               },
             });
@@ -625,7 +603,7 @@ export class SharedQueueConsumer {
         const queue = await prisma.taskQueue.findUnique({
           where: {
             runtimeEnvironmentId_name: {
-              runtimeEnvironmentId: environment.id,
+              runtimeEnvironmentId: resumableAttempt.runtimeEnvironmentId,
               name: resumableRun.queue,
             },
           },
@@ -754,10 +732,6 @@ export class SharedQueueConsumer {
     return;
   }
 
-  #envIdFromQueue(queueName: string) {
-    return queueName.split(":")[1];
-  }
-
   #doMoreWork(intervalInMs = this._options.interval) {
     setTimeout(() => this.#doWork(), intervalInMs);
   }