diff --git a/.changeset/mighty-eggs-grab.md b/.changeset/mighty-eggs-grab.md new file mode 100644 index 000000000..1d5cd94f2 --- /dev/null +++ b/.changeset/mighty-eggs-grab.md @@ -0,0 +1,30 @@ +--- +"@trigger.dev/core-apps": patch +"trigger.dev": patch +"@trigger.dev/core": patch +--- + +Tasks should now be much more robust and resilient to reconnects during crucial operations and other failure scenarios. + +Task runs now have to signal checkpointable state prior to ALL checkpoints. This ensures flushing always happens. + +All important socket.io RPCs will now be retried with backoff. Actions relying on checkpoints will be replayed if we haven't been checkpointed and restored as expected, e.g. after reconnect. + +Other changes: + +- Fix retry check in shared queue +- Fix env var sync spinner +- Heartbeat between retries +- Fix retry prep +- Fix prod worker no tasks detection +- Fail runs above `MAX_TASK_RUN_ATTEMPTS` +- Additional debug logs in all places +- Prevent crashes due to failed socket schema parsing +- Remove core-apps barrel +- Upgrade socket.io-client to fix an ACK memleak +- Additional index failure logs +- Prevent message loss during reconnect +- Prevent burst of heartbeats on reconnect +- Prevent crash on failed cleanup +- Handle at-least-once lazy execute message delivery +- Handle uncaught entry point exceptions diff --git a/.env.example b/.env.example index 26152641b..30f68970d 100644 --- a/.env.example +++ b/.env.example @@ -71,7 +71,7 @@ COORDINATOR_SECRET=coordinator-secret # generate the actual secret with `openssl # OBJECT_STORE_BASE_URL="https://{bucket}.{accountId}.r2.cloudflarestorage.com" # OBJECT_STORE_ACCESS_KEY_ID= # OBJECT_STORE_SECRET_ACCESS_KEY= -# RUNTIME_WAIT_THRESHOLD_IN_MS=10000 +# CHECKPOINT_THRESHOLD_IN_MS=10000 # These control the server-side internal telemetry # INTERNAL_OTEL_TRACE_EXPORTER_URL= diff --git a/apps/coordinator/package.json b/apps/coordinator/package.json index 4f7c56395..1c1584223 100644 --- a/apps/coordinator/package.json +++ b/apps/coordinator/package.json @@ -21,8 +21,7 @@ "execa": "^8.0.1", "nanoid": "^5.0.6", "prom-client": "^15.1.0", - "socket.io": "4.7.4", - "socket.io-client": "4.7.4" + "socket.io": "4.7.4" }, "devDependencies": { "@types/node": "^18", diff --git a/apps/coordinator/src/chaosMonkey.ts b/apps/coordinator/src/chaosMonkey.ts new file mode 100644 index 000000000..c92b07224 --- /dev/null +++ b/apps/coordinator/src/chaosMonkey.ts @@ -0,0 +1,95 @@ +import type { Execa$ } from "execa"; +import { setTimeout as timeout } from "node:timers/promises"; + +class ChaosMonkeyError extends Error { + constructor(message: string) { + super(message); + this.name = "ChaosMonkeyError"; + } +} + +export class ChaosMonkey { + private chaosEventRate = 0.2; + private delayInSeconds = 45; + + constructor(private enabled = false) { + if (this.enabled) { + console.log("🍌 Chaos monkey enabled"); + } + } + + static Error = ChaosMonkeyError; + + enable() { + this.enabled = true; + console.log("🍌 Chaos monkey enabled"); + } + + disable() { + this.enabled = false; + console.log("🍌 Chaos monkey disabled"); + } + + async call({ + $, + throwErrors = true, + addDelays = true, + }: { + $?: Execa$; + throwErrors?: boolean; + addDelays?: boolean; + } = {}) { + if (!this.enabled) { + return; + } + + const random = Math.random(); + + if (random > this.chaosEventRate) { + // Don't interfere with normal operation + return; + } + + const chaosEvents: Array<() => Promise> = []; + + if (addDelays) { + chaosEvents.push(async () => { + console.log("🍌 Chaos monkey: Add delay"); + + if ($) { + await $`sleep ${this.delayInSeconds}`; + } else { + await timeout(this.delayInSeconds * 1000); + } + }); + } + + if (throwErrors) { + chaosEvents.push(async () => { + console.log("🍌 Chaos monkey: Throw error"); + + if ($) { + await $`false`; + } else { + throw new ChaosMonkey.Error("🍌 Chaos monkey: Throw error"); + } + }); + } + + if (chaosEvents.length === 0) { + console.error("🍌 Chaos monkey: No events selected"); + return; + } + + const randomIndex = Math.floor(Math.random() * chaosEvents.length); + + const chaosEvent = chaosEvents[randomIndex]; + + if (!chaosEvent) { + console.error("🍌 Chaos monkey: No event found"); + return; + } + + await chaosEvent(); + } +} diff --git a/apps/coordinator/src/checkpointer.ts b/apps/coordinator/src/checkpointer.ts new file mode 100644 index 000000000..797764f9c --- /dev/null +++ b/apps/coordinator/src/checkpointer.ts @@ -0,0 +1,587 @@ +import { ExponentialBackoff } from "@trigger.dev/core-apps/backoff"; +import { isExecaChildProcess, testDockerCheckpoint } from "@trigger.dev/core-apps/checkpoints"; +import { SimpleLogger } from "@trigger.dev/core-apps/logger"; +import { $ } from "execa"; +import { nanoid } from "nanoid"; +import fs from "node:fs/promises"; +import { ChaosMonkey } from "./chaosMonkey"; + +type CheckpointerInitializeReturn = { + canCheckpoint: boolean; + willSimulate: boolean; +}; + +type CheckpointAndPushOptions = { + runId: string; + leaveRunning?: boolean; + projectRef: string; + deploymentVersion: string; + shouldHeartbeat?: boolean; +}; + +type CheckpointAndPushResult = + | { success: true; checkpoint: CheckpointData } + | { + success: false; + reason?: "CANCELED" | "DISABLED" | "ERROR" | "IN_PROGRESS" | "NO_SUPPORT" | "SKIP_RETRYING"; + }; + +type CheckpointData = { + location: string; + docker: boolean; +}; + +type CheckpointerOptions = { + dockerMode: boolean; + forceSimulate: boolean; + heartbeat: (runId: string) => void; + registryHost?: string; + registryNamespace?: string; + registryTlsVerify?: boolean; + disableCheckpointSupport?: boolean; + checkpointPath?: string; + simulateCheckpointFailure?: boolean; + simulateCheckpointFailureSeconds?: number; + simulatePushFailure?: boolean; + simulatePushFailureSeconds?: number; + chaosMonkey?: ChaosMonkey; +}; + +async function getFileSize(filePath: string): Promise { + try { + const stats = await fs.stat(filePath); + return stats.size; + } catch (error) { + console.error("Error getting file size:", error); + return -1; + } +} + +async function getParsedFileSize(filePath: string) { + const sizeInBytes = await getFileSize(filePath); + + let message = `Size in bytes: ${sizeInBytes}`; + + if (sizeInBytes > 1024 * 1024) { + const sizeInMB = (sizeInBytes / 1024 / 1024).toFixed(2); + message = `Size in MB (rounded): ${sizeInMB}`; + } else if (sizeInBytes > 1024) { + const sizeInKB = (sizeInBytes / 1024).toFixed(2); + message = `Size in KB (rounded): ${sizeInKB}`; + } + + return { + path: filePath, + sizeInBytes, + message, + }; +} + +export class Checkpointer { + #initialized = false; + #canCheckpoint = false; + #dockerMode: boolean; + + #logger = new SimpleLogger("[checkptr]"); + #abortControllers = new Map(); + #failedCheckpoints = new Map(); + #waitingForRetry = new Set(); + + private registryHost: string; + private registryNamespace: string; + private registryTlsVerify: boolean; + + private disableCheckpointSupport: boolean; + private checkpointPath: string; + + private simulateCheckpointFailure: boolean; + private simulateCheckpointFailureSeconds: number; + private simulatePushFailure: boolean; + private simulatePushFailureSeconds: number; + + private chaosMonkey: ChaosMonkey; + + constructor(private opts: CheckpointerOptions) { + this.#dockerMode = opts.dockerMode; + + this.registryHost = opts.registryHost ?? "localhost:5000"; + this.registryNamespace = opts.registryNamespace ?? "trigger"; + this.registryTlsVerify = opts.registryTlsVerify ?? true; + + this.disableCheckpointSupport = opts.disableCheckpointSupport ?? false; + this.checkpointPath = opts.checkpointPath ?? "/checkpoints"; + + this.simulateCheckpointFailure = opts.simulateCheckpointFailure ?? false; + this.simulateCheckpointFailureSeconds = opts.simulateCheckpointFailureSeconds ?? 300; + this.simulatePushFailure = opts.simulatePushFailure ?? false; + this.simulatePushFailureSeconds = opts.simulatePushFailureSeconds ?? 300; + + this.chaosMonkey = opts.chaosMonkey ?? new ChaosMonkey(!!process.env.CHAOS_MONKEY_ENABLED); + } + + async init(): Promise { + if (this.#initialized) { + return this.#getInitReturn(this.#canCheckpoint); + } + + this.#logger.log(`${this.#dockerMode ? "Docker" : "Kubernetes"} mode`); + + if (this.#dockerMode) { + const testCheckpoint = await testDockerCheckpoint(); + + if (testCheckpoint.ok) { + return this.#getInitReturn(true); + } + + this.#logger.error(testCheckpoint.message, testCheckpoint.error ?? ""); + return this.#getInitReturn(false); + } else { + try { + await $`buildah login --get-login ${this.registryHost}`; + } catch (error) { + this.#logger.error(`No checkpoint support: Not logged in to registry ${this.registryHost}`); + return this.#getInitReturn(false); + } + } + + return this.#getInitReturn(true); + } + + #getInitReturn(canCheckpoint: boolean): CheckpointerInitializeReturn { + this.#canCheckpoint = canCheckpoint; + + if (canCheckpoint) { + if (!this.#initialized) { + this.#logger.log("Full checkpoint support!"); + } + } + + this.#initialized = true; + + const willSimulate = this.#dockerMode && (!this.#canCheckpoint || this.opts.forceSimulate); + + if (willSimulate) { + this.#logger.log("Simulation mode enabled. Containers will be paused, not checkpointed.", { + forceSimulate: this.opts.forceSimulate, + }); + } + + return { + canCheckpoint, + willSimulate, + }; + } + + #getImageRef(projectRef: string, deploymentVersion: string, shortCode: string) { + return `${this.registryHost}/${this.registryNamespace}/${projectRef}:${deploymentVersion}.prod-${shortCode}`; + } + + #getExportLocation(projectRef: string, deploymentVersion: string, shortCode: string) { + const basename = `${projectRef}-${deploymentVersion}-${shortCode}`; + + if (this.#dockerMode) { + return basename; + } else { + return `${this.checkpointPath}/${basename}.tar`; + } + } + + async checkpointAndPush(opts: CheckpointAndPushOptions): Promise { + const start = performance.now(); + this.#logger.log(`checkpointAndPush() start`, { start, opts }); + + let interval: NodeJS.Timer | undefined; + + if (opts.shouldHeartbeat) { + interval = setInterval(() => { + this.#logger.log("Sending heartbeat", { runId: opts.runId }); + this.opts.heartbeat(opts.runId); + }, 20_000); + } + + try { + const result = await this.#checkpointAndPushWithBackoff(opts); + + const end = performance.now(); + this.#logger.log(`checkpointAndPush() end`, { + start, + end, + diff: end - start, + opts, + success: result.success, + }); + + if (!result.success) { + return; + } + + return result.checkpoint; + } finally { + if (opts.shouldHeartbeat) { + clearInterval(interval); + } + } + } + + isCheckpointing(runId: string) { + return this.#abortControllers.has(runId) || this.#waitingForRetry.has(runId); + } + + cancelCheckpoint(runId: string): boolean { + // If the last checkpoint failed, pretend we canceled it + // This ensures tasks don't wait for external resume messages to continue + if (this.#hasFailedCheckpoint(runId)) { + this.#clearFailedCheckpoint(runId); + return true; + } + + if (this.#waitingForRetry.has(runId)) { + this.#waitingForRetry.delete(runId); + return true; + } + + const controller = this.#abortControllers.get(runId); + + if (!controller) { + this.#logger.debug("Nothing to cancel", { runId }); + return false; + } + + controller.abort("cancelCheckpointing()"); + this.#abortControllers.delete(runId); + + return true; + } + + async #checkpointAndPushWithBackoff({ + runId, + leaveRunning = true, // This mirrors kubernetes behaviour more accurately + projectRef, + deploymentVersion, + }: CheckpointAndPushOptions): Promise { + this.#logger.log("Checkpointing with backoff", { + runId, + leaveRunning, + projectRef, + deploymentVersion, + }); + + const backoff = new ExponentialBackoff() + .type("EqualJitter") + .base(3) + .max(3 * 3600) + .maxElapsed(48 * 3600); + + for await (const { delay, retry } of backoff) { + try { + if (retry > 0) { + this.#logger.error("Retrying checkpoint", { + runId, + retry, + delay, + }); + + this.#waitingForRetry.add(runId); + await new Promise((resolve) => setTimeout(resolve, delay.milliseconds)); + + if (!this.#waitingForRetry.has(runId)) { + this.#logger.log("Checkpoint canceled while waiting for retry", { runId }); + return { success: false, reason: "CANCELED" }; + } else { + this.#waitingForRetry.delete(runId); + } + } + + const result = await this.#checkpointAndPush({ + runId, + leaveRunning, + projectRef, + deploymentVersion, + }); + + if (result.success) { + return result; + } + + if (result.reason === "CANCELED") { + this.#logger.log("Checkpoint canceled, won't retry", { runId }); + // Don't fail the checkpoint, as it was canceled + return result; + } + + if (result.reason === "IN_PROGRESS") { + this.#logger.log("Checkpoint already in progress, won't retry", { runId }); + this.#failCheckpoint(runId, result.reason); + return result; + } + + if (result.reason === "NO_SUPPORT") { + this.#logger.log("No checkpoint support, won't retry", { runId }); + this.#failCheckpoint(runId, result.reason); + return result; + } + + if (result.reason === "DISABLED") { + this.#logger.log("Checkpoint support disabled, won't retry", { runId }); + this.#failCheckpoint(runId, result.reason); + return result; + } + + if (result.reason === "SKIP_RETRYING") { + this.#logger.log("Skipping retrying", { runId }); + return result; + } + + continue; + } catch (error) { + this.#logger.error("Checkpoint error", { + retry, + runId, + delay, + error: error instanceof Error ? error.message : error, + }); + } + } + + this.#logger.error(`Checkpoint failed after exponential backoff`, { + runId, + leaveRunning, + projectRef, + deploymentVersion, + }); + this.#failCheckpoint(runId, "ERROR"); + + return { success: false, reason: "ERROR" }; + } + + async #checkpointAndPush({ + runId, + leaveRunning = true, // This mirrors kubernetes behaviour more accurately + projectRef, + deploymentVersion, + }: CheckpointAndPushOptions): Promise { + await this.init(); + + const options = { + runId, + leaveRunning, + projectRef, + deploymentVersion, + }; + + if (!this.#dockerMode && !this.#canCheckpoint) { + this.#logger.error("No checkpoint support. Simulation requires docker."); + return { success: false, reason: "NO_SUPPORT" }; + } + + if (this.isCheckpointing(runId)) { + this.#logger.error("Checkpoint procedure already in progress", { options }); + return { success: false, reason: "IN_PROGRESS" }; + } + + // This is a new checkpoint, clear any last failure for this run + this.#clearFailedCheckpoint(runId); + + if (this.disableCheckpointSupport) { + this.#logger.error("Checkpoint support disabled", { options }); + return { success: false, reason: "DISABLED" }; + } + + const controller = new AbortController(); + this.#abortControllers.set(runId, controller); + + const $$ = $({ signal: controller.signal }); + + const shortCode = nanoid(8); + const imageRef = this.#getImageRef(projectRef, deploymentVersion, shortCode); + const exportLocation = this.#getExportLocation(projectRef, deploymentVersion, shortCode); + + const cleanup = async () => { + if (this.#dockerMode) { + return; + } + + try { + await $`rm ${exportLocation}`; + this.#logger.log("Deleted checkpoint archive", { exportLocation }); + + await $`buildah rmi ${imageRef}`; + this.#logger.log("Deleted checkpoint image", { imageRef }); + } catch (error) { + this.#logger.error("Failure during checkpoint cleanup", { exportLocation, error }); + } + }; + + try { + await this.chaosMonkey.call({ $: $$ }); + + this.#logger.log("Checkpointing:", { options }); + + const containterName = this.#getRunContainerName(runId); + + // Create checkpoint (docker) + if (this.#dockerMode) { + try { + if (this.opts.forceSimulate || !this.#canCheckpoint) { + this.#logger.log("Simulating checkpoint"); + this.#logger.debug(await $$`docker pause ${containterName}`); + } else { + if (this.simulateCheckpointFailure) { + if (performance.now() < this.simulateCheckpointFailureSeconds * 1000) { + this.#logger.error("Simulating checkpoint failure", { options }); + throw new Error("SIMULATE_CHECKPOINT_FAILURE"); + } + } + + if (leaveRunning) { + this.#logger.debug( + await $$`docker checkpoint create --leave-running ${containterName} ${exportLocation}` + ); + } else { + this.#logger.debug( + await $$`docker checkpoint create ${containterName} ${exportLocation}` + ); + } + } + } catch (error) { + this.#logger.error("Failed while creating docker checkpoint", { exportLocation }); + throw error; + } + + this.#logger.log("checkpoint created:", { + runId, + location: exportLocation, + }); + + return { + success: true, + checkpoint: { + location: exportLocation, + docker: true, + }, + }; + } + + // Create checkpoint (CRI) + if (!this.#canCheckpoint) { + this.#logger.error("No checkpoint support in kubernetes mode."); + return { success: false, reason: "SKIP_RETRYING" }; + } + + const containerId = this.#logger.debug( + // @ts-expect-error + await $$`crictl ps` + .pipeStdout($$({ stdin: "pipe" })`grep ${containterName}`) + .pipeStdout($$({ stdin: "pipe" })`cut -f1 ${"-d "}`) + ); + + if (!containerId.stdout) { + this.#logger.error("could not find container id", { options, containterName }); + return { success: false, reason: "SKIP_RETRYING" }; + } + + const start = performance.now(); + + if (this.simulateCheckpointFailure) { + if (performance.now() < this.simulateCheckpointFailureSeconds * 1000) { + this.#logger.error("Simulating checkpoint failure", { options }); + throw new Error("SIMULATE_CHECKPOINT_FAILURE"); + } + } + + // Create checkpoint + this.#logger.debug(await $$`crictl checkpoint --export=${exportLocation} ${containerId}`); + const postCheckpoint = performance.now(); + + // Print checkpoint size + const size = await getParsedFileSize(exportLocation); + this.#logger.log("checkpoint archive created", { size, options }); + + // Create image from checkpoint + const container = this.#logger.debug(await $$`buildah from scratch`); + const postFrom = performance.now(); + + this.#logger.debug(await $$`buildah add ${container} ${exportLocation} /`); + const postAdd = performance.now(); + + this.#logger.debug( + await $$`buildah config --annotation=io.kubernetes.cri-o.annotations.checkpoint.name=counter ${container}` + ); + const postConfig = performance.now(); + + this.#logger.debug(await $$`buildah commit ${container} ${imageRef}`); + const postCommit = performance.now(); + + this.#logger.debug(await $$`buildah rm ${container}`); + const postRm = performance.now(); + + if (this.simulatePushFailure) { + if (performance.now() < this.simulatePushFailureSeconds * 1000) { + this.#logger.error("Simulating push failure", { options }); + throw new Error("SIMULATE_PUSH_FAILURE"); + } + } + + // Push checkpoint image + this.#logger.debug( + await $$`buildah push --tls-verify=${String(this.registryTlsVerify)} ${imageRef}` + ); + const postPush = performance.now(); + + const perf = { + "crictl checkpoint": postCheckpoint - start, + "buildah from": postFrom - postCheckpoint, + "buildah add": postAdd - postFrom, + "buildah config": postConfig - postAdd, + "buildah commit": postCommit - postConfig, + "buildah rm": postRm - postCommit, + "buildah push": postPush - postRm, + }; + + this.#logger.log("Checkpointed and pushed image to:", { location: imageRef, perf }); + + return { + success: true, + checkpoint: { + location: imageRef, + docker: false, + }, + }; + } catch (error) { + if (isExecaChildProcess(error)) { + if (error.isCanceled) { + this.#logger.error("Checkpoint canceled", { options, error }); + + return { success: false, reason: "CANCELED" }; + } + + this.#logger.error("Checkpoint command error", { options, error }); + + return { success: false, reason: "ERROR" }; + } + + this.#logger.error("Unhandled checkpoint error", { options, error }); + + return { success: false, reason: "ERROR" }; + } finally { + this.#abortControllers.delete(runId); + await cleanup(); + } + } + + #failCheckpoint(runId: string, error: unknown) { + this.#failedCheckpoints.set(runId, error); + } + + #clearFailedCheckpoint(runId: string) { + this.#failedCheckpoints.delete(runId); + } + + #hasFailedCheckpoint(runId: string) { + return this.#failedCheckpoints.has(runId); + } + + #getRunContainerName(suffix: string) { + return `task-run-${suffix}`; + } +} diff --git a/apps/coordinator/src/index.ts b/apps/coordinator/src/index.ts index b7cab90c5..5be829d5f 100644 --- a/apps/coordinator/src/index.ts +++ b/apps/coordinator/src/index.ts @@ -1,7 +1,4 @@ import { createServer } from "node:http"; -import fs from "node:fs/promises"; -import { $, type ExecaChildProcess } from "execa"; -import { nanoid } from "nanoid"; import { Server } from "socket.io"; import { CoordinatorToPlatformMessages, @@ -9,11 +6,14 @@ import { PlatformToCoordinatorMessages, ProdWorkerSocketData, ProdWorkerToCoordinatorMessages, + WaitReason, } from "@trigger.dev/core/v3"; import { ZodNamespace } from "@trigger.dev/core/v3/zodNamespace"; import { ZodSocketConnection } from "@trigger.dev/core/v3/zodSocket"; -import { HttpReply, getTextBody, SimpleLogger, testDockerCheckpoint } from "@trigger.dev/core-apps"; -import { ExponentialBackoff } from "./backoff"; +import { HttpReply, getTextBody } from "@trigger.dev/core-apps/http"; +import { SimpleLogger } from "@trigger.dev/core-apps/logger"; +import { ChaosMonkey } from "./chaosMonkey"; +import { Checkpointer } from "./checkpointer"; import { collectDefaultMetrics, register, Gauge } from "prom-client"; collectDefaultMetrics(); @@ -21,31 +21,26 @@ collectDefaultMetrics(); const HTTP_SERVER_PORT = Number(process.env.HTTP_SERVER_PORT || 8020); const NODE_NAME = process.env.NODE_NAME || "coordinator"; const DEFAULT_RETRY_DELAY_THRESHOLD_IN_MS = 30_000; -const CHAOS_MONKEY_ENABLED = !!process.env.CHAOS_MONKEY_ENABLED; -const FORCE_CHECKPOINT_SIMULATION = ["1", "true"].includes( - process.env.FORCE_CHECKPOINT_SIMULATION ?? "true" -); -const DISABLE_CHECKPOINT_SUPPORT = ["1", "true"].includes( - process.env.DISABLE_CHECKPOINT_SUPPORT ?? "false" -); -const SIMULATE_PUSH_FAILURE = ["1", "true"].includes(process.env.SIMULATE_PUSH_FAILURE ?? "false"); -const SIMULATE_PUSH_FAILURE_SECONDS = parseInt( - process.env.SIMULATE_PUSH_FAILURE_SECONDS ?? "300", - 10 -); -const SIMULATE_CHECKPOINT_FAILURE = ["1", "true"].includes( - process.env.SIMULATE_CHECKPOINT_FAILURE ?? "false" -); -const SIMULATE_CHECKPOINT_FAILURE_SECONDS = parseInt( - process.env.SIMULATE_CHECKPOINT_FAILURE_SECONDS ?? "300", - 10 -); +const boolFromEnv = (env: string, defaultValue: boolean): boolean => { + const value = process.env[env]; -const REGISTRY_HOST = process.env.REGISTRY_HOST || "localhost:5000"; -const REGISTRY_NAMESPACE = process.env.REGISTRY_NAMESPACE || "trigger"; -const CHECKPOINT_PATH = process.env.CHECKPOINT_PATH || "/checkpoints"; -const REGISTRY_TLS_VERIFY = process.env.REGISTRY_TLS_VERIFY === "false" ? "false" : "true"; + if (!value) { + return defaultValue; + } + + return ["1", "true"].includes(value); +}; + +const numFromEnv = (env: string, defaultValue: number): number => { + const value = process.env[env]; + + if (!value) { + return defaultValue; + } + + return parseInt(value, 10); +}; const PLATFORM_ENABLED = ["1", "true"].includes(process.env.PLATFORM_ENABLED ?? "true"); const PLATFORM_HOST = process.env.PLATFORM_HOST || "127.0.0.1"; @@ -54,547 +49,27 @@ const PLATFORM_SECRET = process.env.PLATFORM_SECRET || "coordinator-secret"; const SECURE_CONNECTION = ["1", "true"].includes(process.env.SECURE_CONNECTION ?? "false"); const logger = new SimpleLogger(`[${NODE_NAME}]`); - -if (CHAOS_MONKEY_ENABLED) { - logger.log("🍌 Chaos monkey enabled"); -} - -type CheckpointerInitializeReturn = { - canCheckpoint: boolean; - willSimulate: boolean; -}; - -type CheckpointAndPushOptions = { - runId: string; - leaveRunning?: boolean; - projectRef: string; - deploymentVersion: string; -}; - -type CheckpointAndPushResult = - | { success: true; checkpoint: CheckpointData } - | { - success: false; - reason?: "CANCELED" | "DISABLED" | "ERROR" | "IN_PROGRESS" | "NO_SUPPORT" | "SKIP_RETRYING"; - }; - -type CheckpointData = { - location: string; - docker: boolean; -}; - -function isExecaChildProcess(maybeExeca: unknown): maybeExeca is Awaited { - return typeof maybeExeca === "object" && maybeExeca !== null && "escapedCommand" in maybeExeca; -} - -async function getFileSize(filePath: string): Promise { - try { - const stats = await fs.stat(filePath); - return stats.size; - } catch (error) { - console.error("Error getting file size:", error); - return -1; - } -} - -async function getParsedFileSize(filePath: string) { - const sizeInBytes = await getFileSize(filePath); - - let message = `Size in bytes: ${sizeInBytes}`; - - if (sizeInBytes > 1024 * 1024) { - const sizeInMB = (sizeInBytes / 1024 / 1024).toFixed(2); - message = `Size in MB (rounded): ${sizeInMB}`; - } else if (sizeInBytes > 1024) { - const sizeInKB = (sizeInBytes / 1024).toFixed(2); - message = `Size in KB (rounded): ${sizeInKB}`; - } - - return { - path: filePath, - sizeInBytes, - message, - }; -} - -class Checkpointer { - #initialized = false; - #canCheckpoint = false; - #dockerMode = !process.env.KUBERNETES_PORT; - - #logger = new SimpleLogger("[checkptr]"); - #abortControllers = new Map(); - #failedCheckpoints = new Map(); - #waitingForRetry = new Set(); - - constructor(private opts = { forceSimulate: false }) {} - - async init(): Promise { - if (this.#initialized) { - return this.#getInitReturn(this.#canCheckpoint); - } - - this.#logger.log(`${this.#dockerMode ? "Docker" : "Kubernetes"} mode`); - - if (this.#dockerMode) { - const testCheckpoint = await testDockerCheckpoint(); - - if (testCheckpoint.ok) { - return this.#getInitReturn(true); - } - - this.#logger.error(testCheckpoint.message, testCheckpoint.error ?? ""); - return this.#getInitReturn(false); - } else { - try { - await $`buildah login --get-login ${REGISTRY_HOST}`; - } catch (error) { - this.#logger.error(`No checkpoint support: Not logged in to registry ${REGISTRY_HOST}`); - return this.#getInitReturn(false); - } - } - - return this.#getInitReturn(true); - } - - #getInitReturn(canCheckpoint: boolean): CheckpointerInitializeReturn { - this.#initialized = true; - this.#canCheckpoint = canCheckpoint; - - if (canCheckpoint) { - this.#logger.log("Full checkpoint support!"); - } - - const willSimulate = this.#dockerMode && (!this.#canCheckpoint || this.opts.forceSimulate); - - if (willSimulate) { - this.#logger.log("Simulation mode enabled. Containers will be paused, not checkpointed.", { - forceSimulate: this.opts.forceSimulate, - }); - } - - return { - canCheckpoint, - willSimulate, - }; - } - - #getImageRef(projectRef: string, deploymentVersion: string, shortCode: string) { - return `${REGISTRY_HOST}/${REGISTRY_NAMESPACE}/${projectRef}:${deploymentVersion}.prod-${shortCode}`; - } - - #getExportLocation(projectRef: string, deploymentVersion: string, shortCode: string) { - const basename = `${projectRef}-${deploymentVersion}-${shortCode}`; - - if (this.#dockerMode) { - return basename; - } else { - return `${CHECKPOINT_PATH}/${basename}.tar`; - } - } - - async checkpointAndPush(opts: CheckpointAndPushOptions): Promise { - const start = performance.now(); - logger.log(`checkpointAndPush() start`, { start, opts }); - - const result = await this.#checkpointAndPushWithBackoff(opts); - - const end = performance.now(); - logger.log(`checkpointAndPush() end`, { - start, - end, - diff: end - start, - opts, - success: result.success, - }); - - if (!result.success) { - return; - } - - return result.checkpoint; - } - - isCheckpointing(runId: string) { - return this.#abortControllers.has(runId) || this.#waitingForRetry.has(runId); - } - - cancelCheckpoint(runId: string): boolean { - // If the last checkpoint failed, pretend we canceled it - // This ensures tasks don't wait for external resume messages to continue - if (this.#hasFailedCheckpoint(runId)) { - this.#clearFailedCheckpoint(runId); - return true; - } - - if (this.#waitingForRetry.has(runId)) { - this.#waitingForRetry.delete(runId); - return true; - } - - const controller = this.#abortControllers.get(runId); - - if (!controller) { - logger.debug("Nothing to cancel", { runId }); - return false; - } - - controller.abort("cancelCheckpointing()"); - this.#abortControllers.delete(runId); - - return true; - } - - async #checkpointAndPushWithBackoff({ - runId, - leaveRunning = true, // This mirrors kubernetes behaviour more accurately - projectRef, - deploymentVersion, - }: CheckpointAndPushOptions): Promise { - this.#logger.log("Checkpointing with backoff", { - runId, - leaveRunning, - projectRef, - deploymentVersion, - }); - - const backoff = new ExponentialBackoff() - .type("EqualJitter") - .base(3) - .max(3 * 3600) - .maxElapsed(48 * 3600); - - for await (const { delay, retry } of backoff) { - try { - if (retry > 0) { - this.#logger.error("Retrying checkpoint", { - runId, - retry, - delay, - }); - - this.#waitingForRetry.add(runId); - await new Promise((resolve) => setTimeout(resolve, delay.milliseconds)); - - if (!this.#waitingForRetry.has(runId)) { - this.#logger.log("Checkpoint canceled while waiting for retry", { runId }); - return { success: false, reason: "CANCELED" }; - } else { - this.#waitingForRetry.delete(runId); - } - } - - const result = await this.#checkpointAndPush({ - runId, - leaveRunning, - projectRef, - deploymentVersion, - }); - - if (result.success) { - return result; - } - - if (result.reason === "CANCELED") { - this.#logger.log("Checkpoint canceled, won't retry", { runId }); - // Don't fail the checkpoint, as it was canceled - return result; - } - - if (result.reason === "IN_PROGRESS") { - this.#logger.log("Checkpoint already in progress, won't retry", { runId }); - this.#failCheckpoint(runId, result.reason); - return result; - } - - if (result.reason === "NO_SUPPORT") { - this.#logger.log("No checkpoint support, won't retry", { runId }); - this.#failCheckpoint(runId, result.reason); - return result; - } - - if (result.reason === "DISABLED") { - this.#logger.log("Checkpoint support disabled, won't retry", { runId }); - this.#failCheckpoint(runId, result.reason); - return result; - } - - if (result.reason === "SKIP_RETRYING") { - this.#logger.log("Skipping retrying", { runId }); - return result; - } - - continue; - } catch (error) { - this.#logger.error("Checkpoint error", { - retry, - runId, - delay, - error: error instanceof Error ? error.message : error, - }); - } - } - - this.#logger.error(`Checkpoint failed after exponential backoff`, { - runId, - leaveRunning, - projectRef, - deploymentVersion, - }); - this.#failCheckpoint(runId, "ERROR"); - - return { success: false, reason: "ERROR" }; - } - - async #checkpointAndPush({ - runId, - leaveRunning = true, // This mirrors kubernetes behaviour more accurately - projectRef, - deploymentVersion, - }: CheckpointAndPushOptions): Promise { - await this.init(); - - const options = { - runId, - leaveRunning, - projectRef, - deploymentVersion, - }; - - if (!this.#dockerMode && !this.#canCheckpoint) { - this.#logger.error("No checkpoint support. Simulation requires docker."); - return { success: false, reason: "NO_SUPPORT" }; - } - - if (this.#abortControllers.has(runId)) { - logger.error("Checkpoint procedure already in progress", { options }); - return { success: false, reason: "IN_PROGRESS" }; - } - - // This is a new checkpoint, clear any last failure for this run - this.#clearFailedCheckpoint(runId); - - if (DISABLE_CHECKPOINT_SUPPORT) { - this.#logger.error("Checkpoint support disabled", { options }); - return { success: false, reason: "DISABLED" }; - } - - const controller = new AbortController(); - this.#abortControllers.set(runId, controller); - - const $$ = $({ signal: controller.signal }); - - const shortCode = nanoid(8); - const imageRef = this.#getImageRef(projectRef, deploymentVersion, shortCode); - const exportLocation = this.#getExportLocation(projectRef, deploymentVersion, shortCode); - - const cleanup = async () => { - if (this.#dockerMode) { - return; - } - - try { - await $`rm ${exportLocation}`; - this.#logger.log("Deleted checkpoint archive", { exportLocation }); - - await $`buildah rmi ${imageRef}`; - this.#logger.log("Deleted checkpoint image", { imageRef }); - } catch (error) { - this.#logger.error("Failure during checkpoint cleanup", { exportLocation, error }); - } - }; - - try { - if (CHAOS_MONKEY_ENABLED) { - console.log("🍌 Chaos monkey wreaking havoc"); - - const random = Math.random(); - - if (random < 0.33) { - // Fake long checkpoint duration - await $$`sleep 300`; - } else if (random < 0.66) { - // Fake checkpoint error - await $$`false`; - } else { - // no-op - } - } - - this.#logger.log("Checkpointing:", { options }); - - const containterName = this.#getRunContainerName(runId); - - // Create checkpoint (docker) - if (this.#dockerMode) { - try { - if (this.opts.forceSimulate || !this.#canCheckpoint) { - this.#logger.log("Simulating checkpoint"); - this.#logger.debug(await $$`docker pause ${containterName}`); - } else { - if (SIMULATE_CHECKPOINT_FAILURE) { - if (performance.now() < SIMULATE_CHECKPOINT_FAILURE_SECONDS * 1000) { - this.#logger.error("Simulating checkpoint failure", { options }); - throw new Error("SIMULATE_CHECKPOINT_FAILURE"); - } - } - - if (leaveRunning) { - this.#logger.debug( - await $$`docker checkpoint create --leave-running ${containterName} ${exportLocation}` - ); - } else { - this.#logger.debug( - await $$`docker checkpoint create ${containterName} ${exportLocation}` - ); - } - } - } catch (error) { - this.#logger.error("Failed while creating docker checkpoint", { exportLocation }); - throw error; - } - - this.#logger.log("checkpoint created:", { - runId, - location: exportLocation, - }); - - return { - success: true, - checkpoint: { - location: exportLocation, - docker: true, - }, - }; - } - - // Create checkpoint (CRI) - if (!this.#canCheckpoint) { - this.#logger.error("No checkpoint support in kubernetes mode."); - return { success: false, reason: "SKIP_RETRYING" }; - } - - const containerId = this.#logger.debug( - // @ts-expect-error - await $$`crictl ps` - .pipeStdout($$({ stdin: "pipe" })`grep ${containterName}`) - .pipeStdout($$({ stdin: "pipe" })`cut -f1 ${"-d "}`) - ); - - if (!containerId.stdout) { - this.#logger.error("could not find container id", { options, containterName }); - return { success: false, reason: "SKIP_RETRYING" }; - } - - const start = performance.now(); - - if (SIMULATE_CHECKPOINT_FAILURE) { - if (performance.now() < SIMULATE_CHECKPOINT_FAILURE_SECONDS * 1000) { - this.#logger.error("Simulating checkpoint failure", { options }); - throw new Error("SIMULATE_CHECKPOINT_FAILURE"); - } - } - - // Create checkpoint - this.#logger.debug(await $$`crictl checkpoint --export=${exportLocation} ${containerId}`); - const postCheckpoint = performance.now(); - - // Print checkpoint size - const size = await getParsedFileSize(exportLocation); - this.#logger.log("checkpoint archive created", { size, options }); - - // Create image from checkpoint - const container = this.#logger.debug(await $$`buildah from scratch`); - const postFrom = performance.now(); - - this.#logger.debug(await $$`buildah add ${container} ${exportLocation} /`); - const postAdd = performance.now(); - - this.#logger.debug( - await $$`buildah config --annotation=io.kubernetes.cri-o.annotations.checkpoint.name=counter ${container}` - ); - const postConfig = performance.now(); - - this.#logger.debug(await $$`buildah commit ${container} ${imageRef}`); - const postCommit = performance.now(); - - this.#logger.debug(await $$`buildah rm ${container}`); - const postRm = performance.now(); - - if (SIMULATE_PUSH_FAILURE) { - if (performance.now() < SIMULATE_PUSH_FAILURE_SECONDS * 1000) { - this.#logger.error("Simulating push failure", { options }); - throw new Error("SIMULATE_PUSH_FAILURE"); - } - } - - // Push checkpoint image - this.#logger.debug(await $$`buildah push --tls-verify=${REGISTRY_TLS_VERIFY} ${imageRef}`); - const postPush = performance.now(); - - const perf = { - "crictl checkpoint": postCheckpoint - start, - "buildah from": postFrom - postCheckpoint, - "buildah add": postAdd - postFrom, - "buildah config": postConfig - postAdd, - "buildah commit": postCommit - postConfig, - "buildah rm": postRm - postCommit, - "buildah push": postPush - postRm, - }; - - this.#logger.log("Checkpointed and pushed image to:", { location: imageRef, perf }); - - return { - success: true, - checkpoint: { - location: imageRef, - docker: false, - }, - }; - } catch (error) { - if (isExecaChildProcess(error)) { - if (error.isCanceled) { - this.#logger.error("Checkpoint canceled", { options, error }); - - return { success: false, reason: "CANCELED" }; - } - - this.#logger.error("Checkpoint command error", { options, error }); - - return { success: false, reason: "ERROR" }; - } - - this.#logger.error("Unhandled checkpoint error", { options, error }); - - return { success: false, reason: "ERROR" }; - } finally { - this.#abortControllers.delete(runId); - await cleanup(); - } - } - - #failCheckpoint(runId: string, error: unknown) { - this.#failedCheckpoints.set(runId, error); - } - - #clearFailedCheckpoint(runId: string) { - this.#failedCheckpoints.delete(runId); - } - - #hasFailedCheckpoint(runId: string) { - return this.#failedCheckpoints.has(runId); - } - - #getRunContainerName(suffix: string) { - return `task-run-${suffix}`; - } -} +const chaosMonkey = new ChaosMonkey(!!process.env.CHAOS_MONKEY_ENABLED); class TaskCoordinator { #httpServer: ReturnType; - #checkpointer = new Checkpointer({ forceSimulate: FORCE_CHECKPOINT_SIMULATION }); + #checkpointer = new Checkpointer({ + dockerMode: !process.env.KUBERNETES_PORT, + forceSimulate: boolFromEnv("FORCE_CHECKPOINT_SIMULATION", false), + heartbeat: this.#sendRunHeartbeat.bind(this), + registryHost: process.env.REGISTRY_HOST, + registryNamespace: process.env.REGISTRY_NAMESPACE, + checkpointPath: process.env.CHECKPOINT_PATH, + registryTlsVerify: boolFromEnv("REGISTRY_TLS_VERIFY", true), + disableCheckpointSupport: boolFromEnv("DISABLE_CHECKPOINT_SUPPORT", false), + simulatePushFailure: boolFromEnv("SIMULATE_PUSH_FAILURE", false), + simulatePushFailureSeconds: numFromEnv("SIMULATE_PUSH_FAILURE_SECONDS", 300), + simulateCheckpointFailure: boolFromEnv("SIMULATE_CHECKPOINT_FAILURE", false), + simulateCheckpointFailureSeconds: numFromEnv("SIMULATE_CHECKPOINT_FAILURE_SECONDS", 300), + chaosMonkey, + }); - #prodWorkerNamespace: ZodNamespace< + #prodWorkerNamespace?: ZodNamespace< typeof ProdWorkerToCoordinatorMessages, typeof CoordinatorToProdWorkerMessages, typeof ProdWorkerSocketData @@ -609,7 +84,7 @@ class TaskCoordinator { { resolve: (value: void) => void; reject: (err?: any) => void } >(); - #delayThresholdInMs: number; + #delayThresholdInMs: number = DEFAULT_RETRY_DELAY_THRESHOLD_IN_MS; constructor( private port: number, @@ -617,59 +92,51 @@ class TaskCoordinator { ) { this.#httpServer = this.#createHttpServer(); this.#checkpointer.init(); - this.#delayThresholdInMs = this.#getDelayThreshold(); - - if (process.env.DELAY_THRESHOLD_IN_MS) { - this.#delayThresholdInMs = this.#getDelayThreshold(); - } - - const io = new Server(this.#httpServer); - this.#prodWorkerNamespace = this.#createProdWorkerNamespace(io); - this.#platformSocket = this.#createPlatformSocket(); const connectedTasksTotal = new Gauge({ name: "daemon_connected_tasks_total", // don't change this without updating dashboard config help: "The number of tasks currently connected.", collect: () => { - connectedTasksTotal.set(this.#prodWorkerNamespace.namespace.sockets.size); + connectedTasksTotal.set(this.#prodWorkerNamespace?.namespace.sockets.size ?? 0); }, }); register.registerMetric(connectedTasksTotal); } - #getDelayThreshold() { - if (!process.env.RETRY_DELAY_THRESHOLD_IN_MS) { - return DEFAULT_RETRY_DELAY_THRESHOLD_IN_MS; + #returnValidatedExtraHeaders(headers: Record) { + for (const [key, value] of Object.entries(headers)) { + if (value === undefined) { + throw new Error(`Extra header is undefined: ${key}`); + } } - const threshold = parseInt(process.env.RETRY_DELAY_THRESHOLD_IN_MS); - - if (isNaN(threshold)) { - logger.log( - "RETRY_DELAY_THRESHOLD_IN_MS parses as NaN, must supply integer. Will use default instead.", - { - RETRY_DELAY_THRESHOLD_IN_MS: process.env.RETRY_DELAY_THRESHOLD_IN_MS, - DEFAULT_DELAY_THRESHOLD_IN_MS: DEFAULT_RETRY_DELAY_THRESHOLD_IN_MS, - } - ); - return DEFAULT_RETRY_DELAY_THRESHOLD_IN_MS; - } - - return threshold; + return headers; } + // MARK: SOCKET: PLATFORM #createPlatformSocket() { if (!PLATFORM_ENABLED) { console.log("INFO: platform connection disabled"); return; } + const extraHeaders = this.#returnValidatedExtraHeaders({ + "x-supports-dynamic-config": "yes", + }); + + const host = PLATFORM_HOST; + const port = Number(PLATFORM_WS_PORT); + + logger.log(`connecting to platform: ${host}:${port}`); + logger.debug(`connecting with extra headers`, { extraHeaders }); + const platformConnection = new ZodSocketConnection({ namespace: "coordinator", - host: PLATFORM_HOST, - port: Number(PLATFORM_WS_PORT), + host, + port, secure: SECURE_CONNECTION, + extraHeaders, clientMessages: CoordinatorToPlatformMessages, serverMessages: PlatformToCoordinatorMessages, authToken: PLATFORM_SECRET, @@ -684,6 +151,8 @@ class TaskCoordinator { return; } + await chaosMonkey.call(); + // In case the task resumed faster than we could checkpoint this.#cancelCheckpoint(message.runId); @@ -699,6 +168,8 @@ class TaskCoordinator { return; } + await chaosMonkey.call(); + taskSocket.emit("RESUME_AFTER_DURATION", message); }, REQUEST_ATTEMPT_CANCELLATION: async (message) => { @@ -747,8 +218,19 @@ class TaskCoordinator { return; } + await chaosMonkey.call(); + taskSocket.emit("READY_FOR_RETRY", message); }, + DYNAMIC_CONFIG: async (message) => { + this.#delayThresholdInMs = message.checkpointThresholdInMs; + + // The first time we receive a dynamic config, the worker namespace will be created + if (!this.#prodWorkerNamespace) { + const io = new Server(this.#httpServer); + this.#prodWorkerNamespace = this.#createProdWorkerNamespace(io); + } + }, }, }); @@ -756,7 +238,7 @@ class TaskCoordinator { } async #getRunSocket(runId: string) { - const sockets = await this.#prodWorkerNamespace.fetchSockets(); + const sockets = (await this.#prodWorkerNamespace?.fetchSockets()) ?? []; for (const socket of sockets) { if (socket.data.runId === runId) { @@ -766,7 +248,7 @@ class TaskCoordinator { } async #getAttemptSocket(attemptFriendlyId: string) { - const sockets = await this.#prodWorkerNamespace.fetchSockets(); + const sockets = (await this.#prodWorkerNamespace?.fetchSockets()) ?? []; for (const socket of sockets) { if (socket.data.attemptFriendlyId === attemptFriendlyId) { @@ -775,6 +257,7 @@ class TaskCoordinator { } } + // MARK: SOCKET: WORKERS #createProdWorkerNamespace(io: Server) { const provider = new ZodNamespace({ io, @@ -841,9 +324,19 @@ class TaskCoordinator { return this.#checkpointableTasks.has(socket.data.runId); }; - const readyToCheckpoint = async (): Promise< - { success: true } | { success: false; reason?: string } + const readyToCheckpoint = async ( + reason: WaitReason | "RETRY" + ): Promise< + | { + success: true; + } + | { + success: false; + reason?: string; + } > => { + logger.log("readyToCheckpoint", { runId: socket.data.runId, reason }); + if (checkpointInProgress()) { return { success: false, @@ -851,10 +344,11 @@ class TaskCoordinator { }; } + let timeout: NodeJS.Timeout | undefined = undefined; + const isCheckpointable = new Promise((resolve, reject) => { // We set a reasonable timeout to prevent waiting forever - // TODO: We may also want to cancel the task as it's unlikely to recover - setTimeout(() => reject("timeout"), 10_000); + timeout = setTimeout(() => reject("timeout"), 20_000); this.#checkpointableTasks.set(socket.data.runId, { resolve, reject }); }); @@ -869,30 +363,36 @@ class TaskCoordinator { } catch (error) { logger.error("Error while waiting for checkpointable state", { error }); + await crashRun({ + name: "ReadyForCheckpointError", + message: `Failed to become checkpointable for ${reason}`, + }); + return { success: false, reason: typeof error === "string" ? error : "unknown", }; + } finally { + clearTimeout(timeout); } }; + const updateAttemptFriendlyId = (attemptFriendlyId: string) => { + socket.data.attemptFriendlyId = attemptFriendlyId; + }; + this.#platformSocket?.send("LOG", { metadata: socket.data, text: "connected", }); - socket.on("LOG", (message, callback) => { - logger.log("[LOG]", message.text); + socket.on("TEST", (message, callback) => { + logger.log("[TEST]", { runId: socket.data.runId, message }); callback(); - - this.#platformSocket?.send("LOG", { - version: "v1", - metadata: socket.data, - text: message.text, - }); }); + // Deprecated: Only workers without support for lazy attempts use this socket.on("READY_FOR_EXECUTION", async (message) => { logger.log("[READY_FOR_EXECUTION]", message); @@ -929,7 +429,7 @@ class TaskCoordinator { executionPayload: executionAck.payload, }); - socket.data.attemptFriendlyId = executionAck.payload.execution.attempt.id; + updateAttemptFriendlyId(executionAck.payload.execution.attempt.id); } catch (error) { logger.error("Error", { error }); @@ -943,6 +443,7 @@ class TaskCoordinator { } }); + // MARK: LAZY ATTEMPT socket.on("READY_FOR_LAZY_ATTEMPT", async (message) => { logger.log("[READY_FOR_LAZY_ATTEMPT]", message); @@ -974,11 +475,18 @@ class TaskCoordinator { return; } + await chaosMonkey.call(); + socket.emit("EXECUTE_TASK_RUN_LAZY_ATTEMPT", { version: "v1", lazyPayload: lazyAttempt.lazyPayload, }); } catch (error) { + if (error instanceof ChaosMonkey.Error) { + logger.error("ChaosMonkey error, won't crash run", { runId: socket.data.runId }); + return; + } + logger.error("Error", { error }); await crashRun({ @@ -991,16 +499,24 @@ class TaskCoordinator { } }); + // MARK: RESUME READY socket.on("READY_FOR_RESUME", async (message) => { logger.log("[READY_FOR_RESUME]", message); - socket.data.attemptFriendlyId = message.attemptFriendlyId; + updateAttemptFriendlyId(message.attemptFriendlyId); + this.#platformSocket?.send("READY_FOR_RESUME", message); }); + // MARK: RUN COMPLETED socket.on("TASK_RUN_COMPLETED", async ({ completion, execution }, callback) => { logger.log("completed task", { completionId: completion.id }); + // Cancel all in-progress checkpoints (if any) + this.#cancelCheckpoint(socket.data.runId); + + await chaosMonkey.call({ throwErrors: false }); + const completeWithoutCheckpoint = (shouldExit: boolean) => { this.#platformSocket?.send("TASK_RUN_COMPLETED", { version: "v1", @@ -1045,13 +561,14 @@ class TaskCoordinator { // The worker will then put itself in a checkpointable state callback({ willCheckpointAndRestore: true, shouldExit: false }); - const ready = await readyToCheckpoint(); + const ready = await readyToCheckpoint("RETRY"); if (!ready.success) { logger.error("Failed to become checkpointable", { runId: socket.data.runId, reason: ready.reason, }); + return; } @@ -1059,6 +576,7 @@ class TaskCoordinator { runId: socket.data.runId, projectRef: socket.data.projectRef, deploymentVersion: socket.data.deploymentVersion, + shouldHeartbeat: true, }); if (!checkpoint) { @@ -1081,8 +599,12 @@ class TaskCoordinator { } }); + // MARK: TASK FAILED socket.on("TASK_RUN_FAILED_TO_RUN", async ({ completion }) => { - logger.log("completed task", { completionId: completion.id }); + logger.log("task failed to run", { completionId: completion.id }); + + // Cancel all in-progress checkpoints (if any) + this.#cancelCheckpoint(socket.data.runId); this.#platformSocket?.send("TASK_RUN_FAILED_TO_RUN", { version: "v1", @@ -1094,6 +616,7 @@ class TaskCoordinator { }); }); + // MARK: CHECKPOINT socket.on("READY_FOR_CHECKPOINT", async (message) => { logger.log("[READY_FOR_CHECKPOINT]", message); @@ -1107,6 +630,7 @@ class TaskCoordinator { checkpointable.resolve(); }); + // MARK: CXX CHECKPOINT socket.on("CANCEL_CHECKPOINT", async (message, callback) => { logger.log("[CANCEL_CHECKPOINT]", message); @@ -1121,9 +645,12 @@ class TaskCoordinator { callback({ version: "v2", checkpointCanceled }); }); + // MARK: DURATION WAIT socket.on("WAIT_FOR_DURATION", async (message, callback) => { logger.log("[WAIT_FOR_DURATION]", message); + await chaosMonkey.call({ throwErrors: false }); + if (checkpointInProgress()) { logger.error("Checkpoint already in progress", { runId: socket.data.runId }); callback({ willCheckpointAndRestore: false }); @@ -1140,7 +667,7 @@ class TaskCoordinator { return; } - const ready = await readyToCheckpoint(); + const ready = await readyToCheckpoint("WAIT_FOR_DURATION"); if (!ready.success) { logger.error("Failed to become checkpointable", { @@ -1186,9 +713,18 @@ class TaskCoordinator { } }); + // MARK: TASK WAIT socket.on("WAIT_FOR_TASK", async (message, callback) => { logger.log("[WAIT_FOR_TASK]", message); + await chaosMonkey.call({ throwErrors: false }); + + if (checkpointInProgress()) { + logger.error("Checkpoint already in progress", { runId: socket.data.runId }); + callback({ willCheckpointAndRestore: false }); + return; + } + const { canCheckpoint, willSimulate } = await this.#checkpointer.init(); const willCheckpointAndRestore = canCheckpoint || willSimulate; @@ -1199,6 +735,19 @@ class TaskCoordinator { return; } + // Workers with v1 schemas don't signal when they're ready to checkpoint for dependency waits + if (message.version === "v2") { + const ready = await readyToCheckpoint("WAIT_FOR_TASK"); + + if (!ready.success) { + logger.error("Failed to become checkpointable", { + runId: socket.data.runId, + reason: ready.reason, + }); + return; + } + } + const checkpoint = await this.#checkpointer.checkpointAndPush({ runId: socket.data.runId, projectRef: socket.data.projectRef, @@ -1233,9 +782,18 @@ class TaskCoordinator { } }); + // MARK: BATCH WAIT socket.on("WAIT_FOR_BATCH", async (message, callback) => { logger.log("[WAIT_FOR_BATCH]", message); + await chaosMonkey.call({ throwErrors: false }); + + if (checkpointInProgress()) { + logger.error("Checkpoint already in progress", { runId: socket.data.runId }); + callback({ willCheckpointAndRestore: false }); + return; + } + const { canCheckpoint, willSimulate } = await this.#checkpointer.init(); const willCheckpointAndRestore = canCheckpoint || willSimulate; @@ -1246,6 +804,19 @@ class TaskCoordinator { return; } + // Workers with v1 schemas don't signal when they're ready to checkpoint for dependency waits + if (message.version === "v2") { + const ready = await readyToCheckpoint("WAIT_FOR_BATCH"); + + if (!ready.success) { + logger.error("Failed to become checkpointable", { + runId: socket.data.runId, + reason: ready.reason, + }); + return; + } + } + const checkpoint = await this.#checkpointer.checkpointAndPush({ runId: socket.data.runId, projectRef: socket.data.projectRef, @@ -1281,6 +852,7 @@ class TaskCoordinator { } }); + // MARK: INDEX socket.on("INDEX_TASKS", async (message, callback) => { logger.log("[INDEX_TASKS]", message); @@ -1304,6 +876,7 @@ class TaskCoordinator { callback({ success: !!workerAck?.success }); }); + // MARK: INDEX FAILED socket.on("INDEXING_FAILED", async (message) => { logger.log("[INDEXING_FAILED]", message); @@ -1314,9 +887,12 @@ class TaskCoordinator { }); }); + // MARK: CREATE ATTEMPT socket.on("CREATE_TASK_RUN_ATTEMPT", async (message, callback) => { logger.log("[CREATE_TASK_RUN_ATTEMPT]", message); + await chaosMonkey.call({ throwErrors: false }); + const createAttempt = await this.#platformSocket?.sendWithAck("CREATE_TASK_RUN_ATTEMPT", { runId: message.runId, envId: socket.data.envId, @@ -1324,11 +900,11 @@ class TaskCoordinator { if (!createAttempt?.success) { logger.debug("no ack while creating attempt", message); - callback({ success: false }); + callback({ success: false, reason: createAttempt?.reason }); return; } - socket.data.attemptFriendlyId = createAttempt.executionPayload.execution.attempt.id; + updateAttemptFriendlyId(createAttempt.executionPayload.execution.attempt.id); callback({ success: true, @@ -1341,6 +917,14 @@ class TaskCoordinator { await crashRun(message.error); }); + + socket.on("SET_STATE", async (message) => { + logger.log("[SET_STATE]", message); + + if (message.attemptFriendlyId) { + updateAttemptFriendlyId(message.attemptFriendlyId); + } + }); }, onDisconnect: async (socket, handler, sender, logger) => { this.#platformSocket?.send("LOG", { @@ -1353,7 +937,7 @@ class TaskCoordinator { this.#platformSocket?.send("TASK_HEARTBEAT", message); }, TASK_RUN_HEARTBEAT: async (message) => { - this.#platformSocket?.send("TASK_RUN_HEARTBEAT", message); + this.#sendRunHeartbeat(message.runId); }, }, }); @@ -1361,6 +945,13 @@ class TaskCoordinator { return provider; } + #sendRunHeartbeat(runId: string) { + this.#platformSocket?.send("TASK_RUN_HEARTBEAT", { + version: "v1", + runId, + }); + } + #cancelCheckpoint(runId: string): boolean { const checkpointWait = this.#checkpointableTasks.get(runId); @@ -1377,6 +968,7 @@ class TaskCoordinator { return checkpointCanceled; } + // MARK: HTTP SERVER #createHttpServer() { const httpServer = createServer(async (req, res) => { logger.log(`[${req.method}]`, req.url); diff --git a/apps/docker-provider/package.json b/apps/docker-provider/package.json index 59a1d522b..0c9e4e9a4 100644 --- a/apps/docker-provider/package.json +++ b/apps/docker-provider/package.json @@ -18,8 +18,7 @@ "dependencies": { "@trigger.dev/core": "workspace:*", "@trigger.dev/core-apps": "workspace:*", - "execa": "^8.0.1", - "socket.io-client": "^4.7.4" + "execa": "^8.0.1" }, "devDependencies": { "@types/node": "^18.19.8", diff --git a/apps/docker-provider/src/index.ts b/apps/docker-provider/src/index.ts index a37c54b46..c407e4786 100644 --- a/apps/docker-provider/src/index.ts +++ b/apps/docker-provider/src/index.ts @@ -1,14 +1,13 @@ import { $, type ExecaChildProcess, execa } from "execa"; import { - SimpleLogger, - TaskOperations, ProviderShell, - TaskOperationsRestoreOptions, + TaskOperations, TaskOperationsCreateOptions, TaskOperationsIndexOptions, - isExecaChildProcess, - testDockerCheckpoint, -} from "@trigger.dev/core-apps"; + TaskOperationsRestoreOptions, +} from "@trigger.dev/core-apps/provider"; +import { SimpleLogger } from "@trigger.dev/core-apps/logger"; +import { isExecaChildProcess, testDockerCheckpoint } from "@trigger.dev/core-apps/checkpoints"; import { setTimeout } from "node:timers/promises"; import { PostStartCauses, PreStopCauses } from "@trigger.dev/core/v3"; @@ -54,13 +53,16 @@ class DockerTaskOperations implements TaskOperations { } #getInitReturn(canCheckpoint: boolean): TaskOperationsInitReturn { - this.#initialized = true; this.#canCheckpoint = canCheckpoint; if (canCheckpoint) { - logger.log("Full checkpoint support!"); + if (!this.#initialized) { + logger.log("Full checkpoint support!"); + } } + this.#initialized = true; + const willSimulate = !canCheckpoint || this.opts.forceSimulate; if (willSimulate) { diff --git a/apps/kubernetes-provider/package.json b/apps/kubernetes-provider/package.json index 0e58bf309..a34123bca 100644 --- a/apps/kubernetes-provider/package.json +++ b/apps/kubernetes-provider/package.json @@ -19,8 +19,7 @@ "@kubernetes/client-node": "^0.20.0", "@trigger.dev/core": "workspace:*", "@trigger.dev/core-apps": "workspace:*", - "p-queue": "^8.0.1", - "socket.io-client": "^4.7.4" + "p-queue": "^8.0.1" }, "devDependencies": { "dotenv": "^16.4.2", diff --git a/apps/kubernetes-provider/src/index.ts b/apps/kubernetes-provider/src/index.ts index 9859c6697..60c27d1d7 100644 --- a/apps/kubernetes-provider/src/index.ts +++ b/apps/kubernetes-provider/src/index.ts @@ -1,12 +1,12 @@ import * as k8s from "@kubernetes/client-node"; import { ProviderShell, - SimpleLogger, TaskOperations, TaskOperationsCreateOptions, TaskOperationsIndexOptions, TaskOperationsRestoreOptions, -} from "@trigger.dev/core-apps"; +} from "@trigger.dev/core-apps/provider"; +import { SimpleLogger } from "@trigger.dev/core-apps/logger"; import { MachinePreset, PostStartCauses, diff --git a/apps/kubernetes-provider/src/podCleaner.ts b/apps/kubernetes-provider/src/podCleaner.ts index 29955bab3..04e3bc065 100644 --- a/apps/kubernetes-provider/src/podCleaner.ts +++ b/apps/kubernetes-provider/src/podCleaner.ts @@ -1,5 +1,5 @@ import * as k8s from "@kubernetes/client-node"; -import { SimpleLogger } from "@trigger.dev/core-apps"; +import { SimpleLogger } from "@trigger.dev/core-apps/logger"; type PodCleanerOptions = { runtimeEnv: "local" | "kubernetes"; diff --git a/apps/kubernetes-provider/src/taskMonitor.ts b/apps/kubernetes-provider/src/taskMonitor.ts index 1554fd46c..61fd7a63e 100644 --- a/apps/kubernetes-provider/src/taskMonitor.ts +++ b/apps/kubernetes-provider/src/taskMonitor.ts @@ -1,5 +1,5 @@ import * as k8s from "@kubernetes/client-node"; -import { SimpleLogger } from "@trigger.dev/core-apps"; +import { SimpleLogger } from "@trigger.dev/core-apps/logger"; import { EXIT_CODE_ALREADY_HANDLED, EXIT_CODE_CHILD_NONZERO } from "@trigger.dev/core-apps/process"; import { setTimeout } from "timers/promises"; import PQueue from "p-queue"; @@ -140,6 +140,9 @@ export class TaskMonitor { const exitCode = containerState.exitCode ?? -1; if (exitCode === EXIT_CODE_ALREADY_HANDLED) { + this.#logger.debug("Ignoring pod failure, already handled by worker", { + podName, + }); return; } diff --git a/apps/proxy/src/events/queueEvent.ts b/apps/proxy/src/events/queueEvent.ts index d3b2dcce5..29283e09d 100644 --- a/apps/proxy/src/events/queueEvent.ts +++ b/apps/proxy/src/events/queueEvent.ts @@ -24,6 +24,7 @@ export async function queueEvent(request: Request, env: Env): Promise const anyBody = await request.json(); const body = SendEventBodySchema.safeParse(anyBody); if (!body.success) { + fromZodError(body.error); return json( { error: generateErrorMessage(body.error.issues) }, { diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts index 9552aac4a..48614f71a 100644 --- a/apps/webapp/app/env.server.ts +++ b/apps/webapp/app/env.server.ts @@ -151,7 +151,7 @@ const EnvironmentSchema = z.object({ PROD_OTEL_LOG_EXPORT_TIMEOUT_MILLIS: z.string().default("30000"), PROD_OTEL_LOG_MAX_QUEUE_SIZE: z.string().default("512"), - RUNTIME_WAIT_THRESHOLD_IN_MS: z.coerce.number().int().default(30000), + CHECKPOINT_THRESHOLD_IN_MS: z.coerce.number().int().default(30000), // Internal OTEL environment variables INTERNAL_OTEL_TRACE_EXPORTER_URL: z.string().optional(), diff --git a/apps/webapp/app/v3/environmentVariables/environmentVariablesRepository.server.ts b/apps/webapp/app/v3/environmentVariables/environmentVariablesRepository.server.ts index e30c8560e..91981d585 100644 --- a/apps/webapp/app/v3/environmentVariables/environmentVariablesRepository.server.ts +++ b/apps/webapp/app/v3/environmentVariables/environmentVariablesRepository.server.ts @@ -740,7 +740,7 @@ async function resolveBuiltInProdVariables(runtimeEnvironment: RuntimeEnvironmen }, { key: "TRIGGER_RUNTIME_WAIT_THRESHOLD_IN_MS", - value: String(env.RUNTIME_WAIT_THRESHOLD_IN_MS), + value: String(env.CHECKPOINT_THRESHOLD_IN_MS), }, { key: "TRIGGER_ORG_ID", diff --git a/apps/webapp/app/v3/failedTaskRun.server.ts b/apps/webapp/app/v3/failedTaskRun.server.ts index 79594e73c..49711ac4f 100644 --- a/apps/webapp/app/v3/failedTaskRun.server.ts +++ b/apps/webapp/app/v3/failedTaskRun.server.ts @@ -9,14 +9,19 @@ import { BaseService } from "./services/baseService.server"; const FAILABLE_TASK_RUN_STATUSES: TaskRunStatus[] = ["EXECUTING", "PENDING", "WAITING_FOR_DEPLOY"]; export class FailedTaskRunService extends BaseService { - public async call(runFriendlyId: string, completion: TaskRunFailedExecutionResult) { + public async call(anyRunId: string, completion: TaskRunFailedExecutionResult) { + const isFriendlyId = anyRunId.startsWith("run_"); + const taskRun = await this._prisma.taskRun.findUnique({ - where: { friendlyId: runFriendlyId }, + where: { + friendlyId: isFriendlyId ? anyRunId : undefined, + id: !isFriendlyId ? anyRunId : undefined, + }, }); if (!taskRun) { logger.error("[FailedTaskRunService] Task run not found", { - runFriendlyId, + anyRunId, completion, }); diff --git a/apps/webapp/app/v3/handleSocketIo.server.ts b/apps/webapp/app/v3/handleSocketIo.server.ts index 99cf0b6e4..808e930b5 100644 --- a/apps/webapp/app/v3/handleSocketIo.server.ts +++ b/apps/webapp/app/v3/handleSocketIo.server.ts @@ -1,5 +1,6 @@ import { ClientToSharedQueueMessages, + CoordinatorSocketData, CoordinatorToPlatformMessages, PlatformToCoordinatorMessages, PlatformToProviderMessages, @@ -78,6 +79,7 @@ function createCoordinatorNamespace(io: Server) { authToken: env.COORDINATOR_SECRET, clientMessages: CoordinatorToPlatformMessages, serverMessages: PlatformToCoordinatorMessages, + socketData: CoordinatorSocketData, handlers: { READY_FOR_EXECUTION: async (message) => { const payload = await sharedQueueTasks.getLatestExecutionPayloadFromRun( @@ -238,6 +240,45 @@ function createCoordinatorNamespace(io: Server) { } }, }, + onConnection: async (socket, handler, sender, logger) => { + if (socket.data.supportsDynamicConfig) { + socket.emit("DYNAMIC_CONFIG", { + version: "v1", + checkpointThresholdInMs: env.CHECKPOINT_THRESHOLD_IN_MS, + }); + } + }, + postAuth: async (socket, next, logger) => { + function setSocketDataFromHeader( + dataKey: keyof typeof socket.data, + headerName: string, + required: boolean = true + ) { + const value = socket.handshake.headers[headerName]; + + if (value) { + socket.data[dataKey] = Array.isArray(value) ? value[0] : value; + return; + } + + if (required) { + logger.error("missing required header", { headerName }); + throw new Error("missing header"); + } + } + + try { + setSocketDataFromHeader("supportsDynamicConfig", "x-supports-dynamic-config", false); + } catch (error) { + logger.error("setSocketDataFromHeader error", { error }); + socket.disconnect(true); + return; + } + + logger.debug("success", socket.data); + + next(); + }, }); return coordinator.namespace; diff --git a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts index 9eb1140d2..3e298c14e 100644 --- a/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts +++ b/apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts @@ -268,6 +268,7 @@ export class SharedQueueConsumer { // TODO: For every ACK, decide what should be done with the existing run and attempts. Make sure to check the current statuses first. switch (messageBody.data.type) { + // MARK: EXECUTE case "EXECUTE": { const existingTaskRun = await prisma.taskRun.findUnique({ where: { @@ -478,7 +479,7 @@ export class SharedQueueConsumer { ? lockedTaskRun.attempts[0].number + 1 : 1; - const isRetry = nextAttemptNumber > 1; + const isRetry = lockedTaskRun.status === "WAITING_TO_RESUME" && nextAttemptNumber > 1; try { if (messageBody.data.checkpointEventId) { @@ -493,6 +494,8 @@ export class SharedQueueConsumer { logger.error("Failed to restore checkpoint", { queueMessage: message.data, messageId: message.messageId, + runStatus: lockedTaskRun.status, + isRetry, }); await this.#ackAndDoMoreWork(message.messageId); @@ -503,8 +506,18 @@ export class SharedQueueConsumer { } if (!deployment.worker.supportsLazyAttempts) { - const service = new CreateTaskRunAttemptService(); - await service.call(lockedTaskRun.friendlyId, undefined, false); + try { + const service = new CreateTaskRunAttemptService(); + await service.call(lockedTaskRun.friendlyId, undefined, false); + } catch (error) { + logger.error("Failed to create task run attempt for outdate worker", { + error, + taskRun: lockedTaskRun.id, + }); + + await this.#ackAndDoMoreWork(message.messageId); + return; + } } if (isRetry) { @@ -568,6 +581,7 @@ export class SharedQueueConsumer { break; } + // MARK: DEP RESUME // Resume after dependency completed with no remaining retries case "RESUME": { if (messageBody.data.checkpointEventId) { @@ -728,6 +742,11 @@ export class SharedQueueConsumer { } try { + logger.debug("Broadcasting RESUME_AFTER_DEPENDENCY", { + runId: resumableAttempt.taskRunId, + attemptId: resumableAttempt.id, + }); + // The attempt should still be running so we can broadcast to all coordinators to resume immediately socketIo.coordinatorNamespace.emit("RESUME_AFTER_DEPENDENCY", { version: "v1", @@ -752,6 +771,7 @@ export class SharedQueueConsumer { break; } + // MARK: DURATION RESUME // Resume after duration-based wait case "RESUME_AFTER_DURATION": { try { @@ -785,6 +805,7 @@ export class SharedQueueConsumer { break; } + // MARK: FAIL // Fail for whatever reason, usually runs that have been resumed but stopped heartbeating case "FAIL": { const existingTaskRun = await prisma.taskRun.findUnique({ @@ -1122,6 +1143,11 @@ class SharedQueueTasks { }, include: { lockedBy: true, + _count: { + select: { + attempts: true, + }, + }, }, }); @@ -1143,6 +1169,7 @@ class SharedQueueTasks { runId: run.friendlyId, messageId: run.id, isTest: run.isTest, + attemptCount: run._count.attempts, } satisfies TaskRunExecutionLazyAttemptPayload; } diff --git a/apps/webapp/app/v3/services/completeAttempt.server.ts b/apps/webapp/app/v3/services/completeAttempt.server.ts index bcf00fe9d..86e639862 100644 --- a/apps/webapp/app/v3/services/completeAttempt.server.ts +++ b/apps/webapp/app/v3/services/completeAttempt.server.ts @@ -344,6 +344,7 @@ export class CompleteAttemptService extends BaseService { supportsLazyAttempts?: boolean ) { if (checkpointEventId || !supportsLazyAttempts) { + // Workers without lazy attempt support always need to go through the queue, which is where the attempt is created // We have to replace a potential RESUME with EXECUTE to correctly retry the attempt return await marqs?.replaceMessage( run.id, @@ -355,8 +356,9 @@ export class CompleteAttemptService extends BaseService { retryTimestamp ); } else { - // There's no checkpoint so the worker is still running and waiting for a retry message - // It supports lazy attempts so we can bypass the queue and send the message directly to the worker + // There's no checkpoint and the worker supports lazy attempts + // This means the worker is still running and waiting for a retry message + // It supports lazy attempts so we can bypass the queue and send the message directly to it RetryAttemptService.enqueue(run.id, this._prisma, new Date(retryTimestamp)); } } diff --git a/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts b/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts index a119a5351..adde86186 100644 --- a/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts +++ b/apps/webapp/app/v3/services/createTaskRunAttempt.server.ts @@ -7,6 +7,8 @@ import { BaseService, ServiceValidationError } from "./baseService.server"; import { TaskRun, TaskRunAttempt } from "@trigger.dev/database"; import { machinePresetFromConfig } from "../machinePresets.server"; import { workerQueue } from "~/services/worker.server"; +import { MAX_TASK_RUN_ATTEMPTS } from "~/consts"; +import { CrashTaskRunService } from "./crashTaskRun.server"; export class CreateTaskRunAttemptService extends BaseService { public async call( @@ -93,6 +95,17 @@ export class CreateTaskRunAttemptService extends BaseService { const nextAttemptNumber = taskRun.attempts[0] ? taskRun.attempts[0].number + 1 : 1; + if (nextAttemptNumber > MAX_TASK_RUN_ATTEMPTS) { + const service = new CrashTaskRunService(this._prisma); + await service.call(taskRun.id, { + reason: taskRun.lockedBy.worker.supportsLazyAttempts + ? "Max attempts reached." + : "Max attempts reached. Please upgrade your CLI and SDK.", + }); + + throw new ServiceValidationError("Max attempts reached", 400); + } + const taskRunAttempt = await $transaction(this._prisma, async (tx) => { const taskRunAttempt = await tx.taskRunAttempt.create({ data: { diff --git a/apps/webapp/app/v3/services/resumeAttempt.server.ts b/apps/webapp/app/v3/services/resumeAttempt.server.ts index 42c14ba23..842faae3c 100644 --- a/apps/webapp/app/v3/services/resumeAttempt.server.ts +++ b/apps/webapp/app/v3/services/resumeAttempt.server.ts @@ -12,6 +12,7 @@ import { socketIo } from "../handleSocketIo.server"; import { SharedQueueMessageBody, sharedQueueTasks } from "../marqs/sharedQueueConsumer.server"; import { BaseService } from "./baseService.server"; import { TaskRunAttempt } from "@trigger.dev/database"; +import { isFinalRunStatus } from "../taskStatus"; export class ResumeAttemptService extends BaseService { public async call( @@ -80,10 +81,11 @@ export class ResumeAttemptService extends BaseService { return; } - if (attempt.taskRun.status !== "WAITING_TO_RESUME") { + if (isFinalRunStatus(attempt.taskRun.status)) { logger.error("Run is not resumable", { attemptId: attempt.id, runId: attempt.taskRunId, + status: attempt.taskRun.status, }); return; } diff --git a/packages/cli-v3/package.json b/packages/cli-v3/package.json index 415078096..80e76878f 100644 --- a/packages/cli-v3/package.json +++ b/packages/cli-v3/package.json @@ -117,7 +117,6 @@ "react-error-boundary": "^4.0.12", "semver": "^7.5.0", "simple-git": "^3.19.0", - "socket.io-client": "^4.7.4", "source-map-support": "^0.5.21", "terminal-link": "^3.0.0", "tiny-invariant": "^1.2.0", diff --git a/packages/cli-v3/src/commands/deploy.ts b/packages/cli-v3/src/commands/deploy.ts index 487ce21b2..a644215ef 100644 --- a/packages/cli-v3/src/commands/deploy.ts +++ b/packages/cli-v3/src/commands/deploy.ts @@ -1412,10 +1412,10 @@ async function resolveEnvironmentVariables( const projectConfig = config.config; return await tracer.startActiveSpan("resolveEnvironmentVariables", async (span) => { - try { - const $spinner = spinner(); - $spinner.start("Resolving environment variables"); + const $spinner = spinner(); + $spinner.start("Resolving environment variables"); + try { let processEnv: Record = { ...process.env, }; @@ -1473,10 +1473,14 @@ async function resolveEnvironmentVariables( $spinner.stop("No environment variables to sync"); return; } + } else { + $spinner.stop("No environment variables to sync"); } $spinner.stop("Environment variables resolved"); } catch (e) { + $spinner.stop("Failed to resolve environment variables"); + recordSpanException(span, e); throw e; diff --git a/packages/cli-v3/src/commands/dev.tsx b/packages/cli-v3/src/commands/dev.tsx index 8ca8359de..aeea0cf79 100644 --- a/packages/cli-v3/src/commands/dev.tsx +++ b/packages/cli-v3/src/commands/dev.tsx @@ -298,6 +298,7 @@ function useDev({ websocket.addEventListener("error", (event) => {}); // This is the deprecated task heart beat that uses the friendly attempt ID + // It will only be used if the worker does not support lazy attempts backgroundWorkerCoordinator.onWorkerTaskHeartbeat.attach( async ({ worker, backgroundWorkerId, id }) => { await sender.send("BACKGROUND_WORKER_MESSAGE", { diff --git a/packages/cli-v3/src/workers/prod/backgroundWorker.ts b/packages/cli-v3/src/workers/prod/backgroundWorker.ts index 88867d22d..76b0ad0a5 100644 --- a/packages/cli-v3/src/workers/prod/backgroundWorker.ts +++ b/packages/cli-v3/src/workers/prod/backgroundWorker.ts @@ -14,7 +14,6 @@ import { TaskRunExecutionLazyAttemptPayload, TaskRunExecutionPayload, TaskRunExecutionResult, - WaitReason, correctErrorStackTrace, } from "@trigger.dev/core/v3"; import { ZodIpcConnection } from "@trigger.dev/core/v3/zodIpc"; @@ -39,6 +38,19 @@ type BackgroundWorkerParams = { debugOtel?: boolean; }; +export type OnWaitForDurationMessage = InferSocketMessageSchema< + typeof ProdChildToWorkerMessages, + "WAIT_FOR_DURATION" +>; +export type OnWaitForTaskMessage = InferSocketMessageSchema< + typeof ProdChildToWorkerMessages, + "WAIT_FOR_TASK" +>; +export type OnWaitForBatchMessage = InferSocketMessageSchema< + typeof ProdChildToWorkerMessages, + "WAIT_FOR_BATCH" +>; + export class ProdBackgroundWorker { private _initialized: boolean = false; @@ -48,21 +60,9 @@ export class ProdBackgroundWorker { public onTaskHeartbeat: Evt = new Evt(); public onTaskRunHeartbeat: Evt = new Evt(); - public onWaitForBatch: Evt< - InferSocketMessageSchema - > = new Evt(); - public onWaitForDuration: Evt< - InferSocketMessageSchema - > = new Evt(); - public onWaitForTask: Evt< - InferSocketMessageSchema - > = new Evt(); - - public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>(); - public checkpointCanceledNotification = Evt.create<{ checkpointCanceled: boolean }>(); - - public onReadyForCheckpoint = Evt.create<{ version?: "v1" }>(); - public onCancelCheckpoint = Evt.create<{ version?: "v1" | "v2"; reason?: WaitReason }>(); + public onWaitForDuration: Evt = new Evt(); + public onWaitForTask: Evt = new Evt(); + public onWaitForBatch: Evt = new Evt(); public onCreateTaskRunAttempt = Evt.create<{ version?: "v1"; runId: string }>(); public attemptCreatedNotification = Evt.create< @@ -133,7 +133,12 @@ export class ProdBackgroundWorker { } async flushTelemetry() { + console.log("Flushing telemetry"); + const start = performance.now(); + await this._taskRunProcess?.cleanup(false); + + console.log("Flushed telemetry", { duration: performance.now() - start }); } async initialize(options?: { env?: Record }) { @@ -302,22 +307,6 @@ export class ProdBackgroundWorker { this.onWaitForTask.post(message); }); - taskRunProcess.onReadyForCheckpoint.attach((message) => { - this.onReadyForCheckpoint.post(message); - }); - - taskRunProcess.onCancelCheckpoint.attach((message) => { - this.onCancelCheckpoint.post(message); - }); - - // Notify down the chain - this.preCheckpointNotification.attach((message) => { - taskRunProcess.preCheckpointNotification.post(message); - }); - this.checkpointCanceledNotification.attach((message) => { - taskRunProcess.checkpointCanceledNotification.post(message); - }); - await taskRunProcess.initialize(); this._taskRunProcess = taskRunProcess; @@ -373,6 +362,8 @@ export class ProdBackgroundWorker { kill = false, initialSignal: number | NodeJS.Signals = "SIGTERM" ) { + console.log("Trying graceful exit", { kill, initialSignal }); + try { const initialExit = taskRunProcess.onExit.waitFor(5_000); @@ -389,6 +380,8 @@ export class ProdBackgroundWorker { } async #tryForcefulExit(taskRunProcess: TaskRunProcess) { + console.log("Trying forceful exit"); + try { const forcedKill = taskRunProcess.onExit.waitFor(5_000); taskRunProcess.kill("SIGKILL"); @@ -524,19 +517,24 @@ export class ProdBackgroundWorker { let execution: ProdTaskRunExecution; try { + const start = performance.now(); + // ..and wait for response - const attemptCreated = await this.attemptCreatedNotification.waitFor(30_000); + const attemptCreated = await this.attemptCreatedNotification.waitFor(120_000); if (!attemptCreated.success) { - throw new Error( - `Failed to create attempt${attemptCreated.reason ? `: ${attemptCreated.reason}` : ""}` - ); + throw new Error(`${attemptCreated.reason ?? "Unknown error"}`); } + console.log("Attempt created", { + number: attemptCreated.execution.attempt.number, + duration: performance.now() - start, + }); + execution = attemptCreated.execution; } catch (error) { console.error("Error while creating attempt", error); - throw new Error(`Failed to create task run attempt: ${error}`); + throw new Error(`Failed to create attempt: ${error}`); } const completion = await this.executeTaskRun( @@ -590,21 +588,11 @@ class TaskRunProcess { new Evt(); public onIsBeingKilled: Evt = new Evt(); - public onWaitForBatch: Evt< - InferSocketMessageSchema - > = new Evt(); - public onWaitForDuration: Evt< - InferSocketMessageSchema - > = new Evt(); - public onWaitForTask: Evt< - InferSocketMessageSchema - > = new Evt(); + public onWaitForDuration: Evt = new Evt(); + public onWaitForTask: Evt = new Evt(); + public onWaitForBatch: Evt = new Evt(); public preCheckpointNotification = Evt.create<{ willCheckpointAndRestore: boolean }>(); - public checkpointCanceledNotification = Evt.create<{ checkpointCanceled: boolean }>(); - - public onReadyForCheckpoint = Evt.create<{ version?: "v1" }>(); - public onCancelCheckpoint = Evt.create<{ version?: "v1" | "v2"; reason?: WaitReason }>(); constructor( private runId: string, @@ -664,6 +652,10 @@ class TaskRunProcess { if (this.messageId) { this.onTaskRunHeartbeat.post(this.messageId); } else { + console.error( + "No message id for task heartbeat, falling back to (deprecated) attempt heartbeat", + { id: message.id } + ); this.onTaskHeartbeat.post(message.id); } }, @@ -675,55 +667,7 @@ class TaskRunProcess { this.onWaitForBatch.post(message); }, WAIT_FOR_DURATION: async (message) => { - // Post to coordinator this.onWaitForDuration.post(message); - - try { - // ..and wait for response - const { willCheckpointAndRestore } = await this.preCheckpointNotification.waitFor( - 30_000 - ); - - return { - willCheckpointAndRestore, - }; - } catch (error) { - console.error("Error while waiting for pre-checkpoint notification", error); - - // Assume we won't get checkpointed - return { - willCheckpointAndRestore: false, - }; - } - }, - READY_FOR_CHECKPOINT: async (message) => { - this.onReadyForCheckpoint.post(message); - }, - CANCEL_CHECKPOINT: async (message) => { - const version = "v2"; - - // Post to coordinator - this.onCancelCheckpoint.post(message); - - try { - // ..and wait for response - const { checkpointCanceled } = await this.checkpointCanceledNotification.waitFor( - 30_000 - ); - - return { - version, - checkpointCanceled, - }; - } catch (error) { - console.error("Error while waiting for checkpoint cancellation", error); - - // Assume it's been canceled - return { - version, - checkpointCanceled: true, - }; - } }, }, }); @@ -764,14 +708,21 @@ class TaskRunProcess { realChildPid: this._child?.pid, }); - await this._ipc?.sendWithAck( - "CLEANUP", - { - flush: true, - kill: killParentProcess, - }, - 30_000 - ); + try { + await this._ipc?.sendWithAck( + "CLEANUP", + { + flush: true, + kill: killParentProcess, + }, + 30_000 + ); + } catch (error) { + console.error("Error while cleaning up task run process", error); + if (killParentProcess) { + process.exit(0); + } + } if (killChildProcess) { this._gracefulExitTimeoutElapsed = true; @@ -815,21 +766,34 @@ class TaskRunProcess { taskRunCompletedNotification(completion: TaskRunExecutionResult) { if (!completion.ok && typeof completion.retry !== "undefined") { + console.error( + "Task run completed with error and wants to retry, won't send task run completed notification" + ); return; } - if (this._child?.connected && !this._isBeingKilled && !this._child.killed) { - this._ipc?.send("TASK_RUN_COMPLETED_NOTIFICATION", { - version: "v2", - completion, - }); + if (!this._child?.connected || this._isBeingKilled || this._child.killed) { + console.error( + "Child process not connected or being killed, can't send task run completed notification" + ); + return; } + + this._ipc?.send("TASK_RUN_COMPLETED_NOTIFICATION", { + version: "v2", + completion, + }); } waitCompletedNotification() { - if (this._child?.connected && !this._isBeingKilled && !this._child.killed) { - this._ipc?.send("WAIT_COMPLETED_NOTIFICATION", {}); + if (!this._child?.connected || this._isBeingKilled || this._child.killed) { + console.error( + "Child process not connected or being killed, can't send wait completed notification" + ); + return; } + + this._ipc?.send("WAIT_COMPLETED_NOTIFICATION", {}); } async #handleExit(code: number | null, signal: NodeJS.Signals | null) { diff --git a/packages/cli-v3/src/workers/prod/entry-point.ts b/packages/cli-v3/src/workers/prod/entry-point.ts index e30388491..023b6dbcb 100644 --- a/packages/cli-v3/src/workers/prod/entry-point.ts +++ b/packages/cli-v3/src/workers/prod/entry-point.ts @@ -3,9 +3,11 @@ import { CoordinatorToProdWorkerMessages, PostStartCauses, PreStopCauses, + ProdTaskRunExecution, ProdWorkerToCoordinatorMessages, TaskResource, TaskRunErrorCodes, + TaskRunExecutionResult, TaskRunFailedExecutionResult, WaitReason, } from "@trigger.dev/core/v3"; @@ -13,11 +15,18 @@ import { InferSocketMessageSchema, ZodSocketConnection } from "@trigger.dev/core import { HttpReply, getRandomPortNumber } from "@trigger.dev/core-apps/http"; import { SimpleLogger } from "@trigger.dev/core-apps/logger"; import { EXIT_CODE_ALREADY_HANDLED, EXIT_CODE_CHILD_NONZERO } from "@trigger.dev/core-apps/process"; +import { ExponentialBackoff } from "@trigger.dev/core-apps/backoff"; +import { + OnWaitForBatchMessage, + OnWaitForTaskMessage, + ProdBackgroundWorker, +} from "./backgroundWorker"; +import { TaskMetadataParseError, UncaughtExceptionError } from "../common/errors"; +import { checkpointSafeTimeout, unboundedTimeout } from "@trigger.dev/core/v3/utils/timers"; +import { randomUUID } from "node:crypto"; import { readFile } from "node:fs/promises"; import { createServer } from "node:http"; -import { ProdBackgroundWorker } from "./backgroundWorker"; -import { TaskMetadataParseError, UncaughtExceptionError } from "../common/errors"; -import { setTimeout } from "node:timers/promises"; +import { setTimeout as timeout } from "node:timers/promises"; declare const __PROJECT_CONFIG__: Config; @@ -30,6 +39,10 @@ const SHORT_HASH = process.env.TRIGGER_CONTENT_HASH!.slice(0, 9); const logger = new SimpleLogger(`[${MACHINE_NAME}][${SHORT_HASH}]`); +const defaultBackoff = new ExponentialBackoff("FullJitter", { + maxRetries: 5, +}); + class ProdWorker { private apiUrl = process.env.TRIGGER_API_URL!; private apiKey = process.env.TRIGGER_SECRET_KEY!; @@ -48,6 +61,42 @@ class ProdWorker { private nextResumeAfter?: WaitReason; private waitForPostStart = false; + private connectionCount = 0; + + private waitForTaskReplay: + | { + idempotencyKey: string; + message: OnWaitForTaskMessage; + attempt: number; + } + | undefined; + private waitForBatchReplay: + | { + idempotencyKey: string; + message: OnWaitForBatchMessage; + attempt: number; + } + | undefined; + private readyForLazyAttemptReplay: + | { + idempotencyKey: string; + } + | undefined; + private submitAttemptCompletionReplay: + | { + idempotencyKey: string; + message: { + execution: ProdTaskRunExecution; + completion: TaskRunExecutionResult; + }; + attempt: number; + } + | undefined; + private durationResumeFallback: + | { + idempotencyKey: string; + } + | undefined; #httpPort: number; #backgroundWorker: ProdBackgroundWorker; @@ -84,7 +133,7 @@ class ProdWorker { }); // Wait for termination grace period minus 5s to give cleanup a chance to complete - await setTimeout(terminationGracePeriodSeconds * 1000 - 5000); + await timeout(terminationGracePeriodSeconds * 1000 - 5000); gracefulExitTimeoutElapsed = true; logger.log("Termination timeout reached, exiting gracefully."); @@ -108,16 +157,11 @@ class ProdWorker { } } - async #reconnect(isPostStart = false, reconnectImmediately = false) { - if (isPostStart) { - this.waitForPostStart = false; - } + async #reconnectAfterPostStart() { + this.waitForPostStart = false; this.#coordinatorSocket.close(); - - if (!reconnectImmediately) { - await setTimeout(1000); - } + this.connectionCount = 0; let coordinatorHost = COORDINATOR_HOST; @@ -145,6 +189,128 @@ class ProdWorker { } } + // MARK: TASK WAIT + async #waitForTaskHandler(message: OnWaitForTaskMessage, replayIdempotencyKey?: string) { + const waitForTask = await defaultBackoff.execute(async ({ retry }) => { + logger.log("Wait for task with backoff", { retry }); + + if (!this.attemptFriendlyId) { + logger.error("Failed to send wait message, attempt friendly ID not set", { message }); + + throw new ExponentialBackoff.StopRetrying("No attempt ID"); + } + + return await this.#coordinatorSocket.socket.timeout(20_000).emitWithAck("WAIT_FOR_TASK", { + version: "v2", + friendlyId: message.friendlyId, + attemptFriendlyId: this.attemptFriendlyId, + }); + }); + + if (!waitForTask.success) { + logger.error("Failed to wait for task with backoff", { + cause: waitForTask.cause, + error: waitForTask.error, + }); + + this.#emitUnrecoverableError( + "WaitForTaskFailed", + `${waitForTask.cause}: ${waitForTask.error}` + ); + + return; + } + + const { willCheckpointAndRestore } = waitForTask.result; + + await this.#prepareForWait("WAIT_FOR_TASK", willCheckpointAndRestore); + + if (willCheckpointAndRestore) { + // We need to replay this on next connection if we don't receive RESUME_AFTER_DEPENDENCY within a reasonable time + if (!this.waitForTaskReplay) { + this.waitForTaskReplay = { + message, + attempt: 1, + idempotencyKey: randomUUID(), + }; + } else { + if ( + replayIdempotencyKey && + replayIdempotencyKey !== this.waitForTaskReplay.idempotencyKey + ) { + logger.error( + "wait for task handler called with mismatched idempotency key, won't overwrite replay request" + ); + return; + } + + this.waitForTaskReplay.attempt++; + } + } + } + + // MARK: BATCH WAIT + async #waitForBatchHandler(message: OnWaitForBatchMessage, replayIdempotencyKey?: string) { + const waitForBatch = await defaultBackoff.execute(async ({ retry }) => { + logger.log("Wait for batch with backoff", { retry }); + + if (!this.attemptFriendlyId) { + logger.error("Failed to send wait message, attempt friendly ID not set", { message }); + + throw new ExponentialBackoff.StopRetrying("No attempt ID"); + } + + return await this.#coordinatorSocket.socket.timeout(20_000).emitWithAck("WAIT_FOR_BATCH", { + version: "v2", + batchFriendlyId: message.batchFriendlyId, + runFriendlyIds: message.runFriendlyIds, + attemptFriendlyId: this.attemptFriendlyId, + }); + }); + + if (!waitForBatch.success) { + logger.error("Failed to wait for batch with backoff", { + cause: waitForBatch.cause, + error: waitForBatch.error, + }); + + this.#emitUnrecoverableError( + "WaitForBatchFailed", + `${waitForBatch.cause}: ${waitForBatch.error}` + ); + + return; + } + + const { willCheckpointAndRestore } = waitForBatch.result; + + await this.#prepareForWait("WAIT_FOR_BATCH", willCheckpointAndRestore); + + if (willCheckpointAndRestore) { + // We need to replay this on next connection if we don't receive RESUME_AFTER_DEPENDENCY within a reasonable time + if (!this.waitForBatchReplay) { + this.waitForBatchReplay = { + message, + attempt: 1, + idempotencyKey: randomUUID(), + }; + } else { + if ( + replayIdempotencyKey && + replayIdempotencyKey !== this.waitForBatchReplay.idempotencyKey + ) { + logger.error( + "wait for task handler called with mismatched idempotency key, won't overwrite replay request" + ); + return; + } + + this.waitForBatchReplay.attempt++; + } + } + } + + // MARK: WORKER CREATION #createBackgroundWorker() { const backgroundWorker = new ProdBackgroundWorker("worker.js", { projectConfig: __PROJECT_CONFIG__, @@ -159,73 +325,71 @@ class ProdWorker { }); backgroundWorker.onTaskHeartbeat.attach((attemptFriendlyId) => { - // TODO: Switch to .send() once coordinator uses zod handler for all messages - this.#coordinatorSocket.socket.emit("TASK_HEARTBEAT", { version: "v1", attemptFriendlyId }); + logger.log("onTaskHeartbeat", { attemptFriendlyId }); + + this.#coordinatorSocket.socket.volatile.emit("TASK_HEARTBEAT", { + version: "v1", + attemptFriendlyId, + }); }); backgroundWorker.onTaskRunHeartbeat.attach((runId) => { - this.#coordinatorSocket.socket.emit("TASK_RUN_HEARTBEAT", { version: "v1", runId }); - }); + logger.log("onTaskRunHeartbeat", { runId }); - // Currently, this is only used for duration waits - backgroundWorker.onReadyForCheckpoint.attach(async (message) => { - await this.#prepareForCheckpoint(); - - this.#coordinatorSocket.socket.emit("READY_FOR_CHECKPOINT", { version: "v1" }); - }); - - // Currently, this is only used for duration waits. Might need adjusting for other use cases. - backgroundWorker.onCancelCheckpoint.attach(async (message) => { - logger.log("onCancelCheckpoint", { message }); - - const { checkpointCanceled } = await this.#coordinatorSocket.socket.emitWithAck( - "CANCEL_CHECKPOINT", - { - version: "v2", - reason: message.reason, - } - ); - - logger.log("onCancelCheckpoint coordinator response", { checkpointCanceled }); - - if (checkpointCanceled) { - if (message.reason === "WAIT_FOR_DURATION") { - // Worker will resume immediately - this.paused = false; - this.nextResumeAfter = undefined; - this.waitForPostStart = false; - } - } - - backgroundWorker.checkpointCanceledNotification.post({ checkpointCanceled }); + this.#coordinatorSocket.socket.volatile.emit("TASK_RUN_HEARTBEAT", { version: "v1", runId }); }); backgroundWorker.onCreateTaskRunAttempt.attach(async (message) => { logger.log("onCreateTaskRunAttempt()", { message }); - const createAttempt = await this.#coordinatorSocket.socket.emitWithAck( - "CREATE_TASK_RUN_ATTEMPT", - { - version: "v1", - runId: message.runId, - } - ); + const createAttempt = await defaultBackoff.execute(async ({ retry }) => { + logger.log("Create task run attempt with backoff", { retry }); + + return await this.#coordinatorSocket.socket + .timeout(15_000) + .emitWithAck("CREATE_TASK_RUN_ATTEMPT", { + version: "v1", + runId: message.runId, + }); + }); if (!createAttempt.success) { backgroundWorker.attemptCreatedNotification.post({ success: false, - reason: createAttempt.reason, + reason: `Failed to create attempt with backoff due to ${createAttempt.cause}. ${createAttempt.error}`, + }); + return; + } + + if (!createAttempt.result.success) { + backgroundWorker.attemptCreatedNotification.post({ + success: false, + reason: createAttempt.result.reason, }); return; } backgroundWorker.attemptCreatedNotification.post({ success: true, - execution: createAttempt.executionPayload.execution, + execution: createAttempt.result.executionPayload.execution, }); }); backgroundWorker.attemptCreatedNotification.attach((message) => { + logger.log("attemptCreatedNotification", { + success: message.success, + ...(message.success + ? { + attempt: message.execution.attempt, + queue: message.execution.queue, + worker: message.execution.worker, + machine: message.execution.machine, + } + : { + reason: message.reason, + }), + }); + if (!message.success) { return; } @@ -235,67 +399,114 @@ class ProdWorker { }); backgroundWorker.onWaitForDuration.attach(async (message) => { - if (!this.attemptFriendlyId) { - logger.error("Failed to send wait message, attempt friendly ID not set", { message }); + logger.log("onWaitForDuration", { ...message, drift: Date.now() - message.now }); - this.#emitUnrecoverableError( - "NoAttemptId", - "Attempt ID not set before waiting for duration" - ); + noResume: { + const { ms, waitThresholdInMs } = message; + + const internalTimeout = unboundedTimeout(ms, "internal" as const); + const checkpointSafeInternalTimeout = checkpointSafeTimeout(ms); + + if (ms < waitThresholdInMs) { + await internalTimeout; + break noResume; + } + + const waitForDuration = await defaultBackoff.execute(async ({ retry }) => { + logger.log("Wait for duration with backoff", { retry }); + + if (!this.attemptFriendlyId) { + logger.error("Failed to send wait message, attempt friendly ID not set", { message }); + + throw new ExponentialBackoff.StopRetrying("No attempt ID"); + } + + return await this.#coordinatorSocket.socket + .timeout(20_000) + .emitWithAck("WAIT_FOR_DURATION", { + ...message, + attemptFriendlyId: this.attemptFriendlyId, + }); + }); + + if (!waitForDuration.success) { + logger.error("Failed to wait for duration with backoff", { + cause: waitForDuration.cause, + error: waitForDuration.error, + }); + + this.#emitUnrecoverableError( + "WaitForDurationFailed", + `${waitForDuration.cause}: ${waitForDuration.error}` + ); + + return; + } + + const { willCheckpointAndRestore } = waitForDuration.result; + + if (!willCheckpointAndRestore) { + await internalTimeout; + break noResume; + } + + await this.#prepareForWait("WAIT_FOR_DURATION", willCheckpointAndRestore); + // CHECKPOINTING AFTER THIS LINE + + // internalTimeout acts as a backup and will be accurate if the checkpoint never happens + // checkpointSafeInternalTimeout is accurate even after non-simulated restores + await Promise.race([internalTimeout, checkpointSafeInternalTimeout]); + + try { + // The coordinator should cancel any in-progress checkpoints so we don't end up with race conditions + const { checkpointCanceled } = await this.#coordinatorSocket.socket + .timeout(15_000) + .emitWithAck("CANCEL_CHECKPOINT", { + version: "v2", + reason: "WAIT_FOR_DURATION", + }); + + logger.log("onCancelCheckpoint coordinator response", { checkpointCanceled }); + + if (checkpointCanceled) { + // If the checkpoint was canceled, we will never be resumed externally with RESUME_AFTER_DURATION, so it's safe to immediately resume + break noResume; + } + + logger.log("Waiting for external duration resume as we may have been restored"); + + const idempotencyKey = randomUUID(); + this.durationResumeFallback = { idempotencyKey }; + + setTimeout(() => { + if (!this.durationResumeFallback) { + logger.error("Already resumed after duration, skipping fallback"); + return; + } + + if (this.durationResumeFallback.idempotencyKey !== idempotencyKey) { + logger.error("Duration resume idempotency key mismatch, skipping fallback"); + return; + } + + logger.log("Resuming after duration with fallback"); + + this.#resumeAfterDuration(); + }, 15_000); + } catch (error) { + // If the cancellation times out, we will proceed as if the checkpoint was canceled + logger.debug("Checkpoint cancellation timed out", { error }); + break noResume; + } return; } - const { willCheckpointAndRestore } = await this.#coordinatorSocket.socket.emitWithAck( - "WAIT_FOR_DURATION", - { - ...message, - attemptFriendlyId: this.attemptFriendlyId, - } - ); - - this.#prepareForWait("WAIT_FOR_DURATION", willCheckpointAndRestore); + this.#resumeAfterDuration(); }); - backgroundWorker.onWaitForTask.attach(async (message) => { - if (!this.attemptFriendlyId) { - logger.error("Failed to send wait message, attempt friendly ID not set", { message }); - - this.#emitUnrecoverableError("NoAttemptId", "Attempt ID not set before waiting for task"); - - return; - } - - const { willCheckpointAndRestore } = await this.#coordinatorSocket.socket.emitWithAck( - "WAIT_FOR_TASK", - { - ...message, - attemptFriendlyId: this.attemptFriendlyId, - } - ); - - this.#prepareForWait("WAIT_FOR_TASK", willCheckpointAndRestore); - }); - - backgroundWorker.onWaitForBatch.attach(async (message) => { - if (!this.attemptFriendlyId) { - logger.error("Failed to send wait message, attempt friendly ID not set", { message }); - - this.#emitUnrecoverableError("NoAttemptId", "Attempt ID not set before waiting for batch"); - - return; - } - - const { willCheckpointAndRestore } = await this.#coordinatorSocket.socket.emitWithAck( - "WAIT_FOR_BATCH", - { - ...message, - attemptFriendlyId: this.attemptFriendlyId, - } - ); - - this.#prepareForWait("WAIT_FOR_BATCH", willCheckpointAndRestore); - }); + backgroundWorker.onWaitForTask.attach(this.#waitForTaskHandler.bind(this)); + backgroundWorker.onWaitForBatch.attach(this.#waitForBatchHandler.bind(this)); return backgroundWorker; } @@ -303,31 +514,29 @@ class ProdWorker { async #prepareForWait(reason: WaitReason, willCheckpointAndRestore: boolean) { logger.log(`prepare for ${reason}`, { willCheckpointAndRestore }); - this.#backgroundWorker.preCheckpointNotification.post({ willCheckpointAndRestore }); - - if (willCheckpointAndRestore) { - this.paused = true; - this.nextResumeAfter = reason; - this.waitForPostStart = true; - - if (reason === "WAIT_FOR_TASK" || reason === "WAIT_FOR_BATCH") { - // Duration waits do this via the "ready for checkpoint" event instead - await this.#prepareForCheckpoint(); - } + if (!willCheckpointAndRestore) { + return; } + + this.paused = true; + this.nextResumeAfter = reason; + this.waitForPostStart = true; + + await this.#prepareForCheckpoint(); } + // MARK: RETRY PREP async #prepareForRetry( willCheckpointAndRestore: boolean, shouldExit: boolean, exitCode?: number ) { - logger.log("prepare for retry", { willCheckpointAndRestore, shouldExit }); + logger.log("prepare for retry", { willCheckpointAndRestore, shouldExit, exitCode }); // Graceful shutdown on final attempt if (shouldExit) { if (willCheckpointAndRestore) { - logger.log("WARNING: Will checkpoint but also requested exit. This won't end well."); + logger.error("WARNING: Will checkpoint but also requested exit. This won't end well."); } await this.#exitGracefully(false, exitCode); @@ -340,25 +549,41 @@ class ProdWorker { this.executing = false; this.attemptFriendlyId = undefined; - if (willCheckpointAndRestore) { - this.waitForPostStart = true; - - // We already flush after completion, so we don't need to do it here - this.#prepareForCheckpoint(false); - - this.#coordinatorSocket.socket.emit("READY_FOR_CHECKPOINT", { version: "v1" }); + if (!willCheckpointAndRestore) { return; } + + this.waitForPostStart = true; + + // We already flush after completion, so we don't need to do it here + await this.#prepareForCheckpoint(false); } + // MARK: CHECKPOINT PREP async #prepareForCheckpoint(flush = true) { if (flush) { // Flush before checkpointing so we don't flush the same spans again after restore - await this.#backgroundWorker.flushTelemetry(); + try { + await this.#backgroundWorker.flushTelemetry(); + } catch (error) { + logger.error( + "Failed to flush telemetry while preparing for checkpoint, will proceed anyway", + { error } + ); + } } - // Kill the previous worker process to prevent large checkpoints - await this.#backgroundWorker.forceKillOldTaskRunProcesses(); + try { + // Kill the previous worker process to prevent large checkpoints + await this.#backgroundWorker.forceKillOldTaskRunProcesses(); + } catch (error) { + logger.error( + "Failed to kill previous worker while preparing for checkpoint, will proceed anyway", + { error } + ); + } + + this.#readyForCheckpoint(); } #resumeAfterDuration() { @@ -369,6 +594,152 @@ class ProdWorker { this.#backgroundWorker.waitCompletedNotification(); } + async #readyForLazyAttempt() { + const idempotencyKey = randomUUID(); + + this.readyForLazyAttemptReplay = { + idempotencyKey, + }; + + // Retry if we don't receive EXECUTE_TASK_RUN_LAZY_ATTEMPT in a reasonable time + // ..but we also have to be fast to avoid failing the task due to missing heartbeat + for await (const { delay, retry } of defaultBackoff.min(10).maxRetries(3)) { + if (retry > 0) { + logger.log("retrying ready for lazy attempt", { retry }); + } + + this.#coordinatorSocket.socket.emit("READY_FOR_LAZY_ATTEMPT", { + version: "v1", + runId: this.runId, + totalCompletions: this.completed.size, + }); + + await timeout(delay.milliseconds); + + if (!this.readyForLazyAttemptReplay) { + logger.error("replay ready for lazy attempt cancelled, discarding", { + idempotencyKey, + }); + + return; + } + + if (idempotencyKey !== this.readyForLazyAttemptReplay.idempotencyKey) { + logger.error("replay ready for lazy attempt idempotency key mismatch, discarding", { + idempotencyKey, + newIdempotencyKey: this.readyForLazyAttemptReplay.idempotencyKey, + }); + + return; + } + } + + // Fail the task with a more descriptive message as it likely failed with a generic missing heartbeat error + this.#failRun(this.runId, "Failed to receive execute request in a reasonable time"); + } + + #readyForCheckpoint() { + this.#coordinatorSocket.socket.emit("READY_FOR_CHECKPOINT", { version: "v1" }); + } + + #failRun(anyRunId: string, error: unknown) { + logger.error("Failing run", { anyRunId, error }); + + const completion: TaskRunFailedExecutionResult = { + ok: false, + id: anyRunId, + retry: undefined, + error: + error instanceof Error + ? { + type: "BUILT_IN_ERROR", + name: error.name, + message: error.message, + stackTrace: error.stack ?? "", + } + : { + type: "BUILT_IN_ERROR", + name: "UnknownError", + message: String(error), + stackTrace: "", + }, + }; + + this.#coordinatorSocket.socket.emit("TASK_RUN_FAILED_TO_RUN", { + version: "v1", + completion, + }); + } + + // MARK: ATTEMPT COMPLETION + async #submitAttemptCompletion( + execution: ProdTaskRunExecution, + completion: TaskRunExecutionResult, + replayIdempotencyKey?: string + ) { + const taskRunCompleted = await defaultBackoff.execute(async ({ retry }) => { + logger.log("Submit attempt completion with backoff", { retry }); + + return await this.#coordinatorSocket.socket + .timeout(20_000) + .emitWithAck("TASK_RUN_COMPLETED", { + version: "v1", + execution, + completion, + }); + }); + + if (!taskRunCompleted.success) { + logger.error("Failed to complete lazy attempt with backoff", { + cause: taskRunCompleted.cause, + error: taskRunCompleted.error, + }); + + this.#failRun(execution.run.id, taskRunCompleted.error); + + return; + } + + const { willCheckpointAndRestore, shouldExit } = taskRunCompleted.result; + + logger.log("completion acknowledged", { willCheckpointAndRestore, shouldExit }); + + const exitCode = + !completion.ok && + completion.error.type === "INTERNAL_ERROR" && + completion.error.code === TaskRunErrorCodes.TASK_PROCESS_EXITED_WITH_NON_ZERO_CODE + ? EXIT_CODE_CHILD_NONZERO + : 0; + + await this.#prepareForRetry(willCheckpointAndRestore, shouldExit, exitCode); + + if (willCheckpointAndRestore) { + // We need to replay this on next connection if we don't receive READY_FOR_RETRY within a reasonable time + if (!this.submitAttemptCompletionReplay) { + this.submitAttemptCompletionReplay = { + message: { + execution, + completion, + }, + attempt: 1, + idempotencyKey: randomUUID(), + }; + } else { + if ( + replayIdempotencyKey && + replayIdempotencyKey !== this.submitAttemptCompletionReplay.idempotencyKey + ) { + logger.error( + "attempt completion handler called with mismatched idempotency key, won't overwrite replay request" + ); + return; + } + + this.submitAttemptCompletionReplay.attempt++; + } + } + } + #returnValidatedExtraHeaders(headers: Record) { for (const [key, value] of Object.entries(headers)) { if (value === undefined) { @@ -379,7 +750,7 @@ class ProdWorker { return headers; } - // FIXME: If the the worker can't connect for a while, this runs MANY times - it should only run once + // MARK: COORDINATOR SOCKET #createCoordinatorSocket(host: string) { const extraHeaders = this.#returnValidatedExtraHeaders({ "x-machine-name": MACHINE_NAME, @@ -406,6 +777,10 @@ class ProdWorker { clientMessages: ProdWorkerToCoordinatorMessages, serverMessages: CoordinatorToProdWorkerMessages, extraHeaders, + ioOptions: { + reconnectionDelay: 1000, + reconnectionDelayMax: 3000, + }, handlers: { RESUME_AFTER_DEPENDENCY: async ({ completions }) => { if (!this.paused) { @@ -438,6 +813,17 @@ class ProdWorker { return; } + switch (this.nextResumeAfter) { + case "WAIT_FOR_TASK": { + this.waitForTaskReplay = undefined; + break; + } + case "WAIT_FOR_BATCH": { + this.waitForBatchReplay = undefined; + break; + } + } + this.paused = false; this.nextResumeAfter = undefined; this.waitForPostStart = false; @@ -465,8 +851,11 @@ class ProdWorker { return; } + this.durationResumeFallback = undefined; + this.#resumeAfterDuration(); }, + // Deprecated: This will never get called as this worker supports lazy attempts. It's only here for a quick view of the flow old workers use. EXECUTE_TASK_RUN: async ({ executionPayload }) => { if (this.executing) { logger.error("dropping execute request, already executing"); @@ -495,14 +884,25 @@ class ProdWorker { logger.log("completion acknowledged", { willCheckpointAndRestore, shouldExit }); - this.#prepareForRetry(willCheckpointAndRestore, shouldExit); + await this.#prepareForRetry(willCheckpointAndRestore, shouldExit); }, EXECUTE_TASK_RUN_LAZY_ATTEMPT: async (message) => { + this.readyForLazyAttemptReplay = undefined; + if (this.executing) { logger.error("dropping execute request, already executing"); return; } + const attemptCount = message.lazyPayload.attemptCount ?? 0; + + logger.log("execute attempt counts", { attemptCount, completed: this.completed.size }); + + if (this.completed.size > 0 && this.completed.size >= attemptCount + 1) { + logger.error("dropping execute request, already completed"); + return; + } + this.executing = true; try { @@ -513,48 +913,13 @@ class ProdWorker { this.completed.add(execution.attempt.id); - const { willCheckpointAndRestore, shouldExit } = - await this.#coordinatorSocket.socket.emitWithAck("TASK_RUN_COMPLETED", { - version: "v1", - execution, - completion, - }); - - logger.log("completion acknowledged", { willCheckpointAndRestore, shouldExit }); - - const exitCode = - !completion.ok && - completion.error.type === "INTERNAL_ERROR" && - completion.error.code === TaskRunErrorCodes.TASK_PROCESS_EXITED_WITH_NON_ZERO_CODE - ? EXIT_CODE_CHILD_NONZERO - : 0; - - this.#prepareForRetry(willCheckpointAndRestore, shouldExit, exitCode); + await this.#submitAttemptCompletion(execution, completion); } catch (error) { - const completion: TaskRunFailedExecutionResult = { - ok: false, - id: message.lazyPayload.runId, - retry: undefined, - error: - error instanceof Error - ? { - type: "BUILT_IN_ERROR", - name: error.name, - message: error.message, - stackTrace: error.stack ?? "", - } - : { - type: "BUILT_IN_ERROR", - name: "UnknownError", - message: String(error), - stackTrace: "", - }, - }; - - this.#coordinatorSocket.socket.emit("TASK_RUN_FAILED_TO_RUN", { - version: "v1", - completion, + logger.error("Failed to complete lazy attempt", { + error, }); + + this.#failRun(message.lazyPayload.runId, error); } }, REQUEST_ATTEMPT_CANCELLATION: async (message) => { @@ -570,7 +935,7 @@ class ProdWorker { REQUEST_EXIT: async (message) => { if (message.version === "v2" && message.delayInMs) { logger.log("exit requested with delay", { delayInMs: message.delayInMs }); - await setTimeout(message.delayInMs); + await timeout(message.delayInMs); } this.#coordinatorSocket.close(); @@ -578,157 +943,185 @@ class ProdWorker { }, READY_FOR_RETRY: async (message) => { if (this.completed.size < 1) { + logger.error("Received READY_FOR_RETRY but no completions yet. This is a bug."); return; } - this.#coordinatorSocket.socket.emit("READY_FOR_LAZY_ATTEMPT", { - version: "v1", - runId: this.runId, - totalCompletions: this.completed.size, - }); + this.submitAttemptCompletionReplay = undefined; + + await this.#readyForLazyAttempt(); }, }, + // MARK: ON CONNECTION onConnection: async (socket, handler, sender, logger) => { - logger.log("connected to coordinator", { status: this.#status }); + logger.log("connected to coordinator", { + status: this.#status, + connectionCount: ++this.connectionCount, + }); - if (this.waitForPostStart) { - logger.log("skip connection handler, waiting for post start hook"); - return; - } + // We need to send our current state to the coordinator + socket.emit("SET_STATE", { version: "v1", attemptFriendlyId: this.attemptFriendlyId }); - if (this.paused) { - if (!this.nextResumeAfter) { - logger.error("Missing next resume reason", { status: this.#status }); + try { + if (this.waitForPostStart) { + logger.log("skip connection handler, waiting for post start hook"); + return; + } - this.#emitUnrecoverableError( - "NoNextResume", - "Next resume reason not set while resuming from paused state" - ); + if (this.paused) { + if (!this.nextResumeAfter) { + logger.error("Missing next resume reason", { status: this.#status }); + + this.#emitUnrecoverableError( + "NoNextResume", + "Next resume reason not set while resuming from paused state" + ); + + return; + } + + if (!this.attemptFriendlyId) { + logger.error("Missing friendly ID", { status: this.#status }); + + this.#emitUnrecoverableError( + "NoAttemptId", + "Attempt ID not set while resuming from paused state" + ); + + return; + } + + socket.emit("READY_FOR_RESUME", { + version: "v1", + attemptFriendlyId: this.attemptFriendlyId, + type: this.nextResumeAfter, + }); return; } - if (!this.attemptFriendlyId) { - logger.error("Missing friendly ID", { status: this.#status }); + if (process.env.INDEX_TASKS === "true") { + const failIndex = ( + error: InferSocketMessageSchema< + typeof ProdWorkerToCoordinatorMessages, + "INDEXING_FAILED" + >["error"] + ) => { + socket.emit("INDEXING_FAILED", { + version: "v1", + deploymentId: this.deploymentId, + error, + }); + }; - this.#emitUnrecoverableError( - "NoAttemptId", - "Attempt ID not set while resuming from paused state" - ); + process.removeAllListeners("uncaughtException"); + process.on("uncaughtException", (error) => { + console.error("Uncaught exception while indexing", error); + failIndex(error); + }); + try { + const taskResources = await this.#initializeWorker(); + + const indexTasks = await defaultBackoff.maxRetries(3).execute(async () => { + return await socket.timeout(20_000).emitWithAck("INDEX_TASKS", { + version: "v2", + deploymentId: this.deploymentId, + ...taskResources, + supportsLazyAttempts: true, + }); + }); + + if (!indexTasks.success || !indexTasks.result.success) { + logger.error("indexing failure, shutting down..", { indexTasks }); + process.exit(1); + } else { + logger.info("indexing done, shutting down.."); + process.exit(0); + } + } catch (e) { + const stderr = this.#backgroundWorker.stderr.join("\n"); + + if (e instanceof TaskMetadataParseError) { + logger.error("tasks metadata parse error", { + zodIssues: e.zodIssues, + tasks: e.tasks, + }); + + failIndex({ + name: "TaskMetadataParseError", + message: "There was an error parsing the task metadata", + stack: JSON.stringify({ zodIssues: e.zodIssues, tasks: e.tasks }), + stderr, + }); + } else if (e instanceof UncaughtExceptionError) { + const error = { + name: e.originalError.name, + message: e.originalError.message, + stack: e.originalError.stack, + stderr, + }; + + logger.error("uncaught exception", { originalError: error }); + + failIndex(error); + } else if (e instanceof Error) { + const error = { + name: e.name, + message: e.message, + stack: e.stack, + stderr, + }; + + logger.error("error", { error }); + + failIndex(error); + } else if (typeof e === "string") { + logger.error("string error", { error: { message: e } }); + + failIndex({ + name: "Error", + message: e, + stderr, + }); + } else { + logger.error("unknown error", { error: e }); + + failIndex({ + name: "Error", + message: "Unknown error", + stderr, + }); + } + + await timeout(1000); + + process.exit(EXIT_CODE_ALREADY_HANDLED); + } + } + + if (this.executing) { return; } - socket.emit("READY_FOR_RESUME", { - version: "v1", - attemptFriendlyId: this.attemptFriendlyId, - type: this.nextResumeAfter, + process.removeAllListeners("uncaughtException"); + process.on("uncaughtException", (error) => { + console.error("Uncaught exception during run", error); + this.#failRun(this.runId, error); }); - return; - } - - if (process.env.INDEX_TASKS === "true") { - const failIndex = ( - error: InferSocketMessageSchema< - typeof ProdWorkerToCoordinatorMessages, - "INDEXING_FAILED" - >["error"] - ) => { - socket.emit("INDEXING_FAILED", { - version: "v1", - deploymentId: this.deploymentId, - error, - }); - }; - - try { - const taskResources = await this.#initializeWorker(); - - const { success } = await socket.emitWithAck("INDEX_TASKS", { - version: "v2", - deploymentId: this.deploymentId, - ...taskResources, - supportsLazyAttempts: true, - }); - - if (success) { - logger.info("indexing done, shutting down.."); - process.exit(0); - } else { - logger.info("indexing failure, shutting down.."); - process.exit(1); - } - } catch (e) { - const stderr = this.#backgroundWorker.stderr.join("\n"); - - if (e instanceof TaskMetadataParseError) { - logger.error("tasks metadata parse error", { - zodIssues: e.zodIssues, - tasks: e.tasks, - }); - - failIndex({ - name: "TaskMetadataParseError", - message: "There was an error parsing the task metadata", - stack: JSON.stringify({ zodIssues: e.zodIssues, tasks: e.tasks }), - stderr, - }); - } else if (e instanceof UncaughtExceptionError) { - const error = { - name: e.originalError.name, - message: e.originalError.message, - stack: e.originalError.stack, - stderr, - }; - - logger.error("uncaught exception", { originalError: error }); - - failIndex(error); - } else if (e instanceof Error) { - const error = { - name: e.name, - message: e.message, - stack: e.stack, - stderr, - }; - - logger.error("error", { error }); - - failIndex(error); - } else if (typeof e === "string") { - logger.error("string error", { error: { message: e } }); - - failIndex({ - name: "Error", - message: e, - stderr, - }); - } else { - logger.error("unknown error", { error: e }); - - failIndex({ - name: "Error", - message: "Unknown error", - stderr, - }); - } - - await setTimeout(200); - - process.exit(EXIT_CODE_ALREADY_HANDLED); + await this.#readyForLazyAttempt(); + } catch (error) { + logger.error("connection handler error", { error }); + } finally { + if (this.connectionCount === 1) { + // Skip replays if this is the first connection, including post start + return; } - } - if (this.executing) { - return; + // This is a reconnect, so handle replays + this.#handleReplays(); } - - socket.emit("READY_FOR_LAZY_ATTEMPT", { - version: "v1", - runId: this.runId, - totalCompletions: this.completed.size, - }); }, onError: async (socket, err, logger) => { logger.error("onError", { @@ -737,17 +1130,145 @@ class ProdWorker { message: err.message, }, }); - - await this.#reconnect(); - }, - onDisconnect: async (socket, reason, description, logger) => { - // this.#reconnect(); }, }); return coordinatorConnection; } + // MARK: REPLAYS + async #handleReplays() { + const backoff = new ExponentialBackoff().type("FullJitter").maxRetries(3); + const replayCancellationDelay = 20_000; + + if (this.waitForTaskReplay) { + logger.log("replaying wait for task", { ...this.waitForTaskReplay }); + + const { idempotencyKey, message, attempt } = this.waitForTaskReplay; + + // Give the platform some time to send RESUME_AFTER_DEPENDENCY + await timeout(replayCancellationDelay); + + if (!this.waitForTaskReplay) { + logger.error("wait for task replay cancelled, discarding", { + originalMessage: { idempotencyKey, message, attempt }, + }); + + return; + } + + if (idempotencyKey !== this.waitForTaskReplay.idempotencyKey) { + logger.error("wait for task replay idempotency key mismatch, discarding", { + originalMessage: { idempotencyKey, message, attempt }, + newMessage: this.waitForTaskReplay, + }); + + return; + } + + try { + await backoff.wait(attempt + 1); + + await this.#waitForTaskHandler(message); + } catch (error) { + if (error instanceof ExponentialBackoff.RetryLimitExceeded) { + logger.error("wait for task replay retry limit exceeded", { error }); + } else { + logger.error("wait for task replay error", { error }); + } + } + + return; + } + + if (this.waitForBatchReplay) { + logger.log("replaying wait for batch", { + ...this.waitForBatchReplay, + cancellationDelay: replayCancellationDelay, + }); + + const { idempotencyKey, message, attempt } = this.waitForBatchReplay; + + // Give the platform some time to send RESUME_AFTER_DEPENDENCY + await timeout(replayCancellationDelay); + + if (!this.waitForBatchReplay) { + logger.error("wait for batch replay cancelled, discarding", { + originalMessage: { idempotencyKey, message, attempt }, + }); + + return; + } + + if (idempotencyKey !== this.waitForBatchReplay.idempotencyKey) { + logger.error("wait for batch replay idempotency key mismatch, discarding", { + originalMessage: { idempotencyKey, message, attempt }, + newMessage: this.waitForBatchReplay, + }); + + return; + } + + try { + await backoff.wait(attempt + 1); + + await this.#waitForBatchHandler(message); + } catch (error) { + if (error instanceof ExponentialBackoff.RetryLimitExceeded) { + logger.error("wait for batch replay retry limit exceeded", { error }); + } else { + logger.error("wait for batch replay error", { error }); + } + } + + return; + } + + if (this.submitAttemptCompletionReplay) { + logger.log("replaying attempt completion", { + ...this.submitAttemptCompletionReplay, + cancellationDelay: replayCancellationDelay, + }); + + const { idempotencyKey, message, attempt } = this.submitAttemptCompletionReplay; + + // Give the platform some time to send READY_FOR_RETRY + await timeout(replayCancellationDelay); + + if (!this.submitAttemptCompletionReplay) { + logger.error("attempt completion replay cancelled, discarding", { + originalMessage: { idempotencyKey, message, attempt }, + }); + + return; + } + + if (idempotencyKey !== this.submitAttemptCompletionReplay.idempotencyKey) { + logger.error("attempt completion replay idempotency key mismatch, discarding", { + originalMessage: { idempotencyKey, message, attempt }, + newMessage: this.submitAttemptCompletionReplay, + }); + + return; + } + + try { + await backoff.wait(attempt + 1); + + await this.#submitAttemptCompletion(message.execution, message.completion, idempotencyKey); + } catch (error) { + if (error instanceof ExponentialBackoff.RetryLimitExceeded) { + logger.error("attempt completion replay retry limit exceeded", { error }); + } else { + logger.error("attempt completion replay error", { error }); + } + } + + return; + } + } + + // MARK: HTTP SERVER #createHttpServer() { const httpServer = createServer(async (req, res) => { logger.log(`[${req.method}]`, req.url); @@ -772,20 +1293,15 @@ class ProdWorker { } case "/close": { - await this.#coordinatorSocket.sendWithAck("LOG", { - version: "v1", - text: `[${req.method}] ${req.url}`, - }); - this.#coordinatorSocket.close(); + this.connectionCount = 0; return reply.text("Disconnected from coordinator"); } case "/test": { - await this.#coordinatorSocket.sendWithAck("LOG", { + await this.#coordinatorSocket.socket.timeout(10_000).emitWithAck("TEST", { version: "v1", - text: `[${req.method}] ${req.url}`, }); return reply.text("Received ACK from coordinator"); @@ -828,7 +1344,7 @@ class ProdWorker { break; } case "restore": { - await this.#reconnect(true, true); + await this.#reconnectAfterPostStart(); break; } default: { @@ -868,7 +1384,7 @@ class ProdWorker { this.#httpPort = getRandomPortNumber(); - await setTimeout(100); + await timeout(100); this.start(); }); @@ -887,8 +1403,12 @@ class ProdWorker { const taskResources: Array = []; - if (!this.#backgroundWorker.tasks) { - throw new Error(`Background Worker started without tasks`); + if (!this.#backgroundWorker.tasks || this.#backgroundWorker.tasks.length === 0) { + throw new Error( + `Background Worker started without tasks. Searched in: ${__PROJECT_CONFIG__.triggerDirectories?.join( + ", " + )}` + ); } for (const task of this.#backgroundWorker.tasks) { @@ -932,6 +1452,8 @@ class ProdWorker { nextResumeAfter: this.nextResumeAfter, waitForPostStart: this.waitForPostStart, attemptFriendlyId: this.attemptFriendlyId, + waitForTaskReplay: this.waitForTaskReplay, + waitForBatchReplay: this.waitForBatchReplay, }; } diff --git a/packages/cli-v3/src/workers/prod/worker-facade.ts b/packages/cli-v3/src/workers/prod/worker-facade.ts index aa014ef8a..ef2a6a292 100644 --- a/packages/cli-v3/src/workers/prod/worker-facade.ts +++ b/packages/cli-v3/src/workers/prod/worker-facade.ts @@ -292,6 +292,7 @@ async function asyncHeartbeat(initialDelayInSeconds: number = 30, intervalInSeco while (true) { if (_isRunning && _execution) { try { + // The attempt ID will only be used to heartbeat if the message (run) ID isn't set on the TaskRunProcess await zodIpc.send("TASK_HEARTBEAT", { id: _execution.attempt.id }); } catch (err) { console.error("Failed to send HEARTBEAT message", err); diff --git a/packages/core-apps/package.json b/packages/core-apps/package.json index 9b01a0d40..828f5ed23 100644 --- a/packages/core-apps/package.json +++ b/packages/core-apps/package.json @@ -28,7 +28,6 @@ "@trigger.dev/core": "workspace:*", "@trigger.dev/tsconfig": "workspace:*", "@types/node": "18", - "socket.io-client": "^4.7.4", "typescript": "^5.3.0" }, "engines": { diff --git a/apps/coordinator/src/backoff.ts b/packages/core-apps/src/backoff.ts similarity index 50% rename from apps/coordinator/src/backoff.ts rename to packages/core-apps/src/backoff.ts index 25929da3e..d09e6d75f 100644 --- a/apps/coordinator/src/backoff.ts +++ b/packages/core-apps/src/backoff.ts @@ -1,3 +1,5 @@ +import { setTimeout as timeout } from "node:timers/promises"; + type ExponentialBackoffType = "NoJitter" | "FullJitter" | "EqualJitter"; type ExponentialBackoffOptions = { @@ -16,6 +18,26 @@ class StopRetrying extends Error { } } +class AttemptTimeout extends Error { + constructor(message?: string) { + super(message); + this.name = "AttemptTimeout"; + } +} + +class RetryLimitExceeded extends Error { + constructor(message?: string) { + super(message); + this.name = "RetryLimitExceeded"; + } +} + +type YieldType = T extends AsyncGenerator ? Y : never; + +/** + * Exponential backoff helper class + * - All time units in seconds unless otherwise specified + */ export class ExponentialBackoff { #retries: number = 0; @@ -41,64 +63,44 @@ export class ExponentialBackoff { this.#maxElapsed = opts.maxElapsed ?? Infinity; } - #clone() { - return new ExponentialBackoff(this.#type, { - base: this.#base, - factor: this.#factor, - min: this.#min, - max: this.#max, - maxRetries: this.#maxRetries, - maxElapsed: this.#maxElapsed, + #clone(type?: ExponentialBackoffType, opts: Partial = {}) { + return new ExponentialBackoff(type ?? this.#type, { + base: opts.base ?? this.#base, + factor: opts.factor ?? this.#factor, + min: opts.min ?? this.#min, + max: opts.max ?? this.#max, + maxRetries: opts.maxRetries ?? this.#maxRetries, + maxElapsed: opts.maxElapsed ?? this.#maxElapsed, }); } type(type?: ExponentialBackoffType) { - if (typeof type !== "undefined") { - this.#type = type; - } - return this.#clone(); + return this.#clone(type); } base(base?: number) { - if (typeof base !== "undefined") { - this.#base = base; - } - return this.#clone(); + return this.#clone(undefined, { base }); } factor(factor?: number) { - if (typeof factor !== "undefined") { - this.#factor = factor; - } - return this.#clone(); + return this.#clone(undefined, { factor }); } min(min?: number) { - if (typeof min !== "undefined") { - this.#min = min; - } - return this.#clone(); + return this.#clone(undefined, { min }); } max(max?: number) { - if (typeof max !== "undefined") { - this.#max = max; - } - return this.#clone(); + return this.#clone(undefined, { max }); } maxRetries(maxRetries?: number) { - if (typeof maxRetries !== "undefined") { - this.#maxRetries = maxRetries; - } - return this.#clone(); + return this.#clone(undefined, { maxRetries }); } + // TODO: With .execute(), should this also include the time it takes to execute the callback? maxElapsed(maxElapsed?: number) { - if (typeof maxElapsed !== "undefined") { - this.#maxElapsed = maxElapsed; - } - return this.#clone(); + return this.#clone(undefined, { maxElapsed }); } retries(retries?: number) { @@ -145,6 +147,7 @@ export class ExponentialBackoff { yield* this.retryAsync(); } + /** Returns the delay for the current retry in seconds. */ delay(retries: number = this.#retries, jitter: boolean = true) { if (retries > this.#maxRetries) { console.error( @@ -184,13 +187,31 @@ export class ExponentialBackoff { } } - delay = Math.min(delay, this.#max); - delay = Math.max(delay, this.#min); + // If min/max override the delay, jitter with 20% while respecting min/max + if (delay < this.#min) { + delay = this.#min + Math.random() * (this.#min * 0.2); + } + if (delay > this.#max) { + delay = this.#max - Math.random() * (this.#max * 0.2); + } + delay = Math.round(delay); return delay; } + /** Waits with the appropriate delay for the current retry. */ + async wait(retries: number = this.#retries, jitter: boolean = true) { + if (retries > this.#maxRetries) { + console.error(`Retry limit exceeded: ${retries} > ${this.#maxRetries}`); + throw new RetryLimitExceeded(); + } + + const delay = this.delay(retries, jitter); + + return await timeout(delay * 1000); + } + elapsed(retries: number = this.#retries, jitter: boolean = true) { let elapsed = 0; @@ -243,5 +264,98 @@ export class ExponentialBackoff { throw new StopRetrying(); } + get state() { + return { + retries: this.#retries, + type: this.#type, + base: this.#base, + factor: this.#factor, + min: this.#min, + max: this.#max, + maxRetries: this.#maxRetries, + maxElapsed: this.#maxElapsed, + }; + } + + async execute( + callback: ( + iteratorReturn: YieldType> & { + elapsedMs: number; + } + ) => Promise, + { attemptTimeoutMs = 0 }: { attemptTimeoutMs?: number } = {} + ): Promise< + | { success: true; result: T } + | { success: false; error?: unknown; cause: "StopRetrying" | "Timeout" | "MaxRetries" } + > { + let elapsedMs = 0; + let finalError: unknown = undefined; + + for await (const { delay, retry } of this) { + const start = Date.now(); + + if (retry > 0) { + console.log(`Retrying in ${delay.milliseconds}ms`); + await timeout(delay.milliseconds); + } + + let attemptTimeout: NodeJS.Timeout | undefined = undefined; + + try { + const result = await new Promise(async (resolve, reject) => { + if (attemptTimeoutMs > 0) { + attemptTimeout = setTimeout(() => { + reject(new AttemptTimeout()); + }, attemptTimeoutMs); + } + + try { + const callbackResult = await callback({ delay, retry, elapsedMs }); + + resolve(callbackResult); + } catch (error) { + reject(error); + } + }); + + return { + success: true, + result, + }; + } catch (error) { + finalError = error; + + if (error instanceof StopRetrying) { + return { + success: false, + cause: "StopRetrying", + error: error.message, + }; + } + + if (error instanceof AttemptTimeout) { + continue; + } + } finally { + elapsedMs += Date.now() - start; + clearTimeout(attemptTimeout); + } + } + + if (finalError instanceof AttemptTimeout) { + return { + success: false, + cause: "Timeout", + }; + } else { + return { + success: false, + cause: "MaxRetries", + error: finalError, + }; + } + } + + static RetryLimitExceeded = RetryLimitExceeded; static StopRetrying = StopRetrying; } diff --git a/packages/core-apps/src/index.ts b/packages/core-apps/src/index.ts deleted file mode 100644 index a4efbc947..000000000 --- a/packages/core-apps/src/index.ts +++ /dev/null @@ -1,4 +0,0 @@ -export * from "./http"; -export * from "./logger"; -export * from "./provider"; -export * from "./checkpoints"; diff --git a/packages/core/package.json b/packages/core/package.json index c56e133f5..c359c57b9 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -159,7 +159,7 @@ "@opentelemetry/sdk-trace-node": "^1.22.0", "@opentelemetry/semantic-conventions": "^1.22.0", "humanize-duration": "^3.27.3", - "socket.io-client": "4.7.4", + "socket.io-client": "4.7.5", "superjson": "^2.2.1", "ulidx": "^2.2.1", "zod": "3.22.3", diff --git a/packages/core/src/v3/runtime/prodRuntimeManager.ts b/packages/core/src/v3/runtime/prodRuntimeManager.ts index 19dc0fd88..674ab74e4 100644 --- a/packages/core/src/v3/runtime/prodRuntimeManager.ts +++ b/packages/core/src/v3/runtime/prodRuntimeManager.ts @@ -1,14 +1,11 @@ import { clock } from "../clock-api"; -import { logger } from "../logger-api"; import { BatchTaskRunExecutionResult, ProdChildToWorkerMessages, ProdWorkerToChildMessages, TaskRunContext, - TaskRunExecution, TaskRunExecutionResult, } from "../schemas"; -import { checkpointSafeTimeout, unboundedTimeout } from "../utils/timers"; import { ZodIpcConnection } from "../zodIpc"; import { RuntimeManager } from "./manager"; @@ -24,9 +21,7 @@ export class ProdRuntimeManager implements RuntimeManager { { resolve: (value: BatchTaskRunExecutionResult) => void; reject: (err?: any) => void } > = new Map(); - _waitForDuration: - | { resolve: (value: "external") => void; reject: (err?: any) => void } - | undefined; + _waitForDuration: { resolve: (value: void) => void; reject: (err?: any) => void } | undefined; constructor( private ipc: ZodIpcConnection< @@ -43,64 +38,17 @@ export class ProdRuntimeManager implements RuntimeManager { async waitForDuration(ms: number): Promise { const now = Date.now(); - const internalTimeout = unboundedTimeout(ms, "internal" as const); - const checkpointSafeInternalTimeout = checkpointSafeTimeout(ms); - - if (ms <= this.waitThresholdInMs) { - await internalTimeout; - return; - } - - const externalResume = new Promise<"external">((resolve, reject) => { + const resume = new Promise((resolve, reject) => { this._waitForDuration = { resolve, reject }; }); - const { willCheckpointAndRestore } = await this.ipc.sendWithAck( - "WAIT_FOR_DURATION", - { - ms, - now, - }, - 31_000 - ); + await this.ipc.send("WAIT_FOR_DURATION", { + ms, + now, + waitThresholdInMs: this.waitThresholdInMs, + }); - if (!willCheckpointAndRestore) { - await internalTimeout; - return; - } - - this.ipc.send("READY_FOR_CHECKPOINT", {}); - - // internalTimeout acts as a backup and will be accurate if the checkpoint never happens - // checkpointSafeInternalTimeout is accurate even after non-simulated restores - await Promise.race([internalTimeout, checkpointSafeInternalTimeout]); - - // Resets the clock to the current time - clock.reset(); - - try { - // The coordinator should cancel any in-progress checkpoints - const { checkpointCanceled, version } = await this.ipc.sendWithAck( - "CANCEL_CHECKPOINT", - { - version: "v2", - reason: "WAIT_FOR_DURATION", - }, - 31_000 - ); - - if (checkpointCanceled) { - // There won't be a checkpoint or external resume and we've already completed our internal timeout - return; - } - } catch (error) { - // If the cancellation times out, we will proceed as if the checkpoint was canceled - logger.debug("Checkpoint cancellation timed out", { error }); - return; - } - - // No checkpoint was canceled, so we were checkpointed. We need to wait for the external resume message. - await externalResume; + await resume; } resumeAfterDuration(): void { @@ -111,7 +59,7 @@ export class ProdRuntimeManager implements RuntimeManager { // Resets the clock to the current time clock.reset(); - this._waitForDuration.resolve("external"); + this._waitForDuration.resolve(); this._waitForDuration = undefined; } diff --git a/packages/core/src/v3/schemas/messages.ts b/packages/core/src/v3/schemas/messages.ts index aceb1001d..b57eb9f4b 100644 --- a/packages/core/src/v3/schemas/messages.ts +++ b/packages/core/src/v3/schemas/messages.ts @@ -212,47 +212,15 @@ export const ProdChildToWorkerMessages = { id: z.string(), }), }, - TASK_RUN_HEARTBEAT: { - message: z.object({ - version: z.literal("v1").default("v1"), - id: z.string(), - }), - }, READY_TO_DISPOSE: { message: z.undefined(), }, - READY_FOR_CHECKPOINT: { - message: z.object({ - version: z.literal("v1").default("v1"), - }), - }, - CANCEL_CHECKPOINT: { - message: z - .discriminatedUnion("version", [ - z.object({ - version: z.literal("v1"), - }), - z.object({ - version: z.literal("v2"), - reason: WaitReason.optional(), - }), - ]) - .default({ version: "v1" }), - callback: z.object({ - // TODO: Figure out how best to handle callback schema parsing in zod IPC - version: z.literal("v2") /* .default("v2") */, - checkpointCanceled: z.boolean(), - reason: WaitReason.optional(), - }), - }, WAIT_FOR_DURATION: { message: z.object({ version: z.literal("v1").default("v1"), ms: z.number(), now: z.number(), - }), - callback: z.object({ - willCheckpointAndRestore: z.boolean(), + waitThresholdInMs: z.number(), }), }, WAIT_FOR_TASK: { @@ -457,6 +425,7 @@ export const CoordinatorToPlatformMessages = { }), ]), }, + // Deprecated: Only workers without lazy attempt support will use this READY_FOR_EXECUTION: { message: z.object({ version: z.literal("v1").default("v1"), @@ -624,6 +593,12 @@ export const PlatformToCoordinatorMessages = { runId: z.string(), }), }, + DYNAMIC_CONFIG: { + message: z.object({ + version: z.literal("v1").default("v1"), + checkpointThresholdInMs: z.number(), + }), + }, }; export const ClientToSharedQueueMessages = { @@ -672,10 +647,9 @@ const IndexTasksMessage = z.object({ }); export const ProdWorkerToCoordinatorMessages = { - LOG: { + TEST: { message: z.object({ version: z.literal("v1").default("v1"), - text: z.string(), }), callback: z.void(), }, @@ -698,6 +672,7 @@ export const ProdWorkerToCoordinatorMessages = { }), ]), }, + // Deprecated: Only workers without lazy attempt support will use this READY_FOR_EXECUTION: { message: z.object({ version: z.literal("v1").default("v1"), @@ -784,7 +759,7 @@ export const ProdWorkerToCoordinatorMessages = { }, WAIT_FOR_TASK: { message: z.object({ - version: z.literal("v1").default("v1"), + version: z.enum(["v1", "v2"]).default("v1"), friendlyId: z.string(), // This is the attempt that is waiting attemptFriendlyId: z.string(), @@ -795,7 +770,7 @@ export const ProdWorkerToCoordinatorMessages = { }, WAIT_FOR_BATCH: { message: z.object({ - version: z.literal("v1").default("v1"), + version: z.enum(["v1", "v2"]).default("v1"), batchFriendlyId: z.string(), runFriendlyIds: z.string().array(), // This is the attempt that is waiting @@ -843,6 +818,12 @@ export const ProdWorkerToCoordinatorMessages = { }), }), }, + SET_STATE: { + message: z.object({ + version: z.literal("v1").default("v1"), + attemptFriendlyId: z.string().optional(), + }), + }, }; // TODO: The coordinator can only safely use v1 worker messages, higher versions will need a new flag, e.g. SUPPORTS_VERSIONED_MESSAGES @@ -861,6 +842,7 @@ export const CoordinatorToProdWorkerMessages = { attemptId: z.string(), }), }, + // Deprecated: Only workers without lazy attempt support will use this EXECUTE_TASK_RUN: { message: z.object({ version: z.literal("v1").default("v1"), @@ -908,3 +890,7 @@ export const ProdWorkerSocketData = z.object({ deploymentId: z.string(), deploymentVersion: z.string(), }); + +export const CoordinatorSocketData = z.object({ + supportsDynamicConfig: z.string().optional(), +}); diff --git a/packages/core/src/v3/schemas/schemas.ts b/packages/core/src/v3/schemas/schemas.ts index c7a35f271..82b1a542e 100644 --- a/packages/core/src/v3/schemas/schemas.ts +++ b/packages/core/src/v3/schemas/schemas.ts @@ -220,6 +220,7 @@ export type WaitReason = z.infer; export const TaskRunExecutionLazyAttemptPayload = z.object({ runId: z.string(), + attemptCount: z.number().optional(), messageId: z.string(), isTest: z.boolean(), traceContext: z.record(z.unknown()), diff --git a/packages/core/src/v3/zodSocket.ts b/packages/core/src/v3/zodSocket.ts index 964318586..fca10b391 100644 --- a/packages/core/src/v3/zodSocket.ts +++ b/packages/core/src/v3/zodSocket.ts @@ -1,4 +1,4 @@ -import type { Socket } from "socket.io-client"; +import type { ManagerOptions, Socket, SocketOptions } from "socket.io-client"; import { io } from "socket.io-client"; import { ZodError, z } from "zod"; import { EventEmitterLike, ZodMessageValueSchema } from "./zodMessageHandler"; @@ -67,6 +67,7 @@ export type ZodSocketMessageHandlerOptions; + logger?: StructuredLogger; }; type MessageFromSocketSchema< @@ -90,42 +91,69 @@ const messageSchema = z.object({ export class ZodSocketMessageHandler { #schema: TRPCCatalog; #handlers: ZodSocketMessageHandlers | undefined; + #logger: StructuredLogger; constructor(options: ZodSocketMessageHandlerOptions) { this.#schema = options.schema; this.#handlers = options.handlers; + this.#logger = + options.logger ?? new SimpleStructuredLogger("socket-message-handler", LogLevel.info); } public async handleMessage(message: unknown) { - const parsedMessage = this.parseMessage(message); + const parseResult = this.parseMessage(message); + + if (!parseResult.success) { + this.#logger.error("Failed to parse message, skipping handler", { + rawMessage: message, + error: parseResult.reason, + }); + return; + } if (!this.#handlers) { throw new Error("No handlers provided"); } - const handler = this.#handlers[parsedMessage.type]; + const { type, payload } = parseResult.data; + + const handler = this.#handlers[type]; if (!handler) { - console.error(`No handler for message type: ${String(parsedMessage.type)}`); + console.error(`No handler for message type: ${String(type)}`); return; } - const ack = await handler(parsedMessage.payload); + const ack = await handler(payload); return ack; } - public parseMessage(message: unknown): MessagesFromSocketCatalog { + private parseMessage(message: unknown): + | { + success: true; + data: MessagesFromSocketCatalog; + } + | { + success: false; + reason?: string; + } { const parsedMessage = messageSchema.safeParse(message); if (!parsedMessage.success) { - throw new Error(`Failed to parse message: ${JSON.stringify(parsedMessage.error)}`); + return { + success: false, + reason: `Failed to parse message: ${fromZodError(parsedMessage.error).toString()}`, + }; } const schema = this.#schema[parsedMessage.data.type]["message"]; if (!schema) { - throw new Error(`Unknown message type: ${parsedMessage.data.type}`); + return { + success: false, + reason: `Unknown message type: ${parsedMessage.data.type}`, + }; } const messageWithVersion = { @@ -141,14 +169,18 @@ export class ZodSocketMessageHandler = { schema: TMessageCatalog; socket: ZodSocket; + logger?: StructuredLogger; }; export type GetSocketMessagesWithCallback = { @@ -221,10 +254,12 @@ export type GetSocketMessagesWithoutCallback< export class ZodSocketMessageSender { #schema: TMessageCatalog; #socket: ZodSocket; + #logger: StructuredLogger; constructor(options: ZodSocketMessageSenderOptions) { this.#schema = options.schema; this.#socket = options.socket; + this.#logger = options.logger ?? new SimpleStructuredLogger("zod-socket-sender", LogLevel.info); } public send>( @@ -240,7 +275,10 @@ export class ZodSocketMessageSender = Socket< - ZodMessageCatalogToSocketIoEvents, - ZodMessageCatalogToSocketIoEvents ->; +> = Omit< + Socket< + ZodMessageCatalogToSocketIoEvents, + ZodMessageCatalogToSocketIoEvents + >, + "timeout" +> & { + timeout: ( + timeout: number + ) => Socket< + ZodMessageCatalogToSocketIoEvents, + ZodMessageCatalogToSocketIoEvents + >; +}; interface ZodSocketConnectionOptions< TClientMessages extends ZodSocketMessageCatalogSchema, @@ -295,6 +343,7 @@ interface ZodSocketConnectionOptions< }; handlers?: ZodSocketMessageHandlers; authToken?: string; + ioOptions?: Partial; onConnection?: ( socket: ZodSocket, handler: ZodSocketMessageHandler, @@ -340,6 +389,7 @@ export class ZodSocketConnection< extraHeaders: opts.extraHeaders, reconnectionDelay: 500, reconnectionDelayMax: 1000, + ...opts.ioOptions, }); this.#logger = logger.child({ @@ -355,6 +405,7 @@ export class ZodSocketConnection< this.#sender = new ZodSocketMessageSender({ schema: opts.clientMessages, socket: this.socket, + logger: this.#logger, }); this.socket.on("connect_error", async (error) => { @@ -398,7 +449,3 @@ export class ZodSocketConnection< return this.#sender.sendWithAck.bind(this.#sender); } } - -function createLogger(prefix: string) { - return (...args: any[]) => console.log(prefix, ...args); -} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 34aaa4b17..2f83adcda 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -92,9 +92,6 @@ importers: socket.io: specifier: 4.7.4 version: 4.7.4 - socket.io-client: - specifier: 4.7.4 - version: 4.7.4 devDependencies: '@types/node': specifier: ^18 @@ -123,9 +120,6 @@ importers: execa: specifier: ^8.0.1 version: 8.0.1 - socket.io-client: - specifier: ^4.7.4 - version: 4.7.4 devDependencies: '@types/node': specifier: ^18.19.8 @@ -157,9 +151,6 @@ importers: p-queue: specifier: ^8.0.1 version: 8.0.1 - socket.io-client: - specifier: ^4.7.4 - version: 4.7.4 devDependencies: dotenv: specifier: ^16.4.2 @@ -1624,9 +1615,6 @@ importers: simple-git: specifier: ^3.19.0 version: 3.19.0 - socket.io-client: - specifier: ^4.7.4 - version: 4.7.4 source-map-support: specifier: ^0.5.21 version: 0.5.21 @@ -1758,8 +1746,8 @@ importers: specifier: ^3.27.3 version: 3.27.3 socket.io-client: - specifier: 4.7.4 - version: 4.7.4 + specifier: 4.7.5 + version: 4.7.5 superjson: specifier: ^2.2.1 version: 2.2.1 @@ -1828,9 +1816,6 @@ importers: '@types/node': specifier: '18' version: 18.17.1 - socket.io-client: - specifier: ^4.7.4 - version: 4.7.4 typescript: specifier: ^5.3.0 version: 5.3.3 @@ -20459,6 +20444,7 @@ packages: - bufferutil - supports-color - utf-8-validate + dev: false /engine.io-parser@5.2.2(patch_hash=e6nctogrhpxoivwiwy37ersfu4): resolution: {integrity: sha512-RcyUFKA93/CXH20l4SoVvzZfrSDMOTUS3bWVpTt2FuFP+XYrL8i8oonHP7WInRyVHXh0n/ORtoeiE1os+8qkSw==} @@ -32246,8 +32232,8 @@ packages: - utf-8-validate dev: false - /socket.io-client@4.7.4: - resolution: {integrity: sha512-wh+OkeF0rAVCrABWQBaEjLfb7DVPotMbu0cgWgyR0v6eA4EoVnAwcIeIbcdTE3GT/H3kbdLl7OoH2+asoDRIIg==} + /socket.io-client@4.7.5: + resolution: {integrity: sha512-sJ/tqHOCe7Z50JCBCXrsY3I2k03iOiUe+tj1OmKeD2lXPiGH/RUCdTZFoqVyN7l1MnpIzPrGtLcijffmeouNlQ==} engines: {node: '>=10.0.0'} dependencies: '@socket.io/component-emitter': 3.1.0 @@ -32258,6 +32244,7 @@ packages: - bufferutil - supports-color - utf-8-validate + dev: false /socket.io-parser@4.2.4: resolution: {integrity: sha512-/GbIKmo8ioc+NIWIhwdecY0ge+qVBSMdgxGygevmdHj24bsfgtCmcUUcQ5ZzcylGFHsN3k4HB4Cgkl96KVnuew==} @@ -36393,6 +36380,7 @@ packages: /xmlhttprequest-ssl@2.0.0: resolution: {integrity: sha512-QKxVRxiRACQcVuQEYFsI1hhkrMlrXHPegbbd1yn9UHOmRxY+si12nQYzri3vbzt8VdTTRviqcKxcyllFas5z2A==} engines: {node: '>=0.4.0'} + dev: false /xtend@4.0.2: resolution: {integrity: sha512-LKYU1iAXJXUgAXn9URjiu+MWhyUXHsvfp7mcuYm9dSUKK0/CjtrUwFAxD82/mCWbtLsGjFIad0wIsod4zrTAEQ==}