Merge remote-tracking branch 'origin/main' into feat/set-machine-at-trigger-time

Author: nicktrn

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

@@ -74,6 +74,11 @@ export class DevQueueConsumer {
       return;
     }
 
+    logger.debug("[DevQueueConsumer] Deprecating background worker", {
+      backgroundWorker: backgroundWorker.id,
+      env: this.env.id,
+    });
+
     this._deprecatedWorkers.set(id, backgroundWorker);
     this._backgroundWorkers.delete(id);
   }
@@ -96,9 +101,10 @@ export class DevQueueConsumer {
 
     this._backgroundWorkers.set(backgroundWorker.id, backgroundWorker);
 
-    logger.debug("Registered background worker", {
+    logger.debug("[DevQueueConsumer] Registered background worker", {
       backgroundWorker: backgroundWorker.id,
       inProgressRuns,
+      env: this.env.id,
     });
 
     const subscriber = await devPubSub.subscribe(`backgroundWorker:${backgroundWorker.id}:*`);
@@ -138,6 +144,7 @@ export class DevQueueConsumer {
     logger.debug("[DevQueueConsumer] taskAttemptCompleted()", {
       taskRunCompletion: completion,
       execution,
+      env: this.env.id,
     });
 
     const service = new CompleteAttemptService();
@@ -151,7 +158,7 @@ export class DevQueueConsumer {
   public async taskRunFailed(workerId: string, completion: TaskRunFailedExecutionResult) {
     this._taskFailures++;
 
-    logger.debug("[DevQueueConsumer] taskRunFailed()", { completion });
+    logger.debug("[DevQueueConsumer] taskRunFailed()", { completion, env: this.env.id });
 
     this._inProgressRuns.delete(completion.id);
 
@@ -188,7 +195,7 @@ export class DevQueueConsumer {
       return;
     }
 
-    logger.debug("Stopping dev queue consumer", { env: this.env });
+    logger.debug("[DevQueueConsumer] Stopping dev queue consumer", { env: this.env });
 
     this._enabled = false;
 
@@ -322,20 +329,30 @@ export class DevQueueConsumer {
         env: this.env,
       });
 
-      await marqs?.acknowledgeMessage(message.messageId);
+      await marqs?.acknowledgeMessage(
+        message.messageId,
+        "Failed to parse message.data with MessageBody schema in DevQueueConsumer"
+      );
 
       setTimeout(() => this.#doWork(), 100);
       return;
     }
 
-    const existingTaskRun = await prisma.taskRun.findUnique({
+    const existingTaskRun = await prisma.taskRun.findFirst({
       where: {
         id: message.messageId,
       },
     });
 
     if (!existingTaskRun) {
-      await marqs?.acknowledgeMessage(message.messageId);
+      logger.debug("Failed to find existing task run, acking", {
+        messageId: message.messageId,
+      });
+
+      await marqs?.acknowledgeMessage(
+        message.messageId,
+        "Failed to find task run in DevQueueConsumer"
+      );
       setTimeout(() => this.#doWork(), 100);
       return;
     }
@@ -346,7 +363,18 @@ export class DevQueueConsumer {
       : this.#getLatestBackgroundWorker();
 
     if (!backgroundWorker) {
-      await marqs?.acknowledgeMessage(message.messageId);
+      logger.debug("Failed to find background worker, acking", {
+        messageId: message.messageId,
+        lockedToVersionId: existingTaskRun.lockedToVersionId,
+        deprecatedWorkers: Array.from(this._deprecatedWorkers.keys()),
+        backgroundWorkers: Array.from(this._backgroundWorkers.keys()),
+        latestWorker: this.#getLatestBackgroundWorker(),
+      });
+
+      await marqs?.acknowledgeMessage(
+        message.messageId,
+        "Failed to find background worker in DevQueueConsumer"
+      );
       setTimeout(() => this.#doWork(), 100);
       return;
     }
@@ -363,7 +391,10 @@ export class DevQueueConsumer {
         taskSlugs: backgroundWorker.tasks.map((task) => task.slug),
       });
 
-      await marqs?.acknowledgeMessage(message.messageId);
+      await marqs?.acknowledgeMessage(
+        message.messageId,
+        "No matching background task found in DevQueueConsumer"
+      );
 
       setTimeout(() => this.#doWork(), 100);
       return;
@@ -397,7 +428,10 @@ export class DevQueueConsumer {
         messageId: message.messageId,
       });
 
-      await marqs?.acknowledgeMessage(message.messageId);
+      await marqs?.acknowledgeMessage(
+        message.messageId,
+        "Failed to lock task run in DevQueueConsumer"
+      );
 
       setTimeout(() => this.#doWork(), 100);
       return;
@@ -496,7 +530,10 @@ export class DevQueueConsumer {
         messageId: message.messageId,
         backgroundWorker,
       });
-      await marqs?.acknowledgeMessage(message.messageId);
+      await marqs?.acknowledgeMessage(
+        message.messageId,
+        "Non-lazy attempts are no longer supported in DevQueueConsumer"
+      );
 
       setTimeout(() => this.#doWork(), 100);
     }

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

@@ -291,6 +291,8 @@ export class MarQS {
             orgConcurrencyKey: this.keys.orgCurrentConcurrencyKeyFromQueue(messageQueue),
             messageId: messageData.messageId,
           });
+
+          return;
         }
 
         await this.options.visibilityTimeoutStrategy.heartbeat(
@@ -411,7 +413,7 @@ export class MarQS {
     );
   }
 
-  public async acknowledgeMessage(messageId: string) {
+  public async acknowledgeMessage(messageId: string, reason: string = "unknown") {
     return this.#trace(
       "acknowledgeMessage",
       async (span) => {
@@ -421,6 +423,7 @@ export class MarQS {
           logger.log(`[${this.name}].acknowledgeMessage() message not found`, {
             messageId,
             service: this.name,
+            reason,
           });
           return;
         }
@@ -430,6 +433,7 @@ export class MarQS {
           [SemanticAttributes.MESSAGE_ID]: message.messageId,
           [SemanticAttributes.CONCURRENCY_KEY]: message.concurrencyKey,
           [SemanticAttributes.PARENT_QUEUE]: message.parentQueue,
+          ["marqs.reason"]: reason,
         });
 
         await this.options.visibilityTimeoutStrategy.cancelHeartbeat(messageId);

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

@@ -978,7 +978,7 @@ export class SharedQueueConsumer {
   }
 
   async #ackAndDoMoreWork(messageId: string, intervalInMs?: number) {
-    await marqs?.acknowledgeMessage(messageId);
+    await marqs?.acknowledgeMessage(messageId, "Acking and doing more work in SharedQueueConsumer");
     this.#doMoreWork(intervalInMs);
   }
 

apps/webapp/app/v3/requeueTaskRun.server.ts

@@ -71,7 +71,10 @@ export class RequeueTaskRunService extends BaseService {
       case "WAITING_FOR_DEPLOY": {
         logger.debug("[RequeueTaskRunService] Removing task run from queue", { taskRun });
 
-        await marqs?.acknowledgeMessage(taskRun.id);
+        await marqs?.acknowledgeMessage(
+          taskRun.id,
+          "Run is either DELAYED or WAITING_FOR_DEPLOY so we cannot requeue it in RequeueTaskRunService"
+        );
 
         break;
       }
@@ -93,7 +96,10 @@ export class RequeueTaskRunService extends BaseService {
       case "CANCELED": {
         logger.debug("[RequeueTaskRunService] Task run is completed", { taskRun });
 
-        await marqs?.acknowledgeMessage(taskRun.id);
+        await marqs?.acknowledgeMessage(
+          taskRun.id,
+          "Task run is already completed in RequeueTaskRunService"
+        );
 
         try {
           if (taskRun.runtimeEnvironment.type === "DEVELOPMENT") {

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

@@ -307,7 +307,10 @@ export class CreateCheckpointService extends BaseService {
         checkpointId: checkpoint.id,
         params,
       });
-      await marqs?.acknowledgeMessage(attempt.taskRunId);
+      await marqs?.acknowledgeMessage(
+        attempt.taskRunId,
+        "No checkpoint event in CreateCheckpointService"
+      );
 
       return {
         success: false,

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

@@ -61,7 +61,7 @@ export class FinalizeTaskRunService extends BaseService {
       expiredAt,
       completedAt,
     });
-    await marqs?.acknowledgeMessage(id);
+    await marqs?.acknowledgeMessage(id, "FinalTaskRunService call");
 
     logger.debug("Finalizing run updating run status", {
       id,

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

@@ -172,7 +172,10 @@ export class ResumeAttemptService extends BaseService {
 
       if (!completedAttempt) {
         this._logger.error("Completed attempt not found", { completedAttemptId });
-        await marqs?.acknowledgeMessage(attempt.taskRunId);
+        await marqs?.acknowledgeMessage(
+          attempt.taskRunId,
+          "Cannot find completed attempt in ResumeAttemptService"
+        );
         return;
       }
 
@@ -186,7 +189,10 @@ export class ResumeAttemptService extends BaseService {
 
       if (!resumePayload) {
         logger.error("Failed to get resume payload");
-        await marqs?.acknowledgeMessage(attempt.taskRunId);
+        await marqs?.acknowledgeMessage(
+          attempt.taskRunId,
+          "Failed to get resume payload in ResumeAttemptService"
+        );
         return;
       }
 

scripts/unpack-worker.js

@@ -0,0 +1,68 @@
+const fs = require("fs");
+const zlib = require("zlib");
+const path = require("path");
+
+// Get the file paths from command line arguments
+let [jsonFilePath, destDir] = process.argv.slice(2);
+
+if (!jsonFilePath || !destDir) {
+  console.error("Usage: node script.js <json-file-path> <destination-directory>");
+  process.exit(1);
+}
+
+// Function to decompress the content
+function decompressContent(base64Encoded) {
+  // Decode base64 string to buffer
+  const compressedData = Buffer.from(base64Encoded, "base64");
+
+  // Decompress the data
+  const decompressedData = zlib.inflateSync(compressedData);
+
+  // Convert buffer to string
+  return decompressedData.toString();
+}
+
+try {
+  // Read and parse the JSON file
+  const jsonContent = fs.readFileSync(jsonFilePath, "utf8");
+
+  const data = JSON.parse(jsonContent)[0];
+
+  console.log(data);
+
+  const id = data.id;
+
+  console.log(`Extracting files for: ${id} to ${destDir}`);
+
+  destDir = path.join(destDir, id);
+
+  console.log(`Extracting files to: ${destDir}`);
+
+  // Create the destination directory if it doesn't exist
+  fs.mkdirSync(destDir, { recursive: true });
+
+  // Process each item in the array
+  const sourceFiles = data.metadata.sourceFiles;
+
+  sourceFiles.forEach((file) => {
+    // Decompress the contents
+    const decompressedContent = decompressContent(file.contents);
+
+    // Combine destination directory with file path
+    const fullPath = path.join(destDir, file.filePath);
+
+    // Create directory structure if it doesn't exist
+    const dirPath = path.dirname(fullPath);
+    fs.mkdirSync(dirPath, { recursive: true });
+
+    // Write the decompressed content to the file
+    fs.writeFileSync(fullPath, decompressedContent);
+
+    console.log(`Created file: ${fullPath}`);
+  });
+
+  console.log(`\nAll files have been extracted to: ${destDir}`);
+} catch (error) {
+  console.error(error);
+  process.exit(1);
+}