diff --git a/apps/coordinator/src/index.ts b/apps/coordinator/src/index.ts index 8c2b66163..185b4c84c 100644 --- a/apps/coordinator/src/index.ts +++ b/apps/coordinator/src/index.ts @@ -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 { 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; - #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); - }, }, }); diff --git a/apps/docker-provider/src/index.ts b/apps/docker-provider/src/index.ts index 50095a5ba..ec05aa2bf 100644 --- a/apps/docker-provider/src/index.ts +++ b/apps/docker-provider/src/index.ts @@ -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(); diff --git a/apps/kubernetes-provider/src/index.ts b/apps/kubernetes-provider/src/index.ts index 6ef8750ad..c790497b4 100644 --- a/apps/kubernetes-provider/src/index.ts +++ b/apps/kubernetes-provider/src/index.ts @@ -389,6 +389,7 @@ class KubernetesTaskOperations implements TaskOperations { const provider = new ProviderShell({ tasks: new KubernetesTaskOperations(), + type: "kubernetes", }); provider.listen(); diff --git a/apps/webapp/app/v3/handleSocketIo.server.ts b/apps/webapp/app/v3/handleSocketIo.server.ts index d6a844a89..0697179dc 100644 --- a/apps/webapp/app/v3/handleSocketIo.server.ts +++ b/apps/webapp/app/v3/handleSocketIo.server.ts @@ -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); diff --git a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts index 50d26a1f6..69d848385 100644 --- a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts @@ -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); diff --git a/apps/webapp/app/v3/services/completeAttempt.server.ts b/apps/webapp/app/v3/services/completeAttempt.server.ts index 19bbbcea0..0c5b3c901 100644 --- a/apps/webapp/app/v3/services/completeAttempt.server.ts +++ b/apps/webapp/app/v3/services/completeAttempt.server.ts @@ -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], + }); + } } } diff --git a/apps/webapp/app/v3/services/createCheckpoint.server.ts b/apps/webapp/app/v3/services/createCheckpoint.server.ts index f36f66c79..acebca7dc 100644 --- a/apps/webapp/app/v3/services/createCheckpoint.server.ts +++ b/apps/webapp/app/v3/services/createCheckpoint.server.ts @@ -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; } diff --git a/apps/webapp/app/v3/services/resumeAttempt.server.ts b/apps/webapp/app/v3/services/resumeAttempt.server.ts new file mode 100644 index 000000000..1f4af4b64 --- /dev/null +++ b/apps/webapp/app/v3/services/resumeAttempt.server.ts @@ -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 + ): Promise { + 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; + } + } + } +} diff --git a/packages/cli-v3/src/commands/deploy.ts b/packages/cli-v3/src/commands/deploy.ts index 12178cb7e..be8a7408f 100644 --- a/packages/cli-v3/src/commands/deploy.ts +++ b/packages/cli-v3/src/commands/deploy.ts @@ -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: { diff --git a/packages/cli-v3/src/commands/dev.tsx b/packages/cli-v3/src/commands/dev.tsx index 766c928d9..446a985b9 100644 --- a/packages/cli-v3/src/commands/dev.tsx +++ b/packages/cli-v3/src/commands/dev.tsx @@ -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 diff --git a/packages/cli-v3/src/workers/common/worker-setup.ts b/packages/cli-v3/src/workers/dev/worker-setup.ts similarity index 100% rename from packages/cli-v3/src/workers/common/worker-setup.ts rename to packages/cli-v3/src/workers/dev/worker-setup.ts diff --git a/packages/cli-v3/src/workers/prod/backgroundWorker.ts b/packages/cli-v3/src/workers/prod/backgroundWorker.ts index 1cd5e513a..909e5e6cd 100644 --- a/packages/cli-v3/src/workers/prod/backgroundWorker.ts +++ b/packages/cli-v3/src/workers/prod/backgroundWorker.ts @@ -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 = 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 = 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 { 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; - 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", {}); } } diff --git a/packages/cli-v3/src/workers/prod/entry-point.ts b/packages/cli-v3/src/workers/prod/entry-point.ts index b3d844a30..cf850dc4f 100644 --- a/packages/cli-v3/src/workers/prod/entry-point.ts +++ b/packages/cli-v3/src/workers/prod/entry-point.ts @@ -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"); diff --git a/packages/cli-v3/src/workers/prod/worker-facade.ts b/packages/cli-v3/src/workers/prod/worker-facade.ts index 9d3ecee38..eae5674e6 100644 --- a/packages/cli-v3/src/workers/prod/worker-facade.ts +++ b/packages/cli-v3/src/workers/prod/worker-facade.ts @@ -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; - -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); } diff --git a/packages/cli-v3/src/workers/prod/worker-setup.ts b/packages/cli-v3/src/workers/prod/worker-setup.ts new file mode 100644 index 000000000..535670edc --- /dev/null +++ b/packages/cli-v3/src/workers/prod/worker-setup.ts @@ -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", + }, + }); +}); diff --git a/packages/cli-v3/tsup.workers.config.ts b/packages/cli-v3/tsup.workers.config.ts index ac0872d0e..f582f968d 100644 --- a/packages/cli-v3/tsup.workers.config.ts +++ b/packages/cli-v3/tsup.workers.config.ts @@ -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", ], diff --git a/packages/core-apps/src/provider.ts b/packages/core-apps/src/provider.ts index 32206100d..930f9b373 100644 --- a/packages/core-apps/src/provider.ts +++ b/packages/core-apps/src/provider.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) => {}, }, }); diff --git a/packages/core/src/v3/index.ts b/packages/core/src/v3/index.ts index e0e10d5c0..c54f01cf1 100644 --- a/packages/core/src/v3/index.ts +++ b/packages/core/src/v3/index.ts @@ -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"; diff --git a/packages/core/src/v3/runtime/prodRuntimeManager.ts b/packages/core/src/v3/runtime/prodRuntimeManager.ts index d4e751824..0f9889bb6 100644 --- a/packages/core/src/v3/runtime/prodRuntimeManager.ts +++ b/packages/core/src/v3/runtime/prodRuntimeManager.ts @@ -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 = new Map(); - constructor(private sender: ZodMessageSender) {} + 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 { - 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((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 { @@ -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, }); diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts index 26b4c7a20..ad6706037 100644 --- a/packages/core/src/v3/schemas/messages.ts +++ b/packages/core/src/v3/schemas/messages.ts @@ -187,6 +187,16 @@ export const TaskMetadataWithFilePath = TaskMetadata.extend({ export type TaskMetadataWithFilePath = z.infer; +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"), + }), + }, }; diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts index aef0c229c..efd5395f0 100644 --- a/packages/core/src/v3/schemas/schemas.ts +++ b/packages/core/src/v3/schemas/schemas.ts @@ -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, - }), - ]), }, }; diff --git a/packages/core/src/v3/zodIpc.ts b/packages/core/src/v3/zodIpc.ts new file mode 100644 index 000000000..cfc084643 --- /dev/null +++ b/packages/core/src/v3/zodIpc.ts @@ -0,0 +1,328 @@ +import { randomUUID } from "crypto"; +import { + GetSocketCallbackSchema, + GetSocketMessageSchema, + GetSocketMessagesWithCallback, + GetSocketMessagesWithoutCallback, + MessagesFromSocketCatalog, + SocketMessageHasCallback, + ZodSocketMessageCatalogSchema, +} from "./zodSocket"; +import { z } from "zod"; + +interface ZodIpcMessageSender { + send>( + type: K, + payload: z.input> + ): Promise; + + sendWithAck>( + type: K, + payload: z.input> + ): Promise>>; +} + +type ZodIpcMessageHandlers< + TListenCatalog extends ZodSocketMessageCatalogSchema, + TEmitCatalog extends ZodSocketMessageCatalogSchema, +> = Partial<{ + [K in keyof TListenCatalog]: ( + payload: z.infer>, + sender: ZodIpcMessageSender + ) => Promise< + SocketMessageHasCallback extends true + ? z.input> + : 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; + sender: ZodIpcMessageSender; +}; + +class ZodIpcMessageHandler< + TListenCatalog extends ZodSocketMessageCatalogSchema, + TEmitCatalog extends ZodSocketMessageCatalogSchema, +> { + #schema: TListenCatalog; + #handlers: ZodIpcMessageHandlers | undefined; + #sender: ZodIpcMessageSender; + + constructor(options: ZodIpcMessageHandlerOptions) { + 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 { + 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; + +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; +} + +export class ZodIpcConnection< + TListenCatalog extends ZodSocketMessageCatalogSchema, + TEmitCatalog extends ZodSocketMessageCatalogSchema, +> { + #sessionId?: string; + #messageCounter: number = 0; + + #handler: ZodIpcMessageHandler; + + #acks: Map< + number, + { + resolve: (value: unknown) => void; + reject: (reason?: any) => void; + timeout: NodeJS.Timeout; + } + > = new Map(); + + constructor(private opts: ZodIpcConnectionOptions) { + 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 { + 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>( + type: K, + payload: z.input> + ): Promise { + 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>( + type: K, + payload: z.input>, + timeoutInMs?: number + ): Promise>> { + 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, + }); + }); + } +} diff --git a/packages/core/src/v3/zodNamespace.ts b/packages/core/src/v3/zodNamespace.ts index ee74986dc..d6dd5976c 100644 --- a/packages/core/src/v3/zodNamespace.ts +++ b/packages/core/src/v3/zodNamespace.ts @@ -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) => { diff --git a/packages/core/src/v3/zodSocket.ts b/packages/core/src/v3/zodSocket.ts index 686917a39..68d489796 100644 --- a/packages/core/src/v3/zodSocket.ts +++ b/packages/core/src/v3/zodSocket.ts @@ -22,7 +22,7 @@ export type ZodMessageCatalogToSocketIoEvents>) => 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>; -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>; -type SocketMessageHasCallback< +export type SocketMessageHasCallback< TRPCCatalog extends ZodSocketMessageCatalogSchema, TMessageType extends keyof TRPCCatalog, > = GetSocketCallbackSchema extends never ? false : true; @@ -74,7 +74,7 @@ type MessageFromSocketSchema< payload: z.input>; }; -type MessagesFromSocketCatalog = { +export type MessagesFromSocketCatalog = { [K in keyof TMessageCatalog]: MessageFromSocketSchema; }[keyof TMessageCatalog]; @@ -173,13 +173,15 @@ export type ZodSocketMessageSenderOptions; }; -type GetSocketMessagesWithCallback = { +export type GetSocketMessagesWithCallback = { [K in keyof TMessageCatalog]: SocketMessageHasCallback extends true ? K : never; }[keyof TMessageCatalog]; -type GetSocketMessagesWithoutCallback = { +export type GetSocketMessagesWithoutCallback< + TMessageCatalog extends ZodSocketMessageCatalogSchema, +> = { [K in keyof TMessageCatalog]: SocketMessageHasCallback extends true ? never : K;