diff --git a/apps/coordinator/src/index.ts b/apps/coordinator/src/index.ts index dff8c38b3..e49aa04eb 100644 --- a/apps/coordinator/src/index.ts +++ b/apps/coordinator/src/index.ts @@ -157,16 +157,18 @@ class Checkpointer { return this.#abortControllers.has(runId); } - cancelCheckpoint(runId: string) { + cancelCheckpoint(runId: string): boolean { const controller = this.#abortControllers.get(runId); if (!controller) { logger.debug("Nothing to cancel", { runId }); - return; + return false; } controller.abort("cancelCheckpointing()"); this.#abortControllers.delete(runId); + + return true; } async #checkpointAndPush({ @@ -725,10 +727,18 @@ class TaskCoordinator { checkpointable.resolve(); }); - socket.on("CANCEL_CHECKPOINT", async (message) => { + socket.on("CANCEL_CHECKPOINT", async (message, callback) => { logger.log("[CANCEL_CHECKPOINT]", message); - this.#cancelCheckpoint(socket.data.runId); + if (message.version === "v1") { + this.#cancelCheckpoint(socket.data.runId); + // v1 has no callback + return; + } + + const checkpointCanceled = this.#cancelCheckpoint(socket.data.runId); + + callback({ version: "v2", checkpointCanceled }); }); socket.on("WAIT_FOR_DURATION", async (message, callback) => { @@ -933,7 +943,9 @@ class TaskCoordinator { } // Cancel checkpointing procedure - this.#checkpointer.cancelCheckpoint(runId); + const checkpointCanceled = this.#checkpointer.cancelCheckpoint(runId); + + return checkpointCanceled; } #createHttpServer() { diff --git a/packages/cli-v3/src/workers/prod/backgroundWorker.ts b/packages/cli-v3/src/workers/prod/backgroundWorker.ts index f5c5e85e4..fe3a2ed7a 100644 --- a/packages/cli-v3/src/workers/prod/backgroundWorker.ts +++ b/packages/cli-v3/src/workers/prod/backgroundWorker.ts @@ -13,6 +13,7 @@ import { TaskRunExecution, TaskRunExecutionPayload, TaskRunExecutionResult, + WaitReason, correctErrorStackTrace, } from "@trigger.dev/core/v3"; import { ZodIpcConnection } from "@trigger.dev/core/v3/zodIpc"; @@ -68,8 +69,10 @@ export class ProdBackgroundWorker { > = new Evt(); public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>(); + public checkpointCanceledNotification = Evt.create<{ checkpointCanceled: boolean }>(); + public onReadyForCheckpoint = Evt.create<{ version?: "v1" }>(); - public onCancelCheckpoint = Evt.create<{ version?: "v1" }>(); + public onCancelCheckpoint = Evt.create<{ version?: "v1" | "v2"; reason?: WaitReason }>(); private _onClose: Evt = new Evt(); @@ -251,6 +254,9 @@ export class ProdBackgroundWorker { this.preCheckpointNotification.attach((message) => { taskRunProcess.preCheckpointNotification.post(message); }); + this.checkpointCanceledNotification.attach((message) => { + taskRunProcess.checkpointCanceledNotification.post(message); + }); await taskRunProcess.initialize(); @@ -377,8 +383,10 @@ class TaskRunProcess { > = new Evt(); public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>(); + public checkpointCanceledNotification = Evt.create<{ checkpointCanceled: boolean }>(); + public onReadyForCheckpoint = Evt.create<{ version?: "v1" }>(); - public onCancelCheckpoint = Evt.create<{ version?: "v1" }>(); + public onCancelCheckpoint = Evt.create<{ version?: "v1" | "v2"; reason?: WaitReason }>(); constructor( private execution: ProdTaskRunExecution, @@ -438,13 +446,26 @@ class TaskRunProcess { this.onWaitForBatch.post(message); }, WAIT_FOR_DURATION: async (message) => { + // Post to coordinator this.onWaitForDuration.post(message); - // The coordinator will let us know if a checkpoint is about to happen - // We then pass this back down to the runtime in the child process - const { willCheckpointAndRestore } = await this.preCheckpointNotification.waitFor(); + try { + // ..and wait for response + const { willCheckpointAndRestore } = await this.preCheckpointNotification.waitFor( + 30_000 + ); - return { willCheckpointAndRestore }; + return { + willCheckpointAndRestore, + }; + } catch (error) { + console.error("Error while waiting for pre-checkpoint notification", error); + + // Assume we won't get checkpointed + return { + willCheckpointAndRestore: false, + }; + } }, WAIT_FOR_TASK: async (message) => { this.onWaitForTask.post(message); @@ -453,7 +474,30 @@ class TaskRunProcess { this.onReadyForCheckpoint.post(message); }, CANCEL_CHECKPOINT: async (message) => { + const version = "v2"; + + // Post to coordinator this.onCancelCheckpoint.post(message); + + try { + // ..and wait for response + const { checkpointCanceled } = await this.checkpointCanceledNotification.waitFor( + 30_000 + ); + + return { + version, + checkpointCanceled, + }; + } catch (error) { + console.error("Error while waiting for checkpoint cancellation", error); + + // Assume it's been canceled + return { + version, + checkpointCanceled: true, + }; + } }, }, }); diff --git a/packages/cli-v3/src/workers/prod/entry-point.ts b/packages/cli-v3/src/workers/prod/entry-point.ts index 2c75b8eae..2daa37dfd 100644 --- a/packages/cli-v3/src/workers/prod/entry-point.ts +++ b/packages/cli-v3/src/workers/prod/entry-point.ts @@ -43,7 +43,6 @@ class ProdWorker { private attemptFriendlyId?: string; private nextResumeAfter?: WaitReason; - private waitForPostStart = false; #httpPort: number; #backgroundWorker: ProdBackgroundWorker; @@ -82,18 +81,27 @@ class ProdWorker { this.#coordinatorSocket.socket.emit("READY_FOR_CHECKPOINT", { version: "v1" }); }); + // Currently, this is only used for duration waits. Might need adjusting for other use cases. this.#backgroundWorker.onCancelCheckpoint.attach(async (message) => { - logger.log("onCancelCheckpoint() clearing paused state, don't wait for post start hook", { - paused: this.paused, - nextResumeAfter: this.nextResumeAfter, - waitForPostStart: this.waitForPostStart, - }); + logger.log("onCancelCheckpoint()", { message }); - this.paused = false; - this.nextResumeAfter = undefined; - this.waitForPostStart = false; + const { checkpointCanceled } = await this.#coordinatorSocket.socket.emitWithAck( + "CANCEL_CHECKPOINT", + { + version: "v2", + reason: message.reason, + } + ); - this.#coordinatorSocket.socket.emit("CANCEL_CHECKPOINT", { version: "v1" }); + if (checkpointCanceled) { + if (message.reason === "WAIT_FOR_DURATION") { + // Worker will resume immediately + this.paused = false; + this.nextResumeAfter = undefined; + } + } + + this.#backgroundWorker.checkpointCanceledNotification.post({ checkpointCanceled }); }); this.#backgroundWorker.onWaitForDuration.attach(async (message) => { @@ -182,10 +190,6 @@ class ProdWorker { } async #reconnect(isPostStart = false, reconnectImmediately = false) { - if (isPostStart) { - this.waitForPostStart = false; - } - this.#coordinatorSocket.close(); if (!reconnectImmediately) { @@ -224,7 +228,6 @@ class ProdWorker { if (willCheckpointAndRestore) { this.paused = true; this.nextResumeAfter = reason; - this.waitForPostStart = true; } } @@ -244,7 +247,6 @@ class ProdWorker { this.attemptFriendlyId = undefined; if (willCheckpointAndRestore) { - this.waitForPostStart = true; this.#coordinatorSocket.socket.emit("READY_FOR_CHECKPOINT", { version: "v1" }); return; } @@ -267,6 +269,7 @@ class ProdWorker { return headers; } + // FIXME: If the the worker can't connect for a while, this runs MANY times - it should only run once #createCoordinatorSocket(host: string) { const extraHeaders = this.#returnValidatedExtraHeaders({ "x-machine-name": MACHINE_NAME, @@ -423,8 +426,22 @@ class ProdWorker { }, }, onConnection: async (socket, handler, sender, logger) => { - if (this.waitForPostStart) { - logger.log("skip connection handler, waiting for post start hook"); + if (this.paused) { + if (!this.nextResumeAfter) { + return; + } + + if (!this.attemptFriendlyId) { + logger.error("Missing friendly ID"); + return; + } + + socket.emit("READY_FOR_RESUME", { + version: "v1", + attemptFriendlyId: this.attemptFriendlyId, + type: this.nextResumeAfter, + }); + return; } @@ -519,30 +536,6 @@ class ProdWorker { } } - if (this.paused) { - if (!this.nextResumeAfter) { - return; - } - - if (!this.attemptFriendlyId) { - logger.error("Missing friendly ID"); - return; - } - - if (this.nextResumeAfter === "WAIT_FOR_DURATION") { - this.#resumeAfterDuration(); - return; - } - - socket.emit("READY_FOR_RESUME", { - version: "v1", - attemptFriendlyId: this.attemptFriendlyId, - type: this.nextResumeAfter, - }); - - return; - } - if (this.executing) { return; } @@ -587,7 +580,8 @@ class ProdWorker { case "/status": { return reply.json({ executing: this.executing, - pause: this.paused, + paused: this.paused, + completed: this.completed.size, nextResumeAfter: this.nextResumeAfter, }); } diff --git a/packages/cli-v3/src/workers/prod/worker-facade.ts b/packages/cli-v3/src/workers/prod/worker-facade.ts index 8c9885753..bba267b8f 100644 --- a/packages/cli-v3/src/workers/prod/worker-facade.ts +++ b/packages/cli-v3/src/workers/prod/worker-facade.ts @@ -174,7 +174,7 @@ const zodIpc = new ZodIpcConnection({ prodRuntimeManager.resumeTask(completion, execution); }, WAIT_COMPLETED_NOTIFICATION: async () => { - prodRuntimeManager.resumeAfterRestore(); + prodRuntimeManager.resumeAfterDuration(); }, CLEANUP: async ({ flush, kill }, sender) => { if (kill) { diff --git a/packages/core/src/v3/runtime/prodRuntimeManager.ts b/packages/core/src/v3/runtime/prodRuntimeManager.ts index 297ef0c6f..fdf9373af 100644 --- a/packages/core/src/v3/runtime/prodRuntimeManager.ts +++ b/packages/core/src/v3/runtime/prodRuntimeManager.ts @@ -1,4 +1,5 @@ import { clock } from "../clock-api"; +import { logger } from "../logger-api"; import { BatchTaskRunExecutionResult, ProdChildToWorkerMessages, @@ -23,7 +24,9 @@ export class ProdRuntimeManager implements RuntimeManager { { resolve: (value: BatchTaskRunExecutionResult) => void; reject: (err?: any) => void } > = new Map(); - _waitForRestore: { resolve: (value: "restore") => void; reject: (err?: any) => void } | undefined; + _waitForDuration: + | { resolve: (value: "external") => void; reject: (err?: any) => void } + | undefined; constructor( private ipc: ZodIpcConnection< @@ -40,15 +43,15 @@ export class ProdRuntimeManager implements RuntimeManager { async waitForDuration(ms: number): Promise { const now = Date.now(); - const resolveAfterDuration = unboundedTimeout(ms, "duration" as const); + const internalTimeout = unboundedTimeout(ms, "internal" as const); if (ms <= this.waitThresholdInMs) { - await resolveAfterDuration; + await internalTimeout; return; } - const waitForRestore = new Promise<"restore">((resolve, reject) => { - this._waitForRestore = { resolve, reject }; + const externalResume = new Promise<"external">((resolve, reject) => { + this._waitForDuration = { resolve, reject }; }); const { willCheckpointAndRestore } = await this.ipc.sendWithAck("WAIT_FOR_DURATION", { @@ -57,29 +60,101 @@ export class ProdRuntimeManager implements RuntimeManager { }); if (!willCheckpointAndRestore) { - await resolveAfterDuration; + await internalTimeout; return; } + const getTimes = () => { + return { + date: Date.now(), // ms + clock: clock.preciseNow()[0] * 1000, // seconds + perf: performance.now(), // ms + }; + }; + + const preWait = getTimes(); + this.ipc.send("READY_FOR_CHECKPOINT", {}); - // Don't wait for checkpoint beyond the requested wait duration - await Promise.race([waitForRestore, resolveAfterDuration]); + await internalTimeout; - // The coordinator can then cancel any in-progress checkpoints - this.ipc.send("CANCEL_CHECKPOINT", {}); - } - - resumeAfterRestore(): void { - if (!this._waitForRestore) { - return; - } + // The internal timer is up, let's check for missing time + const postWait = getTimes(); // Resets the clock to the current time clock.reset(); - this._waitForRestore.resolve("restore"); - this._waitForRestore = undefined; + const postReset = getTimes(); + + const diffs = { + t1: { + date: postWait.date - preWait.date, + clock: postWait.clock - preWait.clock, + perf: postWait.perf - preWait.perf, + }, + t2: { + date: postReset.date - postWait.date, + clock: postReset.clock - postWait.clock, + perf: postReset.perf - postWait.perf, + }, + }; + + console.log({ + preWait, + postWait, + postReset, + diffs, + }); + + logger.debug("diffs", { + preWait, + postWait, + postReset, + diffs, + }); + + // The coordinator should cancel any in-progress checkpoints + const { checkpointCanceled, version } = await this.ipc.sendWithAck("CANCEL_CHECKPOINT", { + version: "v2", + reason: "WAIT_FOR_DURATION", + }); + + console.log({ checkpointCanceled, version }); + logger.debug("cancel checkpoint", { checkpointCanceled, version }); + + if (checkpointCanceled) { + // There won't be a checkpoint or external resume and we've already completed our internal timeout + return; + } + + console.log("Waiting for external resume"); + + // No checkpoint was canceled, so we were checkpointed. We need to wait for the external resume message. + await externalResume; + + console.log("Done waiting for external resume"); + } + + resumeAfterDuration(): void { + if (!this._waitForDuration) { + return; + } + + process.stdout.write("pre"); + process.stdout.write(JSON.stringify(clock.preciseNow())); + + console.log("pre", clock.preciseNow()); + + // Resets the clock to the current time + clock.reset(); + + console.log("post", clock.preciseNow()); + + process.stdout.write("post"); + process.stdout.write(JSON.stringify(clock.preciseNow())); + + this._waitForDuration.resolve("external"); + this._waitForDuration = undefined; } async waitUntil(date: Date): Promise { diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts index 78951be69..32c66c808 100644 --- a/packages/core/src/v3/schemas/messages.ts +++ b/packages/core/src/v3/schemas/messages.ts @@ -191,10 +191,20 @@ export const ProdChildToWorkerMessages = { }), }, CANCEL_CHECKPOINT: { - message: z.object({ - version: z.enum(["v1", "v2"]).default("v2"), - }), + message: z + .discriminatedUnion("version", [ + z.object({ + version: z.literal("v1"), + }), + z.object({ + version: z.literal("v2"), + reason: WaitReason.optional(), + }), + ]) + .default({ version: "v1" }), callback: z.object({ + // TODO: Figure out how best to handle callback schema parsing in zod IPC + version: z.literal("v2") /* .default("v2") */, checkpointCanceled: z.boolean(), reason: WaitReason.optional(), }), @@ -597,10 +607,19 @@ export const ProdWorkerToCoordinatorMessages = { }), }, CANCEL_CHECKPOINT: { - message: z.object({ - version: z.enum(["v1", "v2"]).default("v2"), - }), + message: z + .discriminatedUnion("version", [ + z.object({ + version: z.literal("v1"), + }), + z.object({ + version: z.literal("v2"), + reason: WaitReason.optional(), + }), + ]) + .default({ version: "v1" }), callback: z.object({ + version: z.literal("v2").default("v2"), checkpointCanceled: z.boolean(), reason: WaitReason.optional(), }),