diff --git a/packages/cli-v3/src/workers/common/errors.ts b/packages/cli-v3/src/workers/common/errors.ts index 4ba1a8ef3..053ab8d19 100644 --- a/packages/cli-v3/src/workers/common/errors.ts +++ b/packages/cli-v3/src/workers/common/errors.ts @@ -53,3 +53,11 @@ export class SigKillTimeoutProcessError extends Error { this.name = "SigKillTimeoutProcessError"; } } + +export class GracefulExitTimeoutError extends Error { + constructor() { + super("Graceful exit timeout"); + + this.name = "GracefulExitTimeoutError"; + } +} diff --git a/packages/cli-v3/src/workers/prod/backgroundWorker.ts b/packages/cli-v3/src/workers/prod/backgroundWorker.ts index 24bd528fd..5431d9d10 100644 --- a/packages/cli-v3/src/workers/prod/backgroundWorker.ts +++ b/packages/cli-v3/src/workers/prod/backgroundWorker.ts @@ -24,6 +24,7 @@ import { ChildProcess, fork } from "node:child_process"; import { CancelledProcessError, CleanupProcessError, + GracefulExitTimeoutError, SigKillTimeoutProcessError, TaskMetadataParseError, UncaughtExceptionError, @@ -88,7 +89,7 @@ export class ProdBackgroundWorker { private params: BackgroundWorkerParams ) {} - async close() { + async close(gracefulExitTimeoutElapsed = false) { if (this._closed) { return; } @@ -99,7 +100,7 @@ export class ProdBackgroundWorker { this.onTaskRunHeartbeat.detach(); // We need to close the task run process - await this._taskRunProcess?.cleanup(true); + await this._taskRunProcess?.cleanup(true, gracefulExitTimeoutElapsed); } async killTaskRunProcess(flush = true, initialSignal: number | NodeJS.Signals = "SIGTERM") { @@ -113,16 +114,10 @@ export class ProdBackgroundWorker { const currentTaskRunProcess = this._taskRunProcess; - try { - const initialExit = currentTaskRunProcess.onExit.waitFor(5_000); - currentTaskRunProcess.kill(initialSignal); - await initialExit; - } catch (error) { - // Try again with SIGKILL - const forcedExit = currentTaskRunProcess.onExit.waitFor(5_000); - currentTaskRunProcess.kill("SIGKILL"); - await forcedExit; - } + // Try graceful exit but don't wait. We limit the amount of processes during creation instead. + this.#tryGracefulExit(currentTaskRunProcess, true, initialSignal).catch((error) => { + console.error("Error while trying graceful exit", error); + }); this._closed = true; } @@ -329,42 +324,14 @@ export class ProdBackgroundWorker { if (currentTaskRunProcess.isBeingKilled) { if (this._taskRunProcessesBeingKilled.size > 1) { - // If there's more than one being killed, wait for graceful exit - try { - await currentTaskRunProcess.onExit.waitFor(5_000); - } catch (error) { - console.error("TaskRunProcess graceful kill timeout exceeded", error); - - try { - const forcedKill = currentTaskRunProcess.onExit.waitFor(5_000); - currentTaskRunProcess.kill("SIGKILL"); - await forcedKill; - } catch (error) { - console.error("TaskRunProcess forced kill timeout exceeded", error); - throw new SigKillTimeoutProcessError(); - } - } + await this.#tryGracefulExit(currentTaskRunProcess); } else { // If there's only one or none being killed, don't do anything so we can create a fresh one in parallel } } else { // It's not being killed, so kill it if (this._taskRunProcessesBeingKilled.size > 0) { - // If there's one being killed already, wait for graceful exit - try { - await currentTaskRunProcess.onExit.waitFor(5_000); - } catch (error) { - console.error("TaskRunProcess graceful kill timeout exceeded", error); - - try { - const forcedKill = currentTaskRunProcess.onExit.waitFor(5_000); - currentTaskRunProcess.kill("SIGKILL"); - await forcedKill; - } catch (error) { - console.error("TaskRunProcess forced kill timeout exceeded", error); - throw new SigKillTimeoutProcessError(); - } - } + await this.#tryGracefulExit(currentTaskRunProcess); } else { // There's none being killed yet, so we can kill it without waiting. We still set a timeout to kill it forcefully just in case it sticks around. currentTaskRunProcess.kill("SIGTERM", 5_000).catch(() => {}); @@ -372,6 +339,37 @@ export class ProdBackgroundWorker { } } + async #tryGracefulExit( + taskRunProcess: TaskRunProcess, + kill = false, + initialSignal: number | NodeJS.Signals = "SIGTERM" + ) { + try { + const initialExit = taskRunProcess.onExit.waitFor(5_000); + + if (kill) { + taskRunProcess.kill(initialSignal); + } + + await initialExit; + } catch (error) { + console.error("TaskRunProcess graceful kill timeout exceeded", error); + + this.#tryForcefulExit(taskRunProcess); + } + } + + async #tryForcefulExit(taskRunProcess: TaskRunProcess) { + try { + const forcedKill = taskRunProcess.onExit.waitFor(5_000); + taskRunProcess.kill("SIGKILL"); + await forcedKill; + } catch (error) { + console.error("TaskRunProcess forced kill timeout exceeded", error); + throw new SigKillTimeoutProcessError(); + } + } + // We need to fork the process before we can execute any tasks, use a fresh process for each execution async executeTaskRun( payload: ProdTaskRunExecutionPayload, @@ -447,6 +445,19 @@ export class ProdBackgroundWorker { }; } + if (e instanceof GracefulExitTimeoutError) { + return { + id: payload.execution.attempt.id, + ok: false, + retry: undefined, + error: { + type: "INTERNAL_ERROR", + code: TaskRunErrorCodes.GRACEFUL_EXIT_TIMEOUT, + message: "Worker process killed while attempt in progress.", + }, + }; + } + return { id: payload.execution.attempt.id, ok: false, @@ -524,6 +535,7 @@ class TaskRunProcess { private _currentExecution: TaskRunExecution | undefined; private _isBeingKilled: boolean = false; private _isBeingCancelled: boolean = false; + private _gracefulExitTimeoutElapsed: boolean = false; /** * @deprecated use onTaskRunHeartbeat instead @@ -682,7 +694,7 @@ class TaskRunProcess { await this.cleanup(true); } - async cleanup(kill: boolean = false) { + async cleanup(kill = false, gracefulExitTimeoutElapsed = false) { if (kill && this._isBeingKilled) { return; } @@ -692,10 +704,21 @@ class TaskRunProcess { this.onIsBeingKilled.post(this); } + const killChildProcess = gracefulExitTimeoutElapsed && !!this._currentExecution; + + // Kill parent unless graceful exit timeout has elapsed and we're in the middle of an execution + const killParentProcess = kill && !killChildProcess; + await this._ipc?.sendWithAck("CLEANUP", { flush: true, - kill, + kill: killParentProcess, }); + + if (killChildProcess) { + this._gracefulExitTimeoutElapsed = true; + // Kill the child process + await this.kill("SIGKILL"); + } } async executeTaskRun(payload: TaskRunExecutionPayload): Promise { @@ -766,6 +789,9 @@ class TaskRunProcess { if (this._isBeingCancelled) { rejecter(new CancelledProcessError()); + } else if (this._gracefulExitTimeoutElapsed) { + // Order matters, this has to be before the graceful exit timeout + rejecter(new GracefulExitTimeoutError()); } else if (this._isBeingKilled) { rejecter(new CleanupProcessError()); } else { diff --git a/packages/cli-v3/src/workers/prod/entry-point.ts b/packages/cli-v3/src/workers/prod/entry-point.ts index 542a53530..517dbe4b4 100644 --- a/packages/cli-v3/src/workers/prod/entry-point.ts +++ b/packages/cli-v3/src/workers/prod/entry-point.ts @@ -71,6 +71,8 @@ class ProdWorker { logger.log("Received signal", { signal }); if (signal === "SIGTERM") { + let gracefulExitTimeoutElapsed = false; + if (this.executing) { const terminationGracePeriodSeconds = 60 * 60; @@ -80,21 +82,27 @@ class ProdWorker { // Wait for termination grace period minus 5s to give cleanup a chance to complete await setTimeout(terminationGracePeriodSeconds * 1000 - 5000); + gracefulExitTimeoutElapsed = true; logger.log("Termination timeout reached, exiting gracefully."); } else { logger.log("Not executing, exiting immediately."); } - await this.#exitGracefully(); + await this.#exitGracefully(gracefulExitTimeoutElapsed); + return; } logger.log("Unhandled signal", { signal }); } - async #exitGracefully() { - await this.#backgroundWorker.close(); - process.exit(0); + async #exitGracefully(gracefulExitTimeoutElapsed = false) { + await this.#backgroundWorker.close(gracefulExitTimeoutElapsed); + + if (!gracefulExitTimeoutElapsed) { + // TODO: Maybe add a sensible timeout instead of a conditional to avoid zombies + process.exit(0); + } } async #reconnect(isPostStart = false, reconnectImmediately = false) { diff --git a/packages/cli-v3/src/workers/prod/worker-facade.ts b/packages/cli-v3/src/workers/prod/worker-facade.ts index d4cd7a91c..e16496d3a 100644 --- a/packages/cli-v3/src/workers/prod/worker-facade.ts +++ b/packages/cli-v3/src/workers/prod/worker-facade.ts @@ -179,23 +179,6 @@ const zodIpc = new ZodIpcConnection({ CLEANUP: async ({ flush, kill }, sender) => { if (kill) { await tracingSDK.flush(); - - if (_execution) { - // Fail currently executing attempt - await sender.send("TASK_RUN_COMPLETED", { - execution: _execution, - result: { - ok: false, - id: _execution.run.id, - error: { - type: "INTERNAL_ERROR", - code: TaskRunErrorCodes.GRACEFUL_EXIT_TIMEOUT, - message: "Worker process killed while attempt in progress.", - }, - }, - }); - } - // Now we need to exit the process await sender.send("READY_TO_DISPOSE", undefined); } else {