improve wait accuracy

apps/coordinator/src/index.ts

@@ -505,6 +505,12 @@ class TaskCoordinator {
           }
 
           confirmCompletion({ didCheckpoint: true, shouldExit: false, checkpoint });
+
+          if (!checkpoint.docker) {
+            socket.emit("REQUEST_EXIT", {
+              version: "v1",
+            });
+          }
         });
 
         socket.on("WAIT_FOR_DURATION", async (message, callback) => {
@@ -550,6 +556,7 @@ class TaskCoordinator {
             reason: {
               type: "WAIT_FOR_DURATION",
               ms: message.ms,
+              now: message.now,
             },
           });
         });

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

@@ -99,7 +99,7 @@ export class CreateCheckpointService {
         await marqs?.replaceMessage(
           attempt.taskRunId,
           { type: "RESUME_AFTER_DURATION", resumableAttemptId: attempt.id },
-          Date.now() + params.reason.ms
+          params.reason.now + params.reason.ms
         );
         break;
       }

packages/cli-v3/src/workers/prod/backgroundWorker.ts

@@ -17,7 +17,6 @@ import {
 } from "@trigger.dev/core/v3";
 import { Evt } from "evt";
 import { ChildProcess, fork } from "node:child_process";
-import { safeDeleteFileSync } from "../../utilities/fileSystem";
 import { UncaughtExceptionError } from "../common/errors";
 
 class UnexpectedExitError extends Error {
@@ -56,7 +55,7 @@ export class ProdBackgroundWorker {
 
   public onTaskHeartbeat: Evt<string> = new Evt();
 
-  public onWaitForDuration: Evt<{ version?: "v1"; ms: number }> = new Evt();
+  public onWaitForDuration: Evt<{ version?: "v1"; ms: number; now: number }> = new Evt();
   public onWaitForTask: Evt<{ version?: "v1"; id: string }> = new Evt();
   public onWaitForBatch: Evt<{ version?: "v1"; id: string; runs: string[] }> = new Evt();
 
@@ -343,7 +342,7 @@ class TaskRunProcess {
   public onExit: Evt<number> = new Evt();
 
   public onWaitForBatch: Evt<{ version?: "v1"; id: string; runs: string[] }> = new Evt();
-  public onWaitForDuration: Evt<{ version?: "v1"; ms: number }> = new Evt();
+  public onWaitForDuration: Evt<{ version?: "v1"; ms: number; now: number }> = new Evt();
   public onWaitForTask: Evt<{ version?: "v1"; id: string }> = new Evt();
 
   public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>();

packages/core/src/v3/runtime/prodRuntimeManager.ts

@@ -49,6 +49,8 @@ export class ProdRuntimeManager implements RuntimeManager {
   async waitForDuration(ms: number): Promise<void> {
     let timeout: NodeJS.Timeout | undefined;
 
+    const now = Date.now();
+
     const resolveAfterDuration = new Promise((resolve) => {
       timeout = setTimeout(resolve, ms);
     });
@@ -63,7 +65,10 @@ export class ProdRuntimeManager implements RuntimeManager {
     });
 
     // There is a slight delay before actually checkpointing, so this has a chance to return
-    const { willCheckpointAndRestore } = await this.ipc.sendWithAck("WAIT_FOR_DURATION", { ms });
+    const { willCheckpointAndRestore } = await this.ipc.sendWithAck("WAIT_FOR_DURATION", {
+      ms,
+      now,
+    });
 
     if (!willCheckpointAndRestore) {
       await resolveAfterDuration;

packages/core/src/v3/schemas/messages.ts

@@ -262,6 +262,7 @@ export const ProdChildToWorkerMessages = {
     message: z.object({
       version: z.literal("v1").default("v1"),
       ms: z.number(),
+      now: z.number(),
     }),
     callback: z.object({
       willCheckpointAndRestore: z.boolean(),

packages/core/src/v3/schemas/schemas.ts

@@ -204,6 +204,7 @@ export const CoordinatorToPlatformMessages = {
         z.object({
           type: z.literal("WAIT_FOR_DURATION"),
           ms: z.number(),
+          now: z.number(),
         }),
         z.object({
           type: z.literal("WAIT_FOR_BATCH"),
@@ -353,6 +354,7 @@ export const ProdWorkerToCoordinatorMessages = {
     message: z.object({
       version: z.literal("v1").default("v1"),
       ms: z.number(),
+      now: z.number(),
     }),
     callback: z.object({
       willCheckpointAndRestore: z.boolean(),

packages/core/src/v3/zodSocket.ts

@@ -155,13 +155,18 @@ export class ZodSocketMessageHandler<TRPCCatalog extends ZodSocketMessageCatalog
 
         let ack;
 
-        // FIXME: this only works if the message doesn't have genuine payload prop
-        if ("payload" in message) {
-          ack = await this.handleMessage({ type: eventName, ...message });
-        } else {
-          // Handle messages not sent by ZodMessageSender
-          const { version, ...payload } = message;
-          ack = await this.handleMessage({ type: eventName, version, payload });
+        try {
+          // FIXME: this only works if the message doesn't have genuine payload prop
+          if ("payload" in message) {
+            ack = await this.handleMessage({ type: eventName, ...message });
+          } else {
+            // Handle messages not sent by ZodMessageSender
+            const { version, ...payload } = message;
+            ack = await this.handleMessage({ type: eventName, version, payload });
+          }
+        } catch (error) {
+          log.error("Error while handling message", { error });
+          return;
         }
 
         if (callback && typeof callback === "function") {