v3: checkpoint and restore with sim for local dev (#933)
* zod ipc connection with acks * pass sender through to ipc handlers * bits and bobs * add host networking back in * disable verbose logs * restore after tasks and batches.. almost * restore and resume all the things * fix for systems without checkpoint support * Get deployment indexing errors to work with the new ZodIpc stuff --------- Co-authored-by: Eric Allam <eallam@icloud.com>
This commit is contained in:
+156
-42
@@ -30,25 +30,44 @@ const PLATFORM_SECRET = process.env.PLATFORM_SECRET || "coordinator-secret";
|
||||
|
||||
const logger = new SimpleLogger(`[${NODE_NAME}]`);
|
||||
|
||||
type CheckpointerInitializeReturn = {
|
||||
canCheckpoint: boolean;
|
||||
willSimulate: boolean;
|
||||
};
|
||||
|
||||
class Checkpointer {
|
||||
#initialized = false;
|
||||
#canCheckpoint = false;
|
||||
#dockerMode = true;
|
||||
#dockerMode = !process.env.KUBERNETES_PORT;
|
||||
|
||||
#logger = new SimpleLogger("[checkptr]");
|
||||
|
||||
async initialize() {
|
||||
constructor(private opts = { forceSimulate: false }) {}
|
||||
|
||||
async initialize(): Promise<CheckpointerInitializeReturn> {
|
||||
if (this.#initialized) {
|
||||
return;
|
||||
return this.#getInitializeReturn();
|
||||
}
|
||||
|
||||
this.#logger.log(`${this.#dockerMode ? "Docker" : "Kubernetes"} mode`);
|
||||
|
||||
if (this.opts.forceSimulate) {
|
||||
this.#logger.log(
|
||||
"Forced simulation enabled. Will simulate regardless of checkpoint support."
|
||||
);
|
||||
}
|
||||
|
||||
try {
|
||||
await $`criu --version`;
|
||||
} catch (error) {
|
||||
this.#logger.error("No checkpoint support: Missing CRIU binary");
|
||||
if (this.#dockerMode) {
|
||||
this.#logger.error("Will simulate instead");
|
||||
}
|
||||
this.#canCheckpoint = false;
|
||||
this.#initialized = true;
|
||||
return;
|
||||
|
||||
return this.#getInitializeReturn();
|
||||
}
|
||||
|
||||
if (this.#dockerMode) {
|
||||
@@ -58,29 +77,41 @@ class Checkpointer {
|
||||
this.#logger.error(
|
||||
"No checkpoint support: Docker needs to have experimental features enabled"
|
||||
);
|
||||
this.#logger.error("Will simulate instead");
|
||||
this.#canCheckpoint = false;
|
||||
this.#initialized = true;
|
||||
return;
|
||||
|
||||
return this.#getInitializeReturn();
|
||||
}
|
||||
}
|
||||
|
||||
this.#logger.log(
|
||||
`Full checkpoint support with docker ${this.#dockerMode ? "enabled" : "disabled"}`
|
||||
`Full checkpoint support in ${this.#dockerMode ? "docker" : "kubernetes"} mode`
|
||||
);
|
||||
|
||||
this.#initialized = true;
|
||||
this.#canCheckpoint = true;
|
||||
|
||||
return this.#getInitializeReturn();
|
||||
}
|
||||
|
||||
async checkpointAndPush(podName: string) {
|
||||
#getInitializeReturn(): CheckpointerInitializeReturn {
|
||||
return {
|
||||
canCheckpoint: this.#canCheckpoint,
|
||||
willSimulate: this.#dockerMode && (!this.#canCheckpoint || this.opts.forceSimulate),
|
||||
};
|
||||
}
|
||||
|
||||
async checkpointAndPush(podName: string, leaveRunning = false) {
|
||||
await this.initialize();
|
||||
|
||||
if (!this.#canCheckpoint) {
|
||||
if (!this.#dockerMode && !this.#canCheckpoint) {
|
||||
this.#logger.error("No checkpoint support. Simulation requires docker.");
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
const { path } = await this.#checkpointContainer(podName);
|
||||
const { path } = await this.#checkpointContainer(podName, leaveRunning);
|
||||
const { tag } = await this.#buildImage(path, podName);
|
||||
const { destination } = await this.#pushImage(tag);
|
||||
|
||||
@@ -102,20 +133,27 @@ class Checkpointer {
|
||||
}
|
||||
}
|
||||
|
||||
async #checkpointContainer(podName: string) {
|
||||
async #checkpointContainer(podName: string, leaveRunning = false) {
|
||||
await this.initialize();
|
||||
|
||||
if (!this.#canCheckpoint) {
|
||||
throw new Error("No checkpoint support");
|
||||
}
|
||||
|
||||
if (this.#dockerMode) {
|
||||
this.#logger.log("Checkpointing:", podName);
|
||||
|
||||
const path = randomUUID();
|
||||
|
||||
try {
|
||||
this.#logger.debug(await $`docker checkpoint create --leave-running ${podName} ${path}`);
|
||||
if (this.opts.forceSimulate || !this.#canCheckpoint) {
|
||||
this.#logger.log("Simulating checkpoint");
|
||||
this.#logger.debug(await $`docker pause ${podName}`);
|
||||
} else {
|
||||
if (leaveRunning) {
|
||||
this.#logger.debug(
|
||||
await $`docker checkpoint create --leave-running ${podName} ${path}`
|
||||
);
|
||||
} else {
|
||||
this.#logger.debug(await $`docker checkpoint create ${podName} ${path}`);
|
||||
}
|
||||
}
|
||||
} catch (error: any) {
|
||||
this.#logger.error(error.stderr);
|
||||
}
|
||||
@@ -123,6 +161,10 @@ class Checkpointer {
|
||||
return { path };
|
||||
}
|
||||
|
||||
if (!this.#canCheckpoint) {
|
||||
throw new Error("No checkpoint support. Simulation requires docker.");
|
||||
}
|
||||
|
||||
const containerId = this.#logger.debug(
|
||||
// @ts-expect-error
|
||||
await $`crictl ps`
|
||||
@@ -146,15 +188,15 @@ class Checkpointer {
|
||||
async #buildImage(checkpointPath: string, tag: string) {
|
||||
await this.initialize();
|
||||
|
||||
if (!this.#canCheckpoint) {
|
||||
throw new Error("No checkpoint support");
|
||||
}
|
||||
|
||||
if (this.#dockerMode) {
|
||||
// Nothing to do here
|
||||
return { tag };
|
||||
}
|
||||
|
||||
if (!this.#canCheckpoint) {
|
||||
throw new Error("No checkpoint support. Simulation requires docker.");
|
||||
}
|
||||
|
||||
const container = this.#logger.debug(await $`buildah from scratch`);
|
||||
this.#logger.debug(await $`buildah add ${container} ${checkpointPath} /`);
|
||||
this.#logger.debug(
|
||||
@@ -171,15 +213,15 @@ class Checkpointer {
|
||||
async #pushImage(tag: string) {
|
||||
await this.initialize();
|
||||
|
||||
if (!this.#canCheckpoint) {
|
||||
throw new Error("No checkpoint support");
|
||||
}
|
||||
|
||||
if (this.#dockerMode) {
|
||||
// Nothing to do here
|
||||
return { destination: "" };
|
||||
}
|
||||
|
||||
if (!this.#canCheckpoint) {
|
||||
throw new Error("No checkpoint support. Simulation requires docker.");
|
||||
}
|
||||
|
||||
const destination = `${REGISTRY_FQDN}/${REPO_NAME}:${tag}`;
|
||||
this.#logger.debug(await $`buildah push --tls-verify=${REGISTRY_TLS_VERIFY} ${destination}`);
|
||||
|
||||
@@ -191,7 +233,7 @@ class Checkpointer {
|
||||
|
||||
class TaskCoordinator {
|
||||
#httpServer: ReturnType<typeof createServer>;
|
||||
#checkpointer = new Checkpointer();
|
||||
#checkpointer = new Checkpointer({ forceSimulate: true });
|
||||
|
||||
#prodWorkerNamespace: ZodNamespace<
|
||||
typeof ProdWorkerToCoordinatorMessages,
|
||||
@@ -249,6 +291,16 @@ class TaskCoordinator {
|
||||
|
||||
taskSocket.emit("RESUME", message);
|
||||
},
|
||||
RESUME_AFTER_DURATION: async (message) => {
|
||||
const taskSocket = await this.#getAttemptSocket(message.attemptId);
|
||||
|
||||
if (!taskSocket) {
|
||||
logger.log("Socket for attempt not found", { attemptId: message.attemptId });
|
||||
return;
|
||||
}
|
||||
|
||||
taskSocket.emit("RESUME_AFTER_DURATION", message);
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
@@ -341,34 +393,49 @@ class TaskCoordinator {
|
||||
return;
|
||||
}
|
||||
|
||||
// FIXME: shouldn't wait for completion here
|
||||
const completionAck = await socket.emitWithAck("EXECUTE_TASK_RUN", {
|
||||
socket.emit("EXECUTE_TASK_RUN", {
|
||||
version: "v1",
|
||||
executionPayload: executionAck.payload,
|
||||
});
|
||||
});
|
||||
|
||||
if (!completionAck.success) {
|
||||
logger.error("completion unsuccessful", { attemptId: socket.data.attemptId });
|
||||
return;
|
||||
}
|
||||
socket.on("READY_FOR_RESUME", async (message) => {
|
||||
logger.log("[READY_FOR_RESUME]", message);
|
||||
this.#platformSocket?.send("READY_FOR_RESUME", message);
|
||||
});
|
||||
|
||||
logger.log("completed task", { completionId: completionAck.completion.id });
|
||||
socket.on("TASK_RUN_COMPLETED", async (message, callback) => {
|
||||
logger.log("completed task", { completionId: message.completion.id });
|
||||
|
||||
this.#platformSocket?.send("TASK_RUN_COMPLETED", {
|
||||
version: "v1",
|
||||
execution: executionAck.payload.execution,
|
||||
completion: completionAck.completion,
|
||||
execution: message.execution,
|
||||
completion: message.completion,
|
||||
});
|
||||
|
||||
callback();
|
||||
});
|
||||
|
||||
socket.on("WAIT_FOR_DURATION", async (message, callback) => {
|
||||
logger.log("[WAIT_FOR_DURATION]", message);
|
||||
|
||||
const { canCheckpoint, willSimulate } = await this.#checkpointer.initialize();
|
||||
|
||||
callback({ willCheckpointAndRestore: canCheckpoint || willSimulate });
|
||||
|
||||
if (!canCheckpoint) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Wait for attempt to reach checkpointable state
|
||||
// TODO: The worker should let us know when to checkpoint so we don't have to guess
|
||||
await new Promise((resolve) => setTimeout(resolve, 2_000));
|
||||
|
||||
const checkpoint = await this.#checkpointer.checkpointAndPush(socket.data.podName);
|
||||
|
||||
if (!checkpoint) {
|
||||
logger.error("Failed to checkpoint", { podName: socket.data.podName });
|
||||
callback({ success: false });
|
||||
// TODO: We have to let the worker know about failures so it can use its own timer
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -377,10 +444,63 @@ class TaskCoordinator {
|
||||
attemptId: socket.data.attemptId,
|
||||
docker: checkpoint.docker,
|
||||
location: checkpoint.destination,
|
||||
reason: "WAIT_FOR_DURATION",
|
||||
reason: {
|
||||
type: "WAIT_FOR_DURATION",
|
||||
ms: message.ms,
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
callback({ success: true });
|
||||
socket.on("WAIT_FOR_TASK", async (message, callback) => {
|
||||
logger.log("[WAIT_FOR_TASK]", message);
|
||||
|
||||
const { canCheckpoint, willSimulate } = await this.#checkpointer.initialize();
|
||||
|
||||
callback({ willCheckpointAndRestore: canCheckpoint || willSimulate });
|
||||
|
||||
const checkpoint = await this.#checkpointer.checkpointAndPush(socket.data.podName);
|
||||
|
||||
if (!checkpoint) {
|
||||
logger.error("Failed to checkpoint", { podName: socket.data.podName });
|
||||
return;
|
||||
}
|
||||
|
||||
this.#platformSocket?.send("CHECKPOINT_CREATED", {
|
||||
version: "v1",
|
||||
attemptId: socket.data.attemptId,
|
||||
docker: checkpoint.docker,
|
||||
location: checkpoint.destination,
|
||||
reason: {
|
||||
type: "WAIT_FOR_TASK",
|
||||
id: message.id,
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
socket.on("WAIT_FOR_BATCH", async (message, callback) => {
|
||||
logger.log("[WAIT_FOR_BATCH]", message);
|
||||
|
||||
const { canCheckpoint, willSimulate } = await this.#checkpointer.initialize();
|
||||
|
||||
callback({ willCheckpointAndRestore: canCheckpoint || willSimulate });
|
||||
|
||||
const checkpoint = await this.#checkpointer.checkpointAndPush(socket.data.podName);
|
||||
|
||||
if (!checkpoint) {
|
||||
logger.error("Failed to checkpoint", { podName: socket.data.podName });
|
||||
return;
|
||||
}
|
||||
|
||||
this.#platformSocket?.send("CHECKPOINT_CREATED", {
|
||||
version: "v1",
|
||||
attemptId: socket.data.attemptId,
|
||||
docker: checkpoint.docker,
|
||||
location: checkpoint.destination,
|
||||
reason: {
|
||||
type: "WAIT_FOR_BATCH",
|
||||
id: message.id,
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
socket.on("INDEX_TASKS", async (message, callback) => {
|
||||
@@ -428,12 +548,6 @@ class TaskCoordinator {
|
||||
TASK_HEARTBEAT: async (message) => {
|
||||
this.#platformSocket?.send("TASK_HEARTBEAT", message);
|
||||
},
|
||||
WAIT_FOR_BATCH: async (message) => {
|
||||
// this.#checkpointer.checkpointAndPush(socket.data.podName);
|
||||
},
|
||||
WAIT_FOR_TASK: async (message) => {
|
||||
// this.#checkpointer.checkpointAndPush(socket.data.podName);
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
|
||||
@@ -11,6 +11,8 @@ const OTEL_EXPORTER_OTLP_ENDPOINT =
|
||||
const logger = new SimpleLogger(`[${MACHINE_NAME}]`);
|
||||
|
||||
class DockerTaskOperations implements TaskOperations {
|
||||
constructor(private opts = { forceSimulate: false }) {}
|
||||
|
||||
async index(opts: {
|
||||
contentHash: string;
|
||||
imageTag: string;
|
||||
@@ -26,7 +28,7 @@ class DockerTaskOperations implements TaskOperations {
|
||||
});
|
||||
|
||||
const { exitCode } = logger.debug(
|
||||
await $`docker run --rm -e TRIGGER_SECRET_KEY=${opts.apiKey} -e TRIGGER_API_URL=${opts.apiUrl} -e COORDINATOR_HOST=${COORDINATOR_HOST} -e COORDINATOR_PORT=${COORDINATOR_PORT} -e POD_NAME=${containerName} -e TRIGGER_ENV_ID=${opts.envId} -e INDEX_TASKS=true --name=${containerName} ${opts.imageTag}`
|
||||
await $`docker run --network=host --rm -e TRIGGER_SECRET_KEY=${opts.apiKey} -e TRIGGER_API_URL=${opts.apiUrl} -e COORDINATOR_HOST=${COORDINATOR_HOST} -e COORDINATOR_PORT=${COORDINATOR_PORT} -e POD_NAME=${containerName} -e TRIGGER_ENV_ID=${opts.envId} -e INDEX_TASKS=true --name=${containerName} ${opts.imageTag}`
|
||||
);
|
||||
|
||||
if (exitCode !== 0) {
|
||||
@@ -38,7 +40,7 @@ class DockerTaskOperations implements TaskOperations {
|
||||
const containerName = this.#getRunContainerName(opts.attemptId);
|
||||
|
||||
const { exitCode } = logger.debug(
|
||||
await $`docker run -d -e OTEL_EXPORTER_OTLP_ENDPOINT=${OTEL_EXPORTER_OTLP_ENDPOINT} -e COORDINATOR_HOST=${COORDINATOR_HOST} -e COORDINATOR_PORT=${COORDINATOR_PORT} -e POD_NAME=${containerName} -e TRIGGER_ENV_ID=${opts.envId} -e TRIGGER_ATTEMPT_ID=${opts.attemptId} --name=${containerName} ${opts.image}`
|
||||
await $`docker run --network=host -d -e OTEL_EXPORTER_OTLP_ENDPOINT=${OTEL_EXPORTER_OTLP_ENDPOINT} -e COORDINATOR_HOST=${COORDINATOR_HOST} -e COORDINATOR_PORT=${COORDINATOR_PORT} -e POD_NAME=${containerName} -e TRIGGER_ENV_ID=${opts.envId} -e TRIGGER_ATTEMPT_ID=${opts.attemptId} --name=${containerName} ${opts.image}`
|
||||
);
|
||||
|
||||
if (exitCode !== 0) {
|
||||
@@ -46,18 +48,23 @@ class DockerTaskOperations implements TaskOperations {
|
||||
}
|
||||
}
|
||||
|
||||
async restore(opts: {
|
||||
attemptId: string;
|
||||
runId: string;
|
||||
image: string;
|
||||
name: string;
|
||||
checkpointId: string;
|
||||
machine: Machine;
|
||||
}) {
|
||||
async restore(opts: { attemptId: string; checkpointRef: string; machine: Machine }) {
|
||||
const containerName = this.#getRunContainerName(opts.attemptId);
|
||||
|
||||
if (this.opts.forceSimulate) {
|
||||
logger.log("Simulating restore");
|
||||
|
||||
const { exitCode } = logger.debug(await $`docker unpause ${containerName}`);
|
||||
|
||||
if (exitCode !== 0) {
|
||||
throw new Error("docker unpause command failed");
|
||||
}
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
const { exitCode } = logger.debug(
|
||||
await $`docker start --checkpoint=${opts.checkpointId} ${containerName}`
|
||||
await $`docker start --checkpoint=${opts.checkpointRef} ${containerName}`
|
||||
);
|
||||
|
||||
if (exitCode !== 0) {
|
||||
@@ -83,7 +90,8 @@ class DockerTaskOperations implements TaskOperations {
|
||||
}
|
||||
|
||||
const provider = new ProviderShell({
|
||||
tasks: new DockerTaskOperations(),
|
||||
tasks: new DockerTaskOperations({ forceSimulate: true }),
|
||||
type: "docker",
|
||||
});
|
||||
|
||||
provider.listen();
|
||||
|
||||
@@ -389,6 +389,7 @@ class KubernetesTaskOperations implements TaskOperations {
|
||||
|
||||
const provider = new ProviderShell({
|
||||
tasks: new KubernetesTaskOperations(),
|
||||
type: "kubernetes",
|
||||
});
|
||||
|
||||
provider.listen();
|
||||
|
||||
@@ -14,10 +14,10 @@ import { SharedSocketConnection } from "./sharedSocketConnection";
|
||||
import { CreateCheckpointService } from "./services/createCheckpoint.server";
|
||||
import { sharedQueueTasks } from "./marqs/sharedQueueConsumer.server";
|
||||
import { CompleteAttemptService } from "./services/completeAttempt.server";
|
||||
import { CreateBackgroundWorkerService } from "./services/createBackgroundWorker.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import { findEnvironmentById } from "~/models/runtimeEnvironment.server";
|
||||
import { CreateDeployedBackgroundWorkerService } from "./services/createDeployedBackgroundWorker.server";
|
||||
import { ResumeAttemptService } from "./services/resumeAttempt.server";
|
||||
import { DeploymentIndexFailed } from "./services/deploymentIndexFailed.server";
|
||||
|
||||
export const socketIo = singleton("socketIo", initalizeIoServer);
|
||||
@@ -58,6 +58,10 @@ function createCoordinatorNamespace(io: Server) {
|
||||
return { success: true, payload };
|
||||
}
|
||||
},
|
||||
READY_FOR_RESUME: async (message) => {
|
||||
const resumeAttempt = new ResumeAttemptService();
|
||||
await resumeAttempt.call(message);
|
||||
},
|
||||
TASK_RUN_COMPLETED: async (message) => {
|
||||
const completeAttempt = new CompleteAttemptService();
|
||||
await completeAttempt.call(message.completion, message.execution);
|
||||
|
||||
@@ -32,6 +32,9 @@ const MessageBody = z.discriminatedUnion("type", [
|
||||
type: z.literal("RESUME"),
|
||||
completedAttemptIds: z.string().array(),
|
||||
}),
|
||||
z.object({
|
||||
type: z.literal("RESUME_AFTER_DURATION"),
|
||||
}),
|
||||
]);
|
||||
|
||||
type BackgroundWorkerWithTasks = BackgroundWorker & { tasks: BackgroundWorkerTask[] };
|
||||
@@ -500,7 +503,7 @@ export class SharedQueueConsumer {
|
||||
const resumableAttempt = resumableRun?.attempts[0];
|
||||
|
||||
if (!resumableAttempt) {
|
||||
logger.error("Task run attempt to resume not found", {
|
||||
logger.error("Resumable attempt not found", {
|
||||
queueMessage: message.data,
|
||||
messageId: message.messageId,
|
||||
});
|
||||
@@ -509,67 +512,6 @@ export class SharedQueueConsumer {
|
||||
return;
|
||||
}
|
||||
|
||||
const deployment = await prisma.workerDeployment.findFirst({
|
||||
where: {
|
||||
environmentId: resumableRun.runtimeEnvironmentId,
|
||||
projectId: resumableRun.projectId,
|
||||
status: "DEPLOYED",
|
||||
imageReference: {
|
||||
not: null,
|
||||
},
|
||||
},
|
||||
orderBy: {
|
||||
createdAt: "desc",
|
||||
},
|
||||
include: {
|
||||
worker: {
|
||||
include: {
|
||||
tasks: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
if (!deployment || !deployment.worker) {
|
||||
logger.error("No matching deployment found for task run", {
|
||||
queueMessage: message.data,
|
||||
messageId: message.messageId,
|
||||
});
|
||||
await marqs?.acknowledgeMessage(message.messageId);
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
return;
|
||||
}
|
||||
|
||||
if (!deployment.imageReference) {
|
||||
logger.error("Deployment is missing an image reference", {
|
||||
queueMessage: message.data,
|
||||
messageId: message.messageId,
|
||||
deployment: deployment.id,
|
||||
});
|
||||
await marqs?.acknowledgeMessage(message.messageId);
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
return;
|
||||
}
|
||||
|
||||
const backgroundTask = deployment.worker.tasks.find(
|
||||
(task) => task.slug === resumableRun.taskIdentifier
|
||||
);
|
||||
|
||||
if (!backgroundTask) {
|
||||
logger.warn("No matching background task found for task run", {
|
||||
taskRun: resumableRun.id,
|
||||
taskIdentifier: resumableRun.taskIdentifier,
|
||||
deployment: deployment.id,
|
||||
backgroundWorker: deployment.worker.id,
|
||||
taskSlugs: deployment.worker.tasks.map((task) => task.slug),
|
||||
});
|
||||
|
||||
await marqs?.acknowledgeMessage(message.messageId);
|
||||
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
return;
|
||||
}
|
||||
|
||||
const queue = await prisma.taskQueue.findUnique({
|
||||
where: {
|
||||
runtimeEnvironmentId_name: {
|
||||
@@ -590,6 +532,43 @@ export class SharedQueueConsumer {
|
||||
return;
|
||||
}
|
||||
|
||||
if (resumableAttempt.status === "PAUSED") {
|
||||
// We need to restore the attempt from the latest checkpoint before we can resume
|
||||
const latestCheckpoint = resumableAttempt.checkpoints[0];
|
||||
|
||||
if (!latestCheckpoint) {
|
||||
logger.error("No checkpoint found", {
|
||||
queueMessage: message.data,
|
||||
messageId: message.messageId,
|
||||
resumableAttemptId: resumableAttempt.id,
|
||||
});
|
||||
await marqs?.acknowledgeMessage(message.messageId);
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
return;
|
||||
}
|
||||
|
||||
await prisma.taskRunAttempt.update({
|
||||
where: {
|
||||
id: resumableAttempt.id,
|
||||
},
|
||||
data: {
|
||||
status: "EXECUTING",
|
||||
},
|
||||
});
|
||||
|
||||
socketIo.providerNamespace.emit("RESTORE", {
|
||||
version: "v1",
|
||||
id: latestCheckpoint.id,
|
||||
attemptId: latestCheckpoint.attemptId,
|
||||
type: latestCheckpoint.type,
|
||||
location: latestCheckpoint.location,
|
||||
reason: latestCheckpoint.reason ?? undefined,
|
||||
});
|
||||
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
return;
|
||||
}
|
||||
|
||||
const completions: TaskRunExecutionResult[] = [];
|
||||
const executions: TaskRunExecution[] = [];
|
||||
|
||||
@@ -643,30 +622,108 @@ export class SharedQueueConsumer {
|
||||
}
|
||||
|
||||
try {
|
||||
// The attempt should still be running so we can broadcast to all coordinators to resume immediately
|
||||
socketIo.coordinatorNamespace.emit("RESUME", {
|
||||
version: "v1",
|
||||
attemptId: resumableAttempt.id,
|
||||
completions,
|
||||
executions,
|
||||
});
|
||||
} catch (e) {
|
||||
if (e instanceof Error) {
|
||||
this._currentSpan?.recordException(e);
|
||||
} else {
|
||||
this._currentSpan?.recordException(new Error(String(e)));
|
||||
}
|
||||
|
||||
this._endSpanInNextIteration = true;
|
||||
|
||||
// Finally we need to nack the message so it can be retried
|
||||
await marqs?.nackMessage(message.messageId);
|
||||
} finally {
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
}
|
||||
break;
|
||||
}
|
||||
// Resume after duration-based wait
|
||||
case "RESUME_AFTER_DURATION": {
|
||||
const resumableRun = await prisma.taskRun.findFirst({
|
||||
where: {
|
||||
id: message.messageId,
|
||||
},
|
||||
include: {
|
||||
attempts: {
|
||||
orderBy: {
|
||||
createdAt: "desc",
|
||||
},
|
||||
take: 1,
|
||||
include: {
|
||||
checkpoints: {
|
||||
take: 1,
|
||||
orderBy: {
|
||||
createdAt: "desc",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const resumableAttempt = resumableRun?.attempts[0];
|
||||
|
||||
if (!resumableAttempt) {
|
||||
logger.error("Resumable attempt not found", {
|
||||
queueMessage: message.data,
|
||||
messageId: message.messageId,
|
||||
});
|
||||
await marqs?.acknowledgeMessage(message.messageId);
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
return;
|
||||
}
|
||||
|
||||
if (resumableAttempt.status !== "PAUSED") {
|
||||
logger.error("Attempt not paused", {
|
||||
queueMessage: message.data,
|
||||
messageId: message.messageId,
|
||||
});
|
||||
await marqs?.acknowledgeMessage(message.messageId);
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
// We need to restore the attempt from the latest checkpoint before we can resume
|
||||
const latestCheckpoint = resumableAttempt.checkpoints[0];
|
||||
|
||||
if (!latestCheckpoint) {
|
||||
// No checkpoint means the task should still be running
|
||||
// We can broadcast to all coordinators to resume immediately
|
||||
socketIo.coordinatorNamespace.emit("RESUME", {
|
||||
version: "v1",
|
||||
attemptId: resumableAttempt.id,
|
||||
image: deployment.imageReference,
|
||||
completions,
|
||||
executions,
|
||||
});
|
||||
} else {
|
||||
// There's a checkpoint we need to restore first
|
||||
// TODO: Send RESUME message once the restored task has checked in
|
||||
socketIo.providerNamespace.emit("RESTORE", {
|
||||
version: "v1",
|
||||
id: latestCheckpoint.id,
|
||||
attemptId: latestCheckpoint.attemptId,
|
||||
type: latestCheckpoint.type,
|
||||
location: latestCheckpoint.location,
|
||||
reason: latestCheckpoint.reason ?? undefined,
|
||||
logger.error("No checkpoint found", {
|
||||
queueMessage: message.data,
|
||||
messageId: message.messageId,
|
||||
resumableAttemptId: resumableAttempt.id,
|
||||
});
|
||||
await marqs?.acknowledgeMessage(message.messageId);
|
||||
setTimeout(() => this.#doWork(), 100);
|
||||
return;
|
||||
}
|
||||
|
||||
await prisma.taskRunAttempt.update({
|
||||
where: {
|
||||
id: resumableAttempt.id,
|
||||
},
|
||||
data: {
|
||||
status: "EXECUTING",
|
||||
},
|
||||
});
|
||||
|
||||
// The attempt will resume automatically after restore
|
||||
socketIo.providerNamespace.emit("RESTORE", {
|
||||
version: "v1",
|
||||
id: latestCheckpoint.id,
|
||||
attemptId: latestCheckpoint.attemptId,
|
||||
type: latestCheckpoint.type,
|
||||
location: latestCheckpoint.location,
|
||||
reason: latestCheckpoint.reason ?? undefined,
|
||||
});
|
||||
} catch (e) {
|
||||
if (e instanceof Error) {
|
||||
this._currentSpan?.recordException(e);
|
||||
|
||||
@@ -182,7 +182,11 @@ export class CompleteAttemptService extends BaseService {
|
||||
},
|
||||
},
|
||||
include: {
|
||||
dependentTaskAttempt: true,
|
||||
dependentTaskAttempt: {
|
||||
include: {
|
||||
taskRun: true,
|
||||
},
|
||||
},
|
||||
items: {
|
||||
include: {
|
||||
taskRun: {
|
||||
@@ -220,12 +224,27 @@ export class CompleteAttemptService extends BaseService {
|
||||
return "ACKNOWLEDGED";
|
||||
}
|
||||
|
||||
await marqs?.replaceMessage(finalizedBatchRun.dependentTaskAttempt.taskRunId, {
|
||||
type: "RESUME",
|
||||
completedAttemptIds: finalizedBatchRun.items.map(
|
||||
(item) => item.taskRun.attempts[0]?.id
|
||||
),
|
||||
});
|
||||
const dependentRun = finalizedBatchRun.dependentTaskAttempt.taskRun;
|
||||
|
||||
if (finalizedBatchRun.dependentTaskAttempt.status === "PAUSED") {
|
||||
await marqs?.enqueueMessage(
|
||||
environment,
|
||||
dependentRun.queue,
|
||||
dependentRun.id,
|
||||
{
|
||||
type: "RESUME",
|
||||
completedAttemptIds: [taskRunAttempt.id],
|
||||
},
|
||||
dependentRun.concurrencyKey ?? undefined
|
||||
);
|
||||
} else {
|
||||
await marqs?.replaceMessage(dependentRun.id, {
|
||||
type: "RESUME",
|
||||
completedAttemptIds: finalizedBatchRun.items.map(
|
||||
(item) => item.taskRun.attempts[0]?.id
|
||||
),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -263,10 +282,23 @@ export class CompleteAttemptService extends BaseService {
|
||||
return "FAILED";
|
||||
}
|
||||
|
||||
await marqs?.replaceMessage(dependentRun.id, {
|
||||
type: "RESUME",
|
||||
completedAttemptIds: [taskRunAttempt.id],
|
||||
});
|
||||
if (dependency.dependentAttempt.status === "PAUSED") {
|
||||
await marqs?.enqueueMessage(
|
||||
environment,
|
||||
dependentRun.queue,
|
||||
dependentRun.id,
|
||||
{
|
||||
type: "RESUME",
|
||||
completedAttemptIds: [taskRunAttempt.id],
|
||||
},
|
||||
dependentRun.concurrencyKey ?? undefined
|
||||
);
|
||||
} else {
|
||||
await marqs?.replaceMessage(dependentRun.id, {
|
||||
type: "RESUME",
|
||||
completedAttemptIds: [taskRunAttempt.id],
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -34,12 +34,37 @@ export class CreateCheckpointService {
|
||||
attemptId: attempt.id,
|
||||
location: params.location,
|
||||
type: params.docker ? "DOCKER" : "KUBERNETES",
|
||||
reason: params.reason,
|
||||
reason: params.reason.type,
|
||||
},
|
||||
});
|
||||
|
||||
// TODO: Can't heartbeat when checkpointed, so we ACK to prevent automatic requeue
|
||||
// await marqs?.acknowledgeMessage(attempt.taskRunId);
|
||||
await this.#prismaClient.taskRunAttempt.update({
|
||||
where: {
|
||||
id: params.attemptId,
|
||||
},
|
||||
data: {
|
||||
status: "PAUSED",
|
||||
},
|
||||
});
|
||||
|
||||
switch (params.reason.type) {
|
||||
case "WAIT_FOR_DURATION": {
|
||||
await marqs?.replaceMessage(
|
||||
attempt.taskRunId,
|
||||
{ type: "RESUME_AFTER_DURATION" },
|
||||
Date.now() + params.reason.ms
|
||||
);
|
||||
break;
|
||||
}
|
||||
case "WAIT_FOR_TASK":
|
||||
case "WAIT_FOR_BATCH": {
|
||||
await marqs?.acknowledgeMessage(attempt.taskRunId);
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
return checkpoint;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,178 @@
|
||||
import {
|
||||
CoordinatorToPlatformMessages,
|
||||
InferSocketMessageSchema,
|
||||
TaskRunExecution,
|
||||
TaskRunExecutionResult,
|
||||
} from "@trigger.dev/core/v3";
|
||||
import { PrismaClient, prisma } from "~/db.server";
|
||||
import { logger } from "~/services/logger.server";
|
||||
import { marqs } from "../marqs.server";
|
||||
import { socketIo } from "../handleSocketIo.server";
|
||||
import { sharedQueueTasks } from "../marqs/sharedQueueConsumer.server";
|
||||
|
||||
export class ResumeAttemptService {
|
||||
#prismaClient: PrismaClient;
|
||||
|
||||
constructor(prismaClient: PrismaClient = prisma) {
|
||||
this.#prismaClient = prismaClient;
|
||||
}
|
||||
|
||||
public async call(
|
||||
params: InferSocketMessageSchema<typeof CoordinatorToPlatformMessages, "READY_FOR_RESUME">
|
||||
): Promise<void> {
|
||||
logger.debug(`ResumeAttemptService.call()`, params);
|
||||
|
||||
const attempt = await this.#prismaClient.taskRunAttempt.findUnique({
|
||||
where: {
|
||||
id: params.attemptId,
|
||||
},
|
||||
include: {
|
||||
taskRun: true,
|
||||
taskRunDependency: {
|
||||
include: {
|
||||
taskRun: {
|
||||
include: {
|
||||
attempts: {
|
||||
orderBy: {
|
||||
number: "desc",
|
||||
},
|
||||
take: 1,
|
||||
select: {
|
||||
id: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
batchTaskRunDependency: {
|
||||
include: {
|
||||
items: {
|
||||
include: {
|
||||
taskRun: {
|
||||
include: {
|
||||
attempts: {
|
||||
orderBy: {
|
||||
number: "desc",
|
||||
},
|
||||
take: 1,
|
||||
select: {
|
||||
id: true,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
if (!attempt) {
|
||||
logger.error("Could not find attempt", { attemptId: params.attemptId });
|
||||
return;
|
||||
}
|
||||
|
||||
switch (params.type) {
|
||||
case "WAIT_FOR_DURATION": {
|
||||
// Nothing to do, but thanks for checking in!
|
||||
break;
|
||||
}
|
||||
case "WAIT_FOR_TASK":
|
||||
case "WAIT_FOR_BATCH": {
|
||||
let completedAttemptIds: string[] = [];
|
||||
|
||||
if (attempt.taskRunDependency) {
|
||||
const dependentAttempt = attempt.taskRunDependency.taskRun.attempts[0];
|
||||
|
||||
if (!dependentAttempt) {
|
||||
logger.error("No dependent attempt", { attemptId: params.attemptId });
|
||||
return;
|
||||
}
|
||||
|
||||
completedAttemptIds = [dependentAttempt.id];
|
||||
} else if (attempt.batchTaskRunDependency) {
|
||||
const dependentBatchItems = attempt.batchTaskRunDependency.items;
|
||||
|
||||
if (!dependentBatchItems) {
|
||||
logger.error("No dependent batch items", { attemptId: params.attemptId });
|
||||
return;
|
||||
}
|
||||
|
||||
completedAttemptIds = dependentBatchItems.map((item) => item.taskRun.attempts[0]?.id);
|
||||
} else {
|
||||
logger.error("No dependencies", { attemptId: params.attemptId });
|
||||
return;
|
||||
}
|
||||
|
||||
if (completedAttemptIds.length === 0) {
|
||||
logger.error("No completed attempt IDs", { attemptId: params.attemptId });
|
||||
return;
|
||||
}
|
||||
|
||||
const completions: TaskRunExecutionResult[] = [];
|
||||
const executions: TaskRunExecution[] = [];
|
||||
|
||||
for (const completedAttemptId of completedAttemptIds) {
|
||||
const completedAttempt = await prisma.taskRunAttempt.findUnique({
|
||||
where: {
|
||||
id: completedAttemptId,
|
||||
taskRun: {
|
||||
lockedAt: {
|
||||
not: null,
|
||||
},
|
||||
lockedById: {
|
||||
not: null,
|
||||
},
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
if (!completedAttempt) {
|
||||
logger.error("Completed attempt not found", {
|
||||
attemptId: params.attemptId,
|
||||
completedAttemptId,
|
||||
});
|
||||
await marqs?.acknowledgeMessage(attempt.taskRunId);
|
||||
return;
|
||||
}
|
||||
|
||||
const completion = await sharedQueueTasks.getCompletionPayloadFromAttempt(
|
||||
completedAttempt.id
|
||||
);
|
||||
|
||||
if (!completion) {
|
||||
await marqs?.acknowledgeMessage(attempt.taskRunId);
|
||||
return;
|
||||
}
|
||||
|
||||
completions.push(completion);
|
||||
|
||||
const executionPayload = await sharedQueueTasks.getExecutionPayloadFromAttempt(
|
||||
completedAttempt.id,
|
||||
false
|
||||
);
|
||||
|
||||
if (!executionPayload) {
|
||||
await marqs?.acknowledgeMessage(attempt.taskRunId);
|
||||
return;
|
||||
}
|
||||
|
||||
executions.push(executionPayload.execution);
|
||||
}
|
||||
|
||||
socketIo.coordinatorNamespace.emit("RESUME", {
|
||||
version: "v1",
|
||||
attemptId: params.attemptId,
|
||||
completions,
|
||||
executions,
|
||||
});
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -587,12 +587,12 @@ async function compileProject(config: ResolvedConfig, options: DeployCommandOpti
|
||||
);
|
||||
|
||||
const workerSetupPath = new URL(
|
||||
importResolve("./workers/common/worker-setup.js", import.meta.url)
|
||||
importResolve("./workers/prod/worker-setup.js", import.meta.url)
|
||||
).href.replace("file://", "");
|
||||
|
||||
const workerContents = workerFacade
|
||||
.replace("__TASKS__", createTaskFileImports(taskFiles))
|
||||
.replace("__WORKER_SETUP__", `import { tracingSDK, sender } from "${workerSetupPath}";`);
|
||||
.replace("__WORKER_SETUP__", `import { tracingSDK } from "${workerSetupPath}";`);
|
||||
|
||||
const result = await build({
|
||||
stdin: {
|
||||
|
||||
@@ -305,7 +305,7 @@ function useDev({
|
||||
);
|
||||
|
||||
const workerSetupPath = new URL(
|
||||
importResolve("./workers/common/worker-setup.js", import.meta.url)
|
||||
importResolve("./workers/dev/worker-setup.js", import.meta.url)
|
||||
).href.replace("file://", "");
|
||||
|
||||
const entryPointContents = workerFacade
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
import {
|
||||
BackgroundWorkerProperties,
|
||||
CreateBackgroundWorkerResponse,
|
||||
ProdChildToWorkerMessages,
|
||||
ProdTaskRunExecutionPayload,
|
||||
ProdWorkerToChildMessages,
|
||||
SemanticInternalAttributes,
|
||||
TaskMetadataWithFilePath,
|
||||
TaskRunBuiltInError,
|
||||
@@ -9,11 +11,8 @@ import {
|
||||
TaskRunExecution,
|
||||
TaskRunExecutionPayload,
|
||||
TaskRunExecutionResult,
|
||||
ZodMessageHandler,
|
||||
ZodMessageSender,
|
||||
childToWorkerMessages,
|
||||
ZodIpcConnection,
|
||||
correctErrorStackTrace,
|
||||
workerToChildMessages,
|
||||
} from "@trigger.dev/core/v3";
|
||||
import { Evt } from "evt";
|
||||
import { ChildProcess, fork } from "node:child_process";
|
||||
@@ -45,15 +44,14 @@ type BackgroundWorkerParams = {
|
||||
|
||||
export class ProdBackgroundWorker {
|
||||
private _initialized: boolean = false;
|
||||
private _handler = new ZodMessageHandler({
|
||||
schema: childToWorkerMessages,
|
||||
});
|
||||
|
||||
public onTaskHeartbeat: Evt<string> = new Evt();
|
||||
|
||||
public onWaitForBatch: Evt<{ version?: "v1"; id: string; runs: string[] }> = new Evt();
|
||||
public onWaitForDuration: Evt<{ version?: "v1"; ms: number }> = new Evt();
|
||||
public onWaitForTask: Evt<{ version?: "v1"; id: string }> = new Evt();
|
||||
public onWaitForBatch: Evt<{ version?: "v1"; id: string; runs: string[] }> = new Evt();
|
||||
|
||||
public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>();
|
||||
|
||||
private _onClose: Evt<void> = new Evt();
|
||||
|
||||
@@ -116,20 +114,28 @@ export class ProdBackgroundWorker {
|
||||
reject(new Error("Worker timed out"));
|
||||
}, 10_000);
|
||||
|
||||
child.on("message", async (msg: any) => {
|
||||
const message = this._handler.parseMessage(msg);
|
||||
|
||||
if (message.type === "TASKS_READY" && !resolved) {
|
||||
clearTimeout(timeout);
|
||||
resolved = true;
|
||||
resolve(message.payload.tasks);
|
||||
child.kill();
|
||||
} else if (message.type === "UNCAUGHT_EXCEPTION") {
|
||||
clearTimeout(timeout);
|
||||
resolved = true;
|
||||
reject(new UncaughtExceptionError(message.payload.error, message.payload.origin));
|
||||
child.kill();
|
||||
}
|
||||
new ZodIpcConnection({
|
||||
listenSchema: ProdChildToWorkerMessages,
|
||||
emitSchema: ProdWorkerToChildMessages,
|
||||
process: child,
|
||||
handlers: {
|
||||
TASKS_READY: async (message) => {
|
||||
if (!resolved) {
|
||||
clearTimeout(timeout);
|
||||
resolved = true;
|
||||
resolve(message.tasks);
|
||||
child.kill();
|
||||
}
|
||||
},
|
||||
UNCAUGHT_EXCEPTION: async (message) => {
|
||||
if (!resolved) {
|
||||
clearTimeout(timeout);
|
||||
resolved = true;
|
||||
reject(new UncaughtExceptionError(message.error, message.origin));
|
||||
child.kill();
|
||||
}
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
child.stdout?.on("data", (data) => {
|
||||
@@ -170,6 +176,11 @@ export class ProdBackgroundWorker {
|
||||
taskRunProcess.taskRunCompletedNotification(completion, execution);
|
||||
}
|
||||
}
|
||||
async waitCompletedNotification() {
|
||||
for (const taskRunProcess of this._taskRunProcesses.values()) {
|
||||
taskRunProcess.waitCompletedNotification();
|
||||
}
|
||||
}
|
||||
|
||||
async #initializeTaskRunProcess(payload: ProdTaskRunExecutionPayload): Promise<TaskRunProcess> {
|
||||
const metadata = this.getMetadata(
|
||||
@@ -208,6 +219,10 @@ export class ProdBackgroundWorker {
|
||||
this.onWaitForTask.post(message);
|
||||
});
|
||||
|
||||
this.preCheckpointNotification.attach((message) => {
|
||||
taskRunProcess.preCheckpointNotification.post(message);
|
||||
});
|
||||
|
||||
await taskRunProcess.initialize();
|
||||
|
||||
this._taskRunProcesses.set(payload.execution.run.id, taskRunProcess);
|
||||
@@ -291,11 +306,12 @@ export class ProdBackgroundWorker {
|
||||
}
|
||||
|
||||
class TaskRunProcess {
|
||||
private _handler = new ZodMessageHandler({
|
||||
schema: childToWorkerMessages,
|
||||
});
|
||||
private _sender: ZodMessageSender<typeof workerToChildMessages>;
|
||||
private _child: ChildProcess | undefined;
|
||||
private _ipc?: ZodIpcConnection<
|
||||
typeof ProdChildToWorkerMessages,
|
||||
typeof ProdWorkerToChildMessages
|
||||
>;
|
||||
private _child?: ChildProcess;
|
||||
|
||||
private _attemptPromises: Map<
|
||||
string,
|
||||
{ resolver: (value: TaskRunExecutionResult) => void; rejecter: (err?: any) => void }
|
||||
@@ -311,21 +327,14 @@ class TaskRunProcess {
|
||||
public onWaitForDuration: Evt<{ version?: "v1"; ms: number }> = new Evt();
|
||||
public onWaitForTask: Evt<{ version?: "v1"; id: string }> = new Evt();
|
||||
|
||||
public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>();
|
||||
|
||||
constructor(
|
||||
private path: string,
|
||||
private env: NodeJS.ProcessEnv,
|
||||
private metadata: BackgroundWorkerProperties,
|
||||
private worker: BackgroundWorkerParams
|
||||
) {
|
||||
this._sender = new ZodMessageSender({
|
||||
schema: workerToChildMessages,
|
||||
sender: async (message) => {
|
||||
if (this._child?.connected && !this._isBeingKilled && !this._child.killed) {
|
||||
this._child?.send?.(message);
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
) {}
|
||||
|
||||
async initialize() {
|
||||
this._child = fork(this.path, {
|
||||
@@ -339,7 +348,57 @@ class TaskRunProcess {
|
||||
},
|
||||
});
|
||||
|
||||
this._child.on("message", this.#handleMessage.bind(this));
|
||||
this._ipc = new ZodIpcConnection({
|
||||
listenSchema: ProdChildToWorkerMessages,
|
||||
emitSchema: ProdWorkerToChildMessages,
|
||||
process: this._child,
|
||||
handlers: {
|
||||
TASK_RUN_COMPLETED: async (message) => {
|
||||
const { result, execution } = message;
|
||||
|
||||
const promiseStatus = this._attemptStatuses.get(execution.attempt.id);
|
||||
|
||||
if (promiseStatus !== "PENDING") {
|
||||
return;
|
||||
}
|
||||
|
||||
this._attemptStatuses.set(execution.attempt.id, "RESOLVED");
|
||||
|
||||
const attemptPromise = this._attemptPromises.get(execution.attempt.id);
|
||||
|
||||
if (!attemptPromise) {
|
||||
return;
|
||||
}
|
||||
|
||||
const { resolver } = attemptPromise;
|
||||
|
||||
resolver(result);
|
||||
},
|
||||
READY_TO_DISPOSE: async (message) => {
|
||||
this.#kill();
|
||||
},
|
||||
TASK_HEARTBEAT: async (message) => {
|
||||
this.onTaskHeartbeat.post(message.id);
|
||||
},
|
||||
TASKS_READY: async (message) => {},
|
||||
WAIT_FOR_BATCH: async (message) => {
|
||||
this.onWaitForBatch.post(message);
|
||||
},
|
||||
WAIT_FOR_DURATION: async (message) => {
|
||||
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();
|
||||
|
||||
return { willCheckpointAndRestore };
|
||||
},
|
||||
WAIT_FOR_TASK: async (message) => {
|
||||
this.onWaitForTask.post(message);
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
this._child.on("exit", this.#handleExit.bind(this));
|
||||
this._child.stdout?.on("data", this.#handleLog.bind(this));
|
||||
this._child.stderr?.on("data", this.#handleStdErr.bind(this));
|
||||
@@ -350,7 +409,7 @@ class TaskRunProcess {
|
||||
return;
|
||||
}
|
||||
|
||||
await this._sender.send("CLEANUP", {
|
||||
await this._ipc?.send("CLEANUP", {
|
||||
flush: true,
|
||||
kill,
|
||||
});
|
||||
@@ -376,11 +435,13 @@ class TaskRunProcess {
|
||||
|
||||
this._currentExecution = execution;
|
||||
|
||||
await this._sender.send("EXECUTE_TASK_RUN", {
|
||||
execution,
|
||||
traceContext,
|
||||
metadata: this.metadata,
|
||||
});
|
||||
if (this._child?.connected && !this._isBeingKilled && !this._child.killed) {
|
||||
await this._ipc?.send("EXECUTE_TASK_RUN", {
|
||||
execution,
|
||||
traceContext,
|
||||
metadata: this.metadata,
|
||||
});
|
||||
}
|
||||
|
||||
const result = await promise;
|
||||
|
||||
@@ -394,67 +455,17 @@ class TaskRunProcess {
|
||||
return;
|
||||
}
|
||||
|
||||
this._sender.send("TASK_RUN_COMPLETED_NOTIFICATION", {
|
||||
completion,
|
||||
execution,
|
||||
});
|
||||
if (this._child?.connected && !this._isBeingKilled && !this._child.killed) {
|
||||
this._ipc?.send("TASK_RUN_COMPLETED_NOTIFICATION", {
|
||||
completion,
|
||||
execution,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async #handleMessage(msg: any) {
|
||||
const message = this._handler.parseMessage(msg);
|
||||
|
||||
switch (message.type) {
|
||||
case "TASK_RUN_COMPLETED": {
|
||||
const { result, execution } = message.payload;
|
||||
|
||||
const promiseStatus = this._attemptStatuses.get(execution.attempt.id);
|
||||
|
||||
if (promiseStatus !== "PENDING") {
|
||||
return;
|
||||
}
|
||||
|
||||
this._attemptStatuses.set(execution.attempt.id, "RESOLVED");
|
||||
|
||||
const attemptPromise = this._attemptPromises.get(execution.attempt.id);
|
||||
|
||||
if (!attemptPromise) {
|
||||
return;
|
||||
}
|
||||
|
||||
const { resolver } = attemptPromise;
|
||||
|
||||
resolver(result);
|
||||
|
||||
break;
|
||||
}
|
||||
case "READY_TO_DISPOSE": {
|
||||
this.#kill();
|
||||
|
||||
break;
|
||||
}
|
||||
case "TASK_HEARTBEAT": {
|
||||
this.onTaskHeartbeat.post(message.payload.id);
|
||||
|
||||
break;
|
||||
}
|
||||
case "TASKS_READY": {
|
||||
break;
|
||||
}
|
||||
case "WAIT_FOR_BATCH": {
|
||||
this.onWaitForBatch.post(message.payload);
|
||||
|
||||
break;
|
||||
}
|
||||
case "WAIT_FOR_DURATION": {
|
||||
this.onWaitForDuration.post(message.payload);
|
||||
|
||||
break;
|
||||
}
|
||||
case "WAIT_FOR_TASK": {
|
||||
this.onWaitForTask.post(message.payload);
|
||||
|
||||
break;
|
||||
}
|
||||
waitCompletedNotification() {
|
||||
if (this._child?.connected && !this._isBeingKilled && !this._child.killed) {
|
||||
this._ipc?.send("WAIT_COMPLETED_NOTIFICATION", {});
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -30,6 +30,9 @@ class ProdWorker {
|
||||
|
||||
private executing = false;
|
||||
private completed = false;
|
||||
private paused = false;
|
||||
|
||||
private nextResumeAfter: "WAIT_FOR_DURATION" | "WAIT_FOR_TASK" | "WAIT_FOR_BATCH" | undefined;
|
||||
|
||||
#httpPort: number;
|
||||
#backgroundWorker: ProdBackgroundWorker;
|
||||
@@ -55,18 +58,76 @@ class ProdWorker {
|
||||
},
|
||||
contentHash: this.contentHash,
|
||||
});
|
||||
|
||||
this.#backgroundWorker.onTaskHeartbeat.attach((attemptFriendlyId) => {
|
||||
this.#coordinatorSocket.send("TASK_HEARTBEAT", { attemptFriendlyId });
|
||||
});
|
||||
this.#backgroundWorker.onWaitForBatch.attach((message) => {
|
||||
this.#coordinatorSocket.send("WAIT_FOR_BATCH", message);
|
||||
// TODO: Switch to .send() once coordinator uses zod handler for all messages
|
||||
this.#coordinatorSocket.socket.emit("TASK_HEARTBEAT", { version: "v1", attemptFriendlyId });
|
||||
});
|
||||
|
||||
this.#backgroundWorker.onWaitForDuration.attach(async (message) => {
|
||||
const { success } = await this.#coordinatorSocket.sendWithAck("WAIT_FOR_DURATION", message);
|
||||
logger.log("WAIT_FOR_DURATION", { success });
|
||||
// TODO: Switch to .send() once coordinator uses zod handler for all messages
|
||||
const { willCheckpointAndRestore } = await this.#coordinatorSocket.socket.emitWithAck(
|
||||
"WAIT_FOR_DURATION",
|
||||
{ version: "v1", ...message }
|
||||
);
|
||||
|
||||
logger.log("WAIT_FOR_DURATION", { willCheckpointAndRestore });
|
||||
|
||||
this.#backgroundWorker.preCheckpointNotification.post({ willCheckpointAndRestore });
|
||||
|
||||
setTimeout(() => {
|
||||
if (willCheckpointAndRestore) {
|
||||
this.paused = true;
|
||||
this.nextResumeAfter = "WAIT_FOR_DURATION";
|
||||
}
|
||||
// Forcing a reconnect will ensure the connection handler runs to trigger automatic resume
|
||||
this.#coordinatorSocket.close();
|
||||
this.#coordinatorSocket.connect();
|
||||
}, 3_000);
|
||||
});
|
||||
this.#backgroundWorker.onWaitForTask.attach((message) => {
|
||||
this.#coordinatorSocket.send("WAIT_FOR_TASK", message);
|
||||
|
||||
this.#backgroundWorker.onWaitForTask.attach(async (message) => {
|
||||
// TODO: Switch to .send() once coordinator uses zod handler for all messages
|
||||
const { willCheckpointAndRestore } = await this.#coordinatorSocket.socket.emitWithAck(
|
||||
"WAIT_FOR_TASK",
|
||||
{ version: "v1", ...message }
|
||||
);
|
||||
|
||||
logger.log("WAIT_FOR_TASK", { willCheckpointAndRestore });
|
||||
|
||||
this.#backgroundWorker.preCheckpointNotification.post({ willCheckpointAndRestore });
|
||||
|
||||
setTimeout(() => {
|
||||
if (willCheckpointAndRestore) {
|
||||
this.paused = true;
|
||||
this.nextResumeAfter = "WAIT_FOR_TASK";
|
||||
}
|
||||
// Forcing a reconnect will ensure the connection handler runs to trigger automatic resume
|
||||
this.#coordinatorSocket.close();
|
||||
this.#coordinatorSocket.connect();
|
||||
}, 3_000);
|
||||
});
|
||||
|
||||
this.#backgroundWorker.onWaitForBatch.attach(async (message) => {
|
||||
// TODO: Switch to .send() once coordinator uses zod handler for all messages
|
||||
const { willCheckpointAndRestore } = await this.#coordinatorSocket.socket.emitWithAck(
|
||||
"WAIT_FOR_BATCH",
|
||||
{ version: "v1", ...message }
|
||||
);
|
||||
|
||||
logger.log("WAIT_FOR_BATCH", { willCheckpointAndRestore });
|
||||
|
||||
this.#backgroundWorker.preCheckpointNotification.post({ willCheckpointAndRestore });
|
||||
|
||||
setTimeout(() => {
|
||||
if (willCheckpointAndRestore) {
|
||||
this.paused = true;
|
||||
this.nextResumeAfter = "WAIT_FOR_BATCH";
|
||||
}
|
||||
// Forcing a reconnect will ensure the connection handler runs to trigger automatic resume
|
||||
this.#coordinatorSocket.close();
|
||||
this.#coordinatorSocket.connect();
|
||||
}, 3_000);
|
||||
});
|
||||
|
||||
this.#httpPort = port;
|
||||
@@ -112,11 +173,13 @@ class ProdWorker {
|
||||
this.#backgroundWorker.taskRunCompletedNotification(completion, execution);
|
||||
}
|
||||
},
|
||||
RESUME_AFTER_DURATION: async (message) => {
|
||||
this.#backgroundWorker.waitCompletedNotification();
|
||||
},
|
||||
EXECUTE_TASK_RUN: async (message) => {
|
||||
if (this.executing || this.completed) {
|
||||
return {
|
||||
success: false,
|
||||
};
|
||||
logger.error("dropping execute request, already executing or completed");
|
||||
return;
|
||||
}
|
||||
|
||||
this.executing = true;
|
||||
@@ -127,15 +190,13 @@ class ProdWorker {
|
||||
this.completed = true;
|
||||
this.executing = false;
|
||||
|
||||
setTimeout(() => {
|
||||
process.exit(0);
|
||||
}, 2000);
|
||||
|
||||
// TODO: replace ack with emit
|
||||
return {
|
||||
success: true,
|
||||
await this.#coordinatorSocket.socket.emitWithAck("TASK_RUN_COMPLETED", {
|
||||
version: "v1",
|
||||
execution: message.executionPayload.execution,
|
||||
completion,
|
||||
};
|
||||
});
|
||||
|
||||
process.exit(0);
|
||||
},
|
||||
},
|
||||
onConnection: async (socket, handler, sender, logger) => {
|
||||
@@ -209,12 +270,33 @@ class ProdWorker {
|
||||
process.exit(1);
|
||||
}, 200);
|
||||
}
|
||||
} else {
|
||||
socket.emit("READY_FOR_EXECUTION", {
|
||||
}
|
||||
|
||||
if (this.paused) {
|
||||
if (!this.nextResumeAfter) {
|
||||
return;
|
||||
}
|
||||
|
||||
socket.emit("READY_FOR_RESUME", {
|
||||
version: "v1",
|
||||
attemptId: this.attemptId,
|
||||
type: this.nextResumeAfter,
|
||||
});
|
||||
|
||||
this.#backgroundWorker.waitCompletedNotification();
|
||||
this.paused = false;
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
if (this.executing) {
|
||||
return;
|
||||
}
|
||||
|
||||
socket.emit("READY_FOR_EXECUTION", {
|
||||
version: "v1",
|
||||
attemptId: this.attemptId,
|
||||
});
|
||||
},
|
||||
});
|
||||
|
||||
@@ -247,11 +329,14 @@ class ProdWorker {
|
||||
return reply.text(this.contentHash);
|
||||
|
||||
case "/wait":
|
||||
const { success } = await this.#coordinatorSocket.sendWithAck("WAIT_FOR_DURATION", {
|
||||
version: "v1",
|
||||
ms: 60_000,
|
||||
});
|
||||
logger.log("WAIT_FOR_DURATION", { success });
|
||||
const { willCheckpointAndRestore } = await this.#coordinatorSocket.sendWithAck(
|
||||
"WAIT_FOR_DURATION",
|
||||
{
|
||||
version: "v1",
|
||||
ms: 60_000,
|
||||
}
|
||||
);
|
||||
logger.log("WAIT_FOR_DURATION", { willCheckpointAndRestore });
|
||||
// this is required when C/Ring established connections
|
||||
this.#coordinatorSocket.close();
|
||||
return reply.text("sent WAIT");
|
||||
|
||||
@@ -1,4 +1,9 @@
|
||||
import { type TracingSDK } from "@trigger.dev/core/v3";
|
||||
import {
|
||||
ProdChildToWorkerMessages,
|
||||
ProdWorkerToChildMessages,
|
||||
ZodIpcConnection,
|
||||
type TracingSDK,
|
||||
} from "@trigger.dev/core/v3";
|
||||
import "source-map-support/register.js";
|
||||
|
||||
__WORKER_SETUP__;
|
||||
@@ -20,16 +25,12 @@ import {
|
||||
TaskRunExecution,
|
||||
TaskRunExecutionRetry,
|
||||
TriggerTracer,
|
||||
ZodMessageHandler,
|
||||
ZodMessageSender,
|
||||
accessoryAttributes,
|
||||
calculateNextRetryDelay,
|
||||
childToWorkerMessages,
|
||||
logger,
|
||||
parseError,
|
||||
runtime,
|
||||
taskContextManager,
|
||||
workerToChildMessages,
|
||||
type BackgroundWorkerProperties,
|
||||
} from "@trigger.dev/core/v3";
|
||||
import * as packageJson from "../../../package.json";
|
||||
@@ -40,12 +41,6 @@ import { TaskMetadataWithFunctions } from "../../types";
|
||||
const tracer = new TriggerTracer({ tracer: otelTracer, logger: otelLogger });
|
||||
const consoleInterceptor = new ConsoleInterceptor(otelLogger);
|
||||
|
||||
declare const sender: ZodMessageSender<typeof childToWorkerMessages>;
|
||||
|
||||
const prodRuntimeManager = new ProdRuntimeManager(sender);
|
||||
|
||||
runtime.setGlobalRuntimeManager(prodRuntimeManager);
|
||||
|
||||
const otelTaskLogger = new OtelTaskLogger({
|
||||
logger: otelLogger,
|
||||
tracer: tracer,
|
||||
@@ -233,10 +228,12 @@ for (const task of tasks) {
|
||||
let _execution: TaskRunExecution | undefined;
|
||||
let _isRunning = false;
|
||||
|
||||
const handler = new ZodMessageHandler({
|
||||
schema: workerToChildMessages,
|
||||
messages: {
|
||||
EXECUTE_TASK_RUN: async ({ execution, traceContext, metadata }) => {
|
||||
const zodIpc = new ZodIpcConnection({
|
||||
listenSchema: ProdWorkerToChildMessages,
|
||||
emitSchema: ProdChildToWorkerMessages,
|
||||
process,
|
||||
handlers: {
|
||||
EXECUTE_TASK_RUN: async ({ execution, traceContext, metadata }, sender) => {
|
||||
if (_isRunning) {
|
||||
console.error("Worker is already running a task");
|
||||
|
||||
@@ -308,7 +305,10 @@ const handler = new ZodMessageHandler({
|
||||
TASK_RUN_COMPLETED_NOTIFICATION: async ({ completion, execution }) => {
|
||||
prodRuntimeManager.resumeTask(completion, execution);
|
||||
},
|
||||
CLEANUP: async ({ flush, kill }) => {
|
||||
WAIT_COMPLETED_NOTIFICATION: async () => {
|
||||
prodRuntimeManager.resumeAfterRestore();
|
||||
},
|
||||
CLEANUP: async ({ flush, kill }, sender) => {
|
||||
if (kill) {
|
||||
await tracingSDK.flush();
|
||||
// Now we need to exit the process
|
||||
@@ -322,11 +322,11 @@ const handler = new ZodMessageHandler({
|
||||
},
|
||||
});
|
||||
|
||||
process.on("message", async (msg: any) => {
|
||||
await handler.handleMessage(msg);
|
||||
});
|
||||
const prodRuntimeManager = new ProdRuntimeManager(zodIpc);
|
||||
|
||||
sender.send("TASKS_READY", { tasks: getTaskMetadata() }).catch((err) => {
|
||||
runtime.setGlobalRuntimeManager(prodRuntimeManager);
|
||||
|
||||
zodIpc.send("TASKS_READY", { tasks: getTaskMetadata() }).catch((err) => {
|
||||
console.error("Failed to send TASKS_READY message", err);
|
||||
});
|
||||
|
||||
@@ -337,7 +337,7 @@ async function asyncHeartbeat(initialDelayInSeconds: number = 30, intervalInSeco
|
||||
while (true) {
|
||||
if (_isRunning && _execution) {
|
||||
try {
|
||||
await sender.send("TASK_HEARTBEAT", { id: _execution.attempt.id });
|
||||
await zodIpc.send("TASK_HEARTBEAT", { id: _execution.attempt.id });
|
||||
} catch (err) {
|
||||
console.error("Failed to send HEARTBEAT message", err);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
import { Resource } from "@opentelemetry/resources";
|
||||
import { OpenAIInstrumentation } from "@traceloop/instrumentation-openai";
|
||||
import {
|
||||
SemanticInternalAttributes,
|
||||
TracingDiagnosticLogLevel,
|
||||
TracingSDK,
|
||||
ZodMessageSender,
|
||||
childToWorkerMessages,
|
||||
} from "@trigger.dev/core/v3";
|
||||
|
||||
export const tracingSDK = new TracingSDK({
|
||||
url: process.env.OTEL_EXPORTER_OTLP_ENDPOINT ?? "http://0.0.0.0:4318",
|
||||
resource: new Resource({
|
||||
[SemanticInternalAttributes.CLI_VERSION]: "3.0.0",
|
||||
}),
|
||||
instrumentations: [new OpenAIInstrumentation()],
|
||||
diagLogLevel: (process.env.OTEL_LOG_LEVEL as TracingDiagnosticLogLevel) ?? "none",
|
||||
});
|
||||
|
||||
process.on("uncaughtException", (error, origin) => {
|
||||
process.send?.({
|
||||
type: "EVENT",
|
||||
message: {
|
||||
type: "UNCAUGHT_EXCEPTION",
|
||||
payload: {
|
||||
error: {
|
||||
name: error.name,
|
||||
message: error.message,
|
||||
stack: error.stack,
|
||||
},
|
||||
origin,
|
||||
},
|
||||
version: "v1",
|
||||
},
|
||||
});
|
||||
});
|
||||
@@ -6,8 +6,9 @@ export default defineConfig({
|
||||
tsconfig: "tsconfig.json",
|
||||
splitting: false,
|
||||
entry: [
|
||||
"src/workers/dev/worker-setup.ts",
|
||||
"src/workers/dev/worker-facade.ts",
|
||||
"src/workers/common/worker-setup.ts",
|
||||
"src/workers/prod/worker-setup.ts",
|
||||
"src/workers/prod/worker-facade.ts",
|
||||
"src/workers/prod/entry-point.ts",
|
||||
],
|
||||
|
||||
@@ -30,6 +30,7 @@ export interface TaskOperations {
|
||||
|
||||
type ProviderShellOptions = {
|
||||
tasks: TaskOperations;
|
||||
type: "docker" | "kubernetes";
|
||||
host?: string;
|
||||
port?: number;
|
||||
};
|
||||
@@ -111,7 +112,7 @@ export class ProviderShell implements Provider {
|
||||
serverMessages: PlatformToProviderMessages,
|
||||
authToken: PLATFORM_SECRET,
|
||||
extraHeaders: {
|
||||
"x-trigger-provider-type": "docker",
|
||||
"x-trigger-provider-type": this.options.type,
|
||||
},
|
||||
handlers: {
|
||||
DELETE: async (message) => {
|
||||
@@ -139,10 +140,28 @@ export class ProviderShell implements Provider {
|
||||
apiUrl: message.apiUrl,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error("task index failed", error);
|
||||
logger.error("index failed", error);
|
||||
}
|
||||
},
|
||||
RESTORE: async (message) => {
|
||||
if (message.type.toLowerCase() !== this.options.type.toLowerCase()) {
|
||||
logger.error(
|
||||
`restore failed: ${this.options.type} provider can't restore ${message.type} checkpoints`
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
await this.tasks.restore({
|
||||
attemptId: message.attemptId,
|
||||
checkpointRef: message.location,
|
||||
// TODO
|
||||
// machine: message.machine,
|
||||
});
|
||||
} catch (error) {
|
||||
logger.error("restore failed", error);
|
||||
}
|
||||
},
|
||||
RESTORE: async (message) => {},
|
||||
},
|
||||
});
|
||||
|
||||
|
||||
@@ -5,6 +5,7 @@ export * from "./apiClient";
|
||||
export * from "./zodMessageHandler";
|
||||
export * from "./zodNamespace";
|
||||
export * from "./zodSocket";
|
||||
export * from "./zodIpc";
|
||||
export * from "./errors";
|
||||
export * from "./runtime-api";
|
||||
export * from "./logger-api";
|
||||
|
||||
@@ -1,12 +1,13 @@
|
||||
import {
|
||||
BatchTaskRunExecutionResult,
|
||||
ProdChildToWorkerMessages,
|
||||
ProdWorkerToChildMessages,
|
||||
TaskMetadataWithFilePath,
|
||||
TaskRunContext,
|
||||
TaskRunExecution,
|
||||
TaskRunExecutionResult,
|
||||
childToWorkerMessages,
|
||||
} from "../schemas";
|
||||
import { ZodMessageSender } from "../zodMessageHandler";
|
||||
import { ZodIpcConnection } from "../zodIpc";
|
||||
import { RuntimeManager } from "./manager";
|
||||
|
||||
export class ProdRuntimeManager implements RuntimeManager {
|
||||
@@ -20,9 +21,16 @@ export class ProdRuntimeManager implements RuntimeManager {
|
||||
{ resolve: (value: BatchTaskRunExecutionResult) => void; reject: (err?: any) => void }
|
||||
> = new Map();
|
||||
|
||||
_waitForRestore: { resolve: (value?: any) => void; reject: (err?: any) => void } | undefined;
|
||||
|
||||
_tasks: Map<string, TaskMetadataWithFilePath> = new Map();
|
||||
|
||||
constructor(private sender: ZodMessageSender<typeof childToWorkerMessages>) {}
|
||||
constructor(
|
||||
private ipc: ZodIpcConnection<
|
||||
typeof ProdWorkerToChildMessages,
|
||||
typeof ProdChildToWorkerMessages
|
||||
>
|
||||
) {}
|
||||
|
||||
disable(): void {
|
||||
// do nothing
|
||||
@@ -39,15 +47,42 @@ export class ProdRuntimeManager implements RuntimeManager {
|
||||
}
|
||||
|
||||
async waitForDuration(ms: number): Promise<void> {
|
||||
if (ms > 30_000) {
|
||||
// TODO: sender with ack support
|
||||
await this.sender.send("WAIT_FOR_DURATION", { ms });
|
||||
// TODO: resolve after resume signal instead
|
||||
let timeout: NodeJS.Timeout | undefined;
|
||||
|
||||
const resolveAfterDuration = new Promise((resolve) => {
|
||||
timeout = setTimeout(resolve, ms);
|
||||
});
|
||||
|
||||
if (ms < 10_000) {
|
||||
await resolveAfterDuration;
|
||||
return;
|
||||
}
|
||||
|
||||
return new Promise((resolve) => {
|
||||
setTimeout(resolve, ms);
|
||||
const waitForRestore = new Promise<TaskRunExecutionResult>((resolve, reject) => {
|
||||
this._waitForRestore = { resolve, reject };
|
||||
});
|
||||
|
||||
// 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 });
|
||||
|
||||
if (!willCheckpointAndRestore) {
|
||||
await resolveAfterDuration;
|
||||
return;
|
||||
}
|
||||
|
||||
// Checkpointing should happen after this line
|
||||
|
||||
await waitForRestore;
|
||||
clearTimeout(timeout);
|
||||
}
|
||||
|
||||
resumeAfterRestore(): void {
|
||||
if (!this._waitForRestore) {
|
||||
return;
|
||||
}
|
||||
|
||||
this._waitForRestore.resolve();
|
||||
this._waitForRestore = undefined;
|
||||
}
|
||||
|
||||
async waitUntil(date: Date): Promise<void> {
|
||||
@@ -59,7 +94,7 @@ export class ProdRuntimeManager implements RuntimeManager {
|
||||
this._taskWaits.set(params.id, { resolve, reject });
|
||||
});
|
||||
|
||||
await this.sender.send("WAIT_FOR_TASK", {
|
||||
await this.ipc.send("WAIT_FOR_TASK", {
|
||||
id: params.id,
|
||||
});
|
||||
|
||||
@@ -83,7 +118,7 @@ export class ProdRuntimeManager implements RuntimeManager {
|
||||
})
|
||||
);
|
||||
|
||||
await this.sender.send("WAIT_FOR_BATCH", {
|
||||
await this.ipc.send("WAIT_FOR_BATCH", {
|
||||
id: params.id,
|
||||
runs: params.runs,
|
||||
});
|
||||
|
||||
@@ -187,6 +187,16 @@ export const TaskMetadataWithFilePath = TaskMetadata.extend({
|
||||
|
||||
export type TaskMetadataWithFilePath = z.infer<typeof TaskMetadataWithFilePath>;
|
||||
|
||||
export const UncaughtExceptionMessage = z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
error: z.object({
|
||||
name: z.string(),
|
||||
message: z.string(),
|
||||
stack: z.string().optional(),
|
||||
}),
|
||||
origin: z.enum(["uncaughtException", "unhandledRejection"]),
|
||||
});
|
||||
|
||||
export const childToWorkerMessages = {
|
||||
TASK_RUN_COMPLETED: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
@@ -215,13 +225,85 @@ export const childToWorkerMessages = {
|
||||
id: z.string(),
|
||||
runs: z.string().array(),
|
||||
}),
|
||||
UNCAUGHT_EXCEPTION: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
error: z.object({
|
||||
name: z.string(),
|
||||
message: z.string(),
|
||||
stack: z.string().optional(),
|
||||
}),
|
||||
origin: z.enum(["uncaughtException", "unhandledRejection"]),
|
||||
}),
|
||||
UNCAUGHT_EXCEPTION: UncaughtExceptionMessage,
|
||||
};
|
||||
|
||||
export const ProdChildToWorkerMessages = {
|
||||
TASK_RUN_COMPLETED: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
execution: TaskRunExecution,
|
||||
result: TaskRunExecutionResult,
|
||||
}),
|
||||
},
|
||||
TASKS_READY: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
tasks: TaskMetadataWithFilePath.array(),
|
||||
}),
|
||||
},
|
||||
TASK_HEARTBEAT: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
id: z.string(),
|
||||
}),
|
||||
},
|
||||
READY_TO_DISPOSE: {
|
||||
message: z.undefined(),
|
||||
},
|
||||
WAIT_FOR_DURATION: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
ms: z.number(),
|
||||
}),
|
||||
callback: z.object({
|
||||
willCheckpointAndRestore: z.boolean(),
|
||||
}),
|
||||
},
|
||||
WAIT_FOR_TASK: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
id: z.string(),
|
||||
}),
|
||||
},
|
||||
WAIT_FOR_BATCH: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
id: z.string(),
|
||||
runs: z.string().array(),
|
||||
}),
|
||||
},
|
||||
UNCAUGHT_EXCEPTION: {
|
||||
message: UncaughtExceptionMessage,
|
||||
},
|
||||
};
|
||||
|
||||
export const ProdWorkerToChildMessages = {
|
||||
EXECUTE_TASK_RUN: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
execution: TaskRunExecution,
|
||||
traceContext: z.record(z.unknown()),
|
||||
metadata: BackgroundWorkerProperties,
|
||||
}),
|
||||
},
|
||||
TASK_RUN_COMPLETED_NOTIFICATION: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
completion: TaskRunExecutionResult,
|
||||
execution: TaskRunExecution,
|
||||
}),
|
||||
},
|
||||
CLEANUP: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
flush: z.boolean().default(false),
|
||||
kill: z.boolean().default(true),
|
||||
}),
|
||||
},
|
||||
WAIT_COMPLETED_NOTIFICATION: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
}),
|
||||
},
|
||||
};
|
||||
|
||||
@@ -145,6 +145,13 @@ export const CoordinatorToPlatformMessages = {
|
||||
}),
|
||||
]),
|
||||
},
|
||||
READY_FOR_RESUME: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
attemptId: z.string(),
|
||||
type: z.enum(["WAIT_FOR_DURATION", "WAIT_FOR_TASK", "WAIT_FOR_BATCH"]),
|
||||
}),
|
||||
},
|
||||
TASK_RUN_COMPLETED: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
@@ -164,7 +171,20 @@ export const CoordinatorToPlatformMessages = {
|
||||
attemptId: z.string(),
|
||||
docker: z.boolean(),
|
||||
location: z.string(),
|
||||
reason: z.string().optional(),
|
||||
reason: z.discriminatedUnion("type", [
|
||||
z.object({
|
||||
type: z.literal("WAIT_FOR_DURATION"),
|
||||
ms: z.number(),
|
||||
}),
|
||||
z.object({
|
||||
type: z.literal("WAIT_FOR_BATCH"),
|
||||
id: z.string(),
|
||||
}),
|
||||
z.object({
|
||||
type: z.literal("WAIT_FOR_TASK"),
|
||||
id: z.string(),
|
||||
}),
|
||||
]),
|
||||
}),
|
||||
},
|
||||
INDEXING_FAILED: {
|
||||
@@ -185,11 +205,16 @@ export const PlatformToCoordinatorMessages = {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
attemptId: z.string(),
|
||||
image: z.string(),
|
||||
completions: TaskRunExecutionResult.array(),
|
||||
executions: TaskRunExecution.array(),
|
||||
}),
|
||||
},
|
||||
RESUME_AFTER_DURATION: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
attemptId: z.string(),
|
||||
}),
|
||||
},
|
||||
};
|
||||
|
||||
export const ClientToSharedQueueMessages = {
|
||||
@@ -260,37 +285,53 @@ export const ProdWorkerToCoordinatorMessages = {
|
||||
attemptId: z.string(),
|
||||
}),
|
||||
},
|
||||
READY_FOR_RESUME: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
attemptId: z.string(),
|
||||
type: z.enum(["WAIT_FOR_DURATION", "WAIT_FOR_TASK", "WAIT_FOR_BATCH"]),
|
||||
}),
|
||||
},
|
||||
TASK_HEARTBEAT: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
attemptFriendlyId: z.string(),
|
||||
}),
|
||||
},
|
||||
TASK_RUN_COMPLETED: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
execution: ProdTaskRunExecution,
|
||||
completion: TaskRunExecutionResult,
|
||||
}),
|
||||
callback: z.void(),
|
||||
},
|
||||
WAIT_FOR_DURATION: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
ms: z.number(),
|
||||
}),
|
||||
callback: z.object({
|
||||
willCheckpointAndRestore: z.boolean(),
|
||||
}),
|
||||
},
|
||||
WAIT_FOR_TASK: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
id: z.string(),
|
||||
}),
|
||||
callback: z.object({
|
||||
willCheckpointAndRestore: z.boolean(),
|
||||
}),
|
||||
},
|
||||
WAIT_FOR_BATCH: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
id: z.string(),
|
||||
runs: z.string().array(),
|
||||
}),
|
||||
},
|
||||
WAIT_FOR_DURATION: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
ms: z.number(),
|
||||
}),
|
||||
callback: z.discriminatedUnion("success", [
|
||||
z.object({
|
||||
success: z.literal(false),
|
||||
}),
|
||||
z.object({
|
||||
success: z.literal(true),
|
||||
}),
|
||||
]),
|
||||
},
|
||||
WAIT_FOR_TASK: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
id: z.string(),
|
||||
callback: z.object({
|
||||
willCheckpointAndRestore: z.boolean(),
|
||||
}),
|
||||
},
|
||||
INDEXING_FAILED: {
|
||||
@@ -311,25 +352,21 @@ export const CoordinatorToProdWorkerMessages = {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
attemptId: z.string(),
|
||||
image: z.string(),
|
||||
completions: TaskRunExecutionResult.array(),
|
||||
executions: TaskRunExecution.array(),
|
||||
}),
|
||||
},
|
||||
RESUME_AFTER_DURATION: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
attemptId: z.string(),
|
||||
}),
|
||||
},
|
||||
EXECUTE_TASK_RUN: {
|
||||
message: z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
executionPayload: ProdTaskRunExecutionPayload,
|
||||
}),
|
||||
callback: z.discriminatedUnion("success", [
|
||||
z.object({
|
||||
success: z.literal(false),
|
||||
}),
|
||||
z.object({
|
||||
success: z.literal(true),
|
||||
completion: TaskRunExecutionResult,
|
||||
}),
|
||||
]),
|
||||
},
|
||||
};
|
||||
|
||||
|
||||
@@ -0,0 +1,328 @@
|
||||
import { randomUUID } from "crypto";
|
||||
import {
|
||||
GetSocketCallbackSchema,
|
||||
GetSocketMessageSchema,
|
||||
GetSocketMessagesWithCallback,
|
||||
GetSocketMessagesWithoutCallback,
|
||||
MessagesFromSocketCatalog,
|
||||
SocketMessageHasCallback,
|
||||
ZodSocketMessageCatalogSchema,
|
||||
} from "./zodSocket";
|
||||
import { z } from "zod";
|
||||
|
||||
interface ZodIpcMessageSender<TEmitCatalog extends ZodSocketMessageCatalogSchema> {
|
||||
send<K extends GetSocketMessagesWithoutCallback<TEmitCatalog>>(
|
||||
type: K,
|
||||
payload: z.input<GetSocketMessageSchema<TEmitCatalog, K>>
|
||||
): Promise<void>;
|
||||
|
||||
sendWithAck<K extends GetSocketMessagesWithCallback<TEmitCatalog>>(
|
||||
type: K,
|
||||
payload: z.input<GetSocketMessageSchema<TEmitCatalog, K>>
|
||||
): Promise<z.infer<GetSocketCallbackSchema<TEmitCatalog, K>>>;
|
||||
}
|
||||
|
||||
type ZodIpcMessageHandlers<
|
||||
TListenCatalog extends ZodSocketMessageCatalogSchema,
|
||||
TEmitCatalog extends ZodSocketMessageCatalogSchema,
|
||||
> = Partial<{
|
||||
[K in keyof TListenCatalog]: (
|
||||
payload: z.infer<GetSocketMessageSchema<TListenCatalog, K>>,
|
||||
sender: ZodIpcMessageSender<TEmitCatalog>
|
||||
) => Promise<
|
||||
SocketMessageHasCallback<TListenCatalog, K> extends true
|
||||
? z.input<GetSocketCallbackSchema<TListenCatalog, K>>
|
||||
: void
|
||||
>;
|
||||
}>;
|
||||
|
||||
const messageSchema = z.object({
|
||||
version: z.literal("v1").default("v1"),
|
||||
type: z.string(),
|
||||
payload: z.unknown(),
|
||||
});
|
||||
|
||||
type ZodIpcMessageHandlerOptions<
|
||||
TListenCatalog extends ZodSocketMessageCatalogSchema,
|
||||
TEmitCatalog extends ZodSocketMessageCatalogSchema,
|
||||
> = {
|
||||
schema: TListenCatalog;
|
||||
handlers?: ZodIpcMessageHandlers<TListenCatalog, TEmitCatalog>;
|
||||
sender: ZodIpcMessageSender<TEmitCatalog>;
|
||||
};
|
||||
|
||||
class ZodIpcMessageHandler<
|
||||
TListenCatalog extends ZodSocketMessageCatalogSchema,
|
||||
TEmitCatalog extends ZodSocketMessageCatalogSchema,
|
||||
> {
|
||||
#schema: TListenCatalog;
|
||||
#handlers: ZodIpcMessageHandlers<TListenCatalog, TEmitCatalog> | undefined;
|
||||
#sender: ZodIpcMessageSender<TEmitCatalog>;
|
||||
|
||||
constructor(options: ZodIpcMessageHandlerOptions<TListenCatalog, TEmitCatalog>) {
|
||||
this.#schema = options.schema;
|
||||
this.#handlers = options.handlers;
|
||||
this.#sender = options.sender;
|
||||
}
|
||||
|
||||
public async handleMessage(message: unknown) {
|
||||
const parsedMessage = this.parseMessage(message);
|
||||
|
||||
if (!this.#handlers) {
|
||||
throw new Error("No handlers provided");
|
||||
}
|
||||
|
||||
const handler = this.#handlers[parsedMessage.type];
|
||||
|
||||
if (!handler) {
|
||||
// console.error(`No handler for message type: ${String(parsedMessage.type)}`);
|
||||
return;
|
||||
}
|
||||
|
||||
const ack = await handler(parsedMessage.payload, this.#sender);
|
||||
|
||||
return ack;
|
||||
}
|
||||
|
||||
public parseMessage(message: unknown): MessagesFromSocketCatalog<TListenCatalog> {
|
||||
const parsedMessage = messageSchema.safeParse(message);
|
||||
|
||||
if (!parsedMessage.success) {
|
||||
throw new Error(`Failed to parse message: ${JSON.stringify(parsedMessage.error)}`);
|
||||
}
|
||||
|
||||
const schema = this.#schema[parsedMessage.data.type]["message"];
|
||||
|
||||
if (!schema) {
|
||||
throw new Error(`Unknown message type: ${parsedMessage.data.type}`);
|
||||
}
|
||||
|
||||
const parsedPayload = schema.safeParse(parsedMessage.data.payload);
|
||||
|
||||
if (!parsedPayload.success) {
|
||||
throw new Error(`Failed to parse message payload: ${JSON.stringify(parsedPayload.error)}`);
|
||||
}
|
||||
|
||||
return {
|
||||
type: parsedMessage.data.type,
|
||||
payload: parsedPayload.data,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
const Packet = z.discriminatedUnion("type", [
|
||||
z.object({
|
||||
type: z.literal("CONNECT"),
|
||||
sessionId: z.string().optional(),
|
||||
}),
|
||||
z.object({
|
||||
type: z.literal("ACK"),
|
||||
message: z.any(),
|
||||
id: z.number(),
|
||||
}),
|
||||
z.object({
|
||||
type: z.literal("EVENT"),
|
||||
message: z.any(),
|
||||
id: z.number().optional(),
|
||||
}),
|
||||
]);
|
||||
|
||||
type Packet = z.infer<typeof Packet>;
|
||||
|
||||
interface ZodIpcConnectionOptions<
|
||||
TListenCatalog extends ZodSocketMessageCatalogSchema,
|
||||
TEmitCatalog extends ZodSocketMessageCatalogSchema,
|
||||
> {
|
||||
listenSchema: TListenCatalog;
|
||||
emitSchema: TEmitCatalog;
|
||||
process: {
|
||||
send?: (message: any) => any;
|
||||
on?: (event: "message", listener: (message: any) => void) => void;
|
||||
};
|
||||
handlers?: ZodIpcMessageHandlers<TListenCatalog, TEmitCatalog>;
|
||||
}
|
||||
|
||||
export class ZodIpcConnection<
|
||||
TListenCatalog extends ZodSocketMessageCatalogSchema,
|
||||
TEmitCatalog extends ZodSocketMessageCatalogSchema,
|
||||
> {
|
||||
#sessionId?: string;
|
||||
#messageCounter: number = 0;
|
||||
|
||||
#handler: ZodIpcMessageHandler<TListenCatalog, TEmitCatalog>;
|
||||
|
||||
#acks: Map<
|
||||
number,
|
||||
{
|
||||
resolve: (value: unknown) => void;
|
||||
reject: (reason?: any) => void;
|
||||
timeout: NodeJS.Timeout;
|
||||
}
|
||||
> = new Map();
|
||||
|
||||
constructor(private opts: ZodIpcConnectionOptions<TListenCatalog, TEmitCatalog>) {
|
||||
this.#handler = new ZodIpcMessageHandler({
|
||||
schema: opts.listenSchema,
|
||||
handlers: opts.handlers,
|
||||
sender: {
|
||||
send: this.send.bind(this),
|
||||
sendWithAck: this.sendWithAck.bind(this),
|
||||
},
|
||||
});
|
||||
|
||||
this.#registerHandlers();
|
||||
// this.connect();
|
||||
}
|
||||
|
||||
async #registerHandlers() {
|
||||
if (!this.opts.process.on) {
|
||||
return;
|
||||
}
|
||||
|
||||
this.opts.process.on("message", async (message) => {
|
||||
this.#handlePacket(message);
|
||||
});
|
||||
}
|
||||
|
||||
async connect() {
|
||||
this.#sendPacket({ type: "CONNECT" });
|
||||
}
|
||||
|
||||
async #handlePacket(packet: Packet): Promise<void> {
|
||||
const parsedPacket = Packet.safeParse(packet);
|
||||
|
||||
if (!parsedPacket.success) {
|
||||
// console.error("dropping invalid packet", packet);
|
||||
return;
|
||||
}
|
||||
|
||||
// console.log("<-", packet);
|
||||
|
||||
switch (parsedPacket.data.type) {
|
||||
case "ACK": {
|
||||
// Check our list of ACKs and resolve with the message
|
||||
const ack = this.#acks.get(parsedPacket.data.id);
|
||||
|
||||
if (!ack) {
|
||||
return;
|
||||
}
|
||||
|
||||
clearTimeout(ack.timeout);
|
||||
ack.resolve(parsedPacket.data.message);
|
||||
|
||||
break;
|
||||
}
|
||||
case "CONNECT": {
|
||||
if (!parsedPacket.data.sessionId) {
|
||||
// This is a client trying to connect, so we generate and send back a session ID
|
||||
const id = randomUUID();
|
||||
|
||||
await this.#sendPacket({ type: "CONNECT", sessionId: id });
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
// This is a server replying to our connect message
|
||||
if (this.#sessionId) {
|
||||
// We're already connected
|
||||
return;
|
||||
}
|
||||
|
||||
this.#sessionId = parsedPacket.data.sessionId;
|
||||
|
||||
break;
|
||||
}
|
||||
case "EVENT": {
|
||||
const result = await this.#handler.handleMessage(parsedPacket.data.message);
|
||||
|
||||
if (typeof parsedPacket.data.id === "undefined") {
|
||||
return;
|
||||
}
|
||||
|
||||
// There's an ID so we should ACK
|
||||
await this.#sendPacket({
|
||||
type: "ACK",
|
||||
id: parsedPacket.data.id,
|
||||
message: result,
|
||||
});
|
||||
|
||||
break;
|
||||
}
|
||||
default: {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async #sendPacket(packet: Packet) {
|
||||
// console.log("->", packet);
|
||||
await this.opts.process.send?.(packet);
|
||||
}
|
||||
|
||||
async send<K extends GetSocketMessagesWithoutCallback<TEmitCatalog>>(
|
||||
type: K,
|
||||
payload: z.input<GetSocketMessageSchema<TEmitCatalog, K>>
|
||||
): Promise<void> {
|
||||
const schema = this.opts.emitSchema[type]["message"];
|
||||
|
||||
if (!schema) {
|
||||
throw new Error(`Unknown message type: ${type as string}`);
|
||||
}
|
||||
|
||||
const parsedPayload = schema.safeParse(payload);
|
||||
|
||||
if (!parsedPayload.success) {
|
||||
throw new Error(`Failed to parse message payload: ${JSON.stringify(parsedPayload.error)}`);
|
||||
}
|
||||
|
||||
await this.#sendPacket({
|
||||
type: "EVENT",
|
||||
message: {
|
||||
type,
|
||||
payload,
|
||||
version: "v1",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
public async sendWithAck<K extends GetSocketMessagesWithCallback<TEmitCatalog>>(
|
||||
type: K,
|
||||
payload: z.input<GetSocketMessageSchema<TEmitCatalog, K>>,
|
||||
timeoutInMs?: number
|
||||
): Promise<z.infer<GetSocketCallbackSchema<TEmitCatalog, K>>> {
|
||||
const currentId = this.#messageCounter++;
|
||||
|
||||
return new Promise(async (resolve, reject) => {
|
||||
// Timeout if the ACK takes too long to get back to us
|
||||
const timeout = setTimeout(() => {
|
||||
reject("timeout");
|
||||
}, timeoutInMs ?? 2000);
|
||||
|
||||
this.#acks.set(currentId, { resolve, reject, timeout });
|
||||
|
||||
const schema = this.opts.emitSchema[type]["message"];
|
||||
|
||||
if (!schema) {
|
||||
clearTimeout(timeout);
|
||||
return reject(`Unknown message type: ${type as string}`);
|
||||
}
|
||||
|
||||
const parsedPayload = schema.safeParse(payload);
|
||||
|
||||
if (!parsedPayload.success) {
|
||||
clearTimeout(timeout);
|
||||
return reject(`Failed to parse message payload: ${JSON.stringify(parsedPayload.error)}`);
|
||||
}
|
||||
|
||||
await this.#sendPacket({
|
||||
type: "EVENT",
|
||||
message: {
|
||||
type,
|
||||
payload,
|
||||
version: "v1",
|
||||
},
|
||||
id: currentId,
|
||||
});
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -96,6 +96,7 @@ export class ZodNamespace<
|
||||
|
||||
this.namespace = this.io.of(opts.name);
|
||||
|
||||
// FIXME: There's a bug here, this sender should not accept Socket schemas with callbacks
|
||||
this.sender = new ZodMessageSender({
|
||||
schema: opts.serverMessages,
|
||||
sender: async (message) => {
|
||||
|
||||
@@ -22,7 +22,7 @@ export type ZodMessageCatalogToSocketIoEvents<TCatalog extends ZodSocketMessageC
|
||||
: (message: z.infer<GetSocketMessageSchema<TCatalog, K>>) => void;
|
||||
};
|
||||
|
||||
type GetSocketMessageSchema<
|
||||
export type GetSocketMessageSchema<
|
||||
TRPCCatalog extends ZodSocketMessageCatalogSchema,
|
||||
TMessageType extends keyof TRPCCatalog,
|
||||
> = TRPCCatalog[TMessageType]["message"];
|
||||
@@ -32,7 +32,7 @@ export type InferSocketMessageSchema<
|
||||
TMessageType extends keyof TRPCCatalog,
|
||||
> = z.infer<GetSocketMessageSchema<TRPCCatalog, TMessageType>>;
|
||||
|
||||
type GetSocketCallbackSchema<
|
||||
export type GetSocketCallbackSchema<
|
||||
TRPCCatalog extends ZodSocketMessageCatalogSchema,
|
||||
TMessageType extends keyof TRPCCatalog,
|
||||
> = TRPCCatalog[TMessageType] extends { callback: any }
|
||||
@@ -44,7 +44,7 @@ export type InferSocketCallbackSchema<
|
||||
TMessageType extends keyof TRPCCatalog,
|
||||
> = z.infer<GetSocketCallbackSchema<TRPCCatalog, TMessageType>>;
|
||||
|
||||
type SocketMessageHasCallback<
|
||||
export type SocketMessageHasCallback<
|
||||
TRPCCatalog extends ZodSocketMessageCatalogSchema,
|
||||
TMessageType extends keyof TRPCCatalog,
|
||||
> = GetSocketCallbackSchema<TRPCCatalog, TMessageType> extends never ? false : true;
|
||||
@@ -74,7 +74,7 @@ type MessageFromSocketSchema<
|
||||
payload: z.input<GetSocketMessageSchema<TMessageCatalog, K>>;
|
||||
};
|
||||
|
||||
type MessagesFromSocketCatalog<TMessageCatalog extends ZodSocketMessageCatalogSchema> = {
|
||||
export type MessagesFromSocketCatalog<TMessageCatalog extends ZodSocketMessageCatalogSchema> = {
|
||||
[K in keyof TMessageCatalog]: MessageFromSocketSchema<K, TMessageCatalog>;
|
||||
}[keyof TMessageCatalog];
|
||||
|
||||
@@ -173,13 +173,15 @@ export type ZodSocketMessageSenderOptions<TMessageCatalog extends ZodSocketMessa
|
||||
socket: ZodSocket<any, TMessageCatalog>;
|
||||
};
|
||||
|
||||
type GetSocketMessagesWithCallback<TMessageCatalog extends ZodSocketMessageCatalogSchema> = {
|
||||
export type GetSocketMessagesWithCallback<TMessageCatalog extends ZodSocketMessageCatalogSchema> = {
|
||||
[K in keyof TMessageCatalog]: SocketMessageHasCallback<TMessageCatalog, K> extends true
|
||||
? K
|
||||
: never;
|
||||
}[keyof TMessageCatalog];
|
||||
|
||||
type GetSocketMessagesWithoutCallback<TMessageCatalog extends ZodSocketMessageCatalogSchema> = {
|
||||
export type GetSocketMessagesWithoutCallback<
|
||||
TMessageCatalog extends ZodSocketMessageCatalogSchema,
|
||||
> = {
|
||||
[K in keyof TMessageCatalog]: SocketMessageHasCallback<TMessageCatalog, K> extends true
|
||||
? never
|
||||
: K;
|
||||
|
||||
Reference in New Issue
Block a user