github-featured/triggerdotdev--trigger.dev
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

completely switch to platform-led duration wait resumes

Browse Source
This commit is contained in:
nicktrn
2024-04-25 14:11:10 +01:00
parent 35efb088db
commit 77a8a645ed
6 changed files with 223 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/coordinator/src/index.ts
+17 -5
View File
@@ -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() {
packages/cli-v3/src/workers/prod/backgroundWorker.ts
+50 -6
View File
@@ -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<void> = 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,
};
}
},
},
});
packages/cli-v3/src/workers/prod/entry-point.ts
+37 -43
View File
@@ -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,
});
}
packages/cli-v3/src/workers/prod/worker-facade.ts
+1 -1
View File
@@ -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) {
packages/core/src/v3/runtime/prodRuntimeManager.ts
+93 -18
View File
@@ -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<void> {
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<void> {
packages/core/src/v3/schemas/messages.ts
+25 -6
View File
@@ -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(),
}),