v4: eagerly fork child process before warm start

Also fixes an issue where the attempt span events weren't coming through in the partial spans

Author: ericallam

packages/cli-v3/src/entryPoints/dev-run-controller.ts

@@ -622,6 +622,15 @@ export class DevRunController {
         version: this.opts.worker.serverWorker?.version,
         engine: "V2",
       },
+      machine: execution.machine,
+    }).initialize();
+
+    logger.debug("executing task run process", {
+      attemptId: execution.attempt.id,
+      runId: execution.run.id,
+    });
+
+    const completion = await this.taskRunProcess.execute({
       payload: {
         execution,
         traceContext: execution.run.traceContext ?? {},
@@ -630,15 +639,6 @@ export class DevRunController {
       messageId: run.friendlyId,
     });
 
-    await this.taskRunProcess.initialize();
-
-    logger.debug("executing task run process", {
-      attemptId: execution.attempt.id,
-      runId: execution.run.id,
-    });
-
-    const completion = await this.taskRunProcess.execute();
-
     logger.debug("Completed run", completion);
 
     try {

packages/cli-v3/src/entryPoints/dev-run-worker.ts

@@ -36,6 +36,7 @@ import {
   logLevels,
   ManagedRuntimeManager,
   OtelTaskLogger,
+  populateEnv,
   StandardLifecycleHooksManager,
   StandardLocalsManager,
   StandardMetadataManager,
@@ -238,7 +239,13 @@ const zodIpc = new ZodIpcConnection({
   emitSchema: ExecutorToWorkerMessageCatalog,
   process,
   handlers: {
-    EXECUTE_TASK_RUN: async ({ execution, traceContext, metadata, metrics }, sender) => {
+    EXECUTE_TASK_RUN: async ({ execution, traceContext, metadata, metrics, env }, sender) => {
+      if (env) {
+        populateEnv(env, {
+          override: true,
+        });
+      }
+
       log(`[${new Date().toISOString()}] Received EXECUTE_TASK_RUN`, execution);
 
       standardRunTimelineMetricsManager.registerMetricsFromExecution(metrics);

packages/cli-v3/src/entryPoints/managed-run-controller.ts

@@ -836,6 +836,18 @@ class ManagedRunController {
         this.exitProcess(this.successExitCode);
       }
 
+      if (this.taskRunProcess) {
+        logger.debug("waitForNextRun: eagerly recreating task run process with options");
+        this.taskRunProcess = new TaskRunProcess({
+          ...this.taskRunProcess.options,
+          isWarmStart: true,
+        }).initialize();
+      } else {
+        logger.debug(
+          "waitForNextRun: no existing task run process, so we can't eagerly recreate it"
+        );
+      }
+
       // Check the service is up and get additional warm start config
       const connect = await this.warmStartClient.connect();
 
@@ -904,6 +916,9 @@ class ManagedRunController {
 
   private exitProcess(code?: number): never {
     logger.log("Exiting process", { code });
+    if (this.taskRunProcess?.isPreparedForNextRun) {
+      this.taskRunProcess.forceExit();
+    }
     process.exit(code);
   }
 
@@ -980,30 +995,33 @@ class ManagedRunController {
   }: WorkloadRunAttemptStartResponseBody) {
     this.snapshotPoller.start();
 
-    this.taskRunProcess = new TaskRunProcess({
-      workerManifest: this.workerManifest,
-      env: envVars,
-      serverWorker: {
-        id: "unmanaged",
-        contentHash: env.TRIGGER_CONTENT_HASH,
-        version: env.TRIGGER_DEPLOYMENT_VERSION,
-        engine: "V2",
-      },
-      payload: {
-        execution,
-        traceContext: execution.run.traceContext ?? {},
-      },
-      messageId: run.friendlyId,
-    });
-
-    await this.taskRunProcess.initialize();
+    if (!this.taskRunProcess || !this.taskRunProcess.isPreparedForNextRun) {
+      this.taskRunProcess = new TaskRunProcess({
+        workerManifest: this.workerManifest,
+        env: envVars,
+        serverWorker: {
+          id: "unmanaged",
+          contentHash: env.TRIGGER_CONTENT_HASH,
+          version: env.TRIGGER_DEPLOYMENT_VERSION,
+          engine: "V2",
+        },
+        machine: execution.machine,
+      }).initialize();
+    }
 
     logger.log("executing task run process", {
       attemptId: execution.attempt.id,
       runId: execution.run.id,
     });
 
-    const completion = await this.taskRunProcess.execute();
+    const completion = await this.taskRunProcess.execute({
+      payload: {
+        execution,
+        traceContext: execution.run.traceContext ?? {},
+      },
+      messageId: run.friendlyId,
+      env: envVars,
+    });
 
     logger.log("Completed run", completion);
 

packages/cli-v3/src/entryPoints/managed-run-worker.ts

@@ -36,6 +36,7 @@ import {
   logLevels,
   ManagedRuntimeManager,
   OtelTaskLogger,
+  populateEnv,
   ProdUsageManager,
   StandardLifecycleHooksManager,
   StandardLocalsManager,
@@ -248,7 +249,16 @@ const zodIpc = new ZodIpcConnection({
   emitSchema: ExecutorToWorkerMessageCatalog,
   process,
   handlers: {
-    EXECUTE_TASK_RUN: async ({ execution, traceContext, metadata, metrics }, sender) => {
+    EXECUTE_TASK_RUN: async (
+      { execution, traceContext, metadata, metrics, env, isWarmStart },
+      sender
+    ) => {
+      if (env) {
+        populateEnv(env, {
+          override: true,
+        });
+      }
+
       standardRunTimelineMetricsManager.registerMetricsFromExecution(metrics);
 
       console.log(`[${new Date().toISOString()}] Received EXECUTE_TASK_RUN`, execution);
@@ -383,6 +393,7 @@ const zodIpc = new ZodIpcConnection({
           tracingSDK,
           consoleInterceptor,
           retries: config.retries,
+          isWarmStart,
         });
 
         try {

packages/cli-v3/src/executions/taskRunProcess.ts

@@ -1,6 +1,7 @@
 import {
   CompletedWaitpoint,
   ExecutorToWorkerMessageCatalog,
+  MachinePreset,
   ServerBackgroundWorker,
   TaskRunErrorCodes,
   TaskRunExecution,
@@ -48,10 +49,15 @@ export type TaskRunProcessOptions = {
   workerManifest: WorkerManifest;
   serverWorker: ServerBackgroundWorker;
   env: Record<string, string>;
+  machine: MachinePreset;
+  isWarmStart?: boolean;
+  cwd?: string;
+};
+
+export type TaskRunProcessExecuteParams = {
   payload: TaskRunExecutionPayload;
   messageId: string;
-
-  cwd?: string;
+  env?: Record<string, string>;
 };
 
 export class TaskRunProcess {
@@ -79,9 +85,18 @@ export class TaskRunProcess {
   public onWaitForBatch: Evt<OnWaitForBatchMessage> = new Evt();
   public onWait: Evt<OnWaitMessage> = new Evt();
 
-  constructor(public readonly options: TaskRunProcessOptions) {}
+  private _isPreparedForNextRun: boolean = false;
+
+  constructor(public readonly options: TaskRunProcessOptions) {
+    this._isPreparedForNextRun = true;
+  }
+
+  get isPreparedForNextRun() {
+    return this._isPreparedForNextRun;
+  }
 
   async cancel() {
+    this._isPreparedForNextRun = false;
     this._isBeingCancelled = true;
 
     try {
@@ -94,6 +109,8 @@ export class TaskRunProcess {
   }
 
   async cleanup(kill = true) {
+    this._isPreparedForNextRun = false;
+
     try {
       await this.#flush();
     } catch (err) {
@@ -105,25 +122,12 @@ export class TaskRunProcess {
     }
   }
 
-  get runId() {
-    return this.options.payload.execution.run.id;
-  }
+  initialize() {
+    const { env: $env, workerManifest, cwd, machine } = this.options;
 
-  get isTest() {
-    return this.options.payload.execution.run.isTest;
-  }
-
-  get payload(): TaskRunExecutionPayload {
-    return this.options.payload;
-  }
-
-  async initialize() {
-    const { env: $env, workerManifest, cwd, messageId, payload } = this.options;
-
-    const maxOldSpaceSize = nodeOptionsWithMaxOldSpaceSize(undefined, payload.execution.machine);
+    const maxOldSpaceSize = nodeOptionsWithMaxOldSpaceSize(undefined, machine);
 
     const fullEnv = {
-      ...(this.isTest ? { TRIGGER_LOG_LEVEL: "debug" } : {}),
       ...$env,
       OTEL_IMPORT_HOOK_INCLUDES: workerManifest.otelImportHook?.include?.join(","),
       // TODO: this will probably need to use something different for bun (maybe --preload?)
@@ -132,7 +136,7 @@ export class TaskRunProcess {
       TRIGGER_PROCESS_FORK_START_TIME: String(Date.now()),
     };
 
-    logger.debug(`[${this.runId}] initializing task run process`, {
+    logger.debug(`initializing task run process`, {
       env: fullEnv,
       path: workerManifest.workerEntryPoint,
       cwd,
@@ -175,13 +179,13 @@ export class TaskRunProcess {
 
           resolver(result);
         },
-        READY_TO_DISPOSE: async (message) => {
-          logger.debug(`[${this.runId}] task run process is ready to dispose`);
+        READY_TO_DISPOSE: async () => {
+          logger.debug(`task run process is ready to dispose`);
 
           this.onReadyToDispose.post(this);
         },
         TASK_HEARTBEAT: async (message) => {
-          this.onTaskRunHeartbeat.post(messageId);
+          this.onTaskRunHeartbeat.post(message.id);
         },
         WAIT_FOR_TASK: async (message) => {
           this.onWaitForTask.post(message);
@@ -190,14 +194,16 @@ export class TaskRunProcess {
           this.onWaitForBatch.post(message);
         },
         UNCAUGHT_EXCEPTION: async (message) => {
-          logger.debug(`[${this.runId}] uncaught exception in task run process`, { ...message });
+          logger.debug("uncaught exception in task run process", { ...message });
         },
       },
     });
 
     this._child.on("exit", this.#handleExit.bind(this));
     this._child.stdout?.on("data", this.#handleLog.bind(this));
     this._child.stderr?.on("data", this.#handleStdErr.bind(this));
+
+    return this;
   }
 
   async #flush(timeoutInMs: number = 5_000) {
@@ -206,7 +212,9 @@ export class TaskRunProcess {
     await this._ipc?.sendWithAck("FLUSH", { timeoutInMs }, timeoutInMs + 1_000);
   }
 
-  async execute(): Promise<TaskRunExecutionResult> {
+  async execute(params: TaskRunProcessExecuteParams): Promise<TaskRunExecutionResult> {
+    this._isPreparedForNextRun = false;
+
     let resolver: (value: TaskRunExecutionResult) => void;
     let rejecter: (err?: any) => void;
 
@@ -215,19 +223,19 @@ export class TaskRunProcess {
       rejecter = reject;
     });
 
-    this._attemptStatuses.set(this.payload.execution.attempt.id, "PENDING");
+    this._attemptStatuses.set(params.payload.execution.attempt.id, "PENDING");
 
     // @ts-expect-error - We know that the resolver and rejecter are defined
-    this._attemptPromises.set(this.payload.execution.attempt.id, { resolver, rejecter });
+    this._attemptPromises.set(params.payload.execution.attempt.id, { resolver, rejecter });
 
-    const { execution, traceContext, metrics } = this.payload;
+    const { execution, traceContext, metrics } = params.payload;
 
     this._currentExecution = execution;
 
     if (this._child?.connected && !this._isBeingKilled && !this._child.killed) {
       logger.debug(
         `[${new Date().toISOString()}][${
-          this.runId
+          params.payload.execution.run.id
         }] sending EXECUTE_TASK_RUN message to task run process`,
         {
           pid: this.pid,
@@ -239,6 +247,8 @@ export class TaskRunProcess {
         traceContext,
         metadata: this.options.serverWorker,
         metrics,
+        env: params.env,
+        isWarmStart: this.options.isWarmStart,
       });
     }
 
@@ -398,7 +408,7 @@ export class TaskRunProcess {
   }
 
   async kill(signal?: number | NodeJS.Signals, timeoutInMs?: number) {
-    logger.debug(`[${this.runId}] killing task run process`, {
+    logger.debug(`killing task run process`, {
       signal,
       timeoutInMs,
       pid: this.pid,
@@ -417,6 +427,16 @@ export class TaskRunProcess {
     }
   }
 
+  forceExit() {
+    try {
+      this._isBeingKilled = true;
+
+      this._child?.kill("SIGKILL");
+    } catch (error) {
+      logger.debug("forceExit: failed to kill child process", { error });
+    }
+  }
+
   get isBeingKilled() {
     return this._isBeingKilled || this._child?.killed;
   }

packages/core/src/v3/schemas/messages.ts

@@ -215,6 +215,8 @@ export const WorkerToExecutorMessageCatalog = {
       traceContext: z.record(z.unknown()),
       metadata: ServerBackgroundWorker,
       metrics: TaskRunExecutionMetrics.optional(),
+      env: z.record(z.string()).optional(),
+      isWarmStart: z.boolean().optional(),
     }),
   },
   TASK_RUN_COMPLETED_NOTIFICATION: {

packages/core/src/v3/semanticInternalAttributes.ts

@@ -23,6 +23,7 @@ export const SemanticInternalAttributes = {
   MACHINE_PRESET_CPU: "ctx.machine.cpu",
   MACHINE_PRESET_MEMORY: "ctx.machine.memory",
   MACHINE_PRESET_CENTS_PER_MS: "ctx.machine.centsPerMs",
+  SKIP_SPAN_PARTIAL: "$span.skip_partial",
   SPAN_PARTIAL: "$span.partial",
   SPAN_ID: "$span.span_id",
   ENTITY_TYPE: "$entity.type",
@@ -58,4 +59,5 @@ export const SemanticInternalAttributes = {
   SPAN_ATTEMPT: "$span.attempt",
   METRIC_EVENTS: "$metrics.events",
   EXECUTION_ENVIRONMENT: "exec_env",
+  WARM_START: "warm_start",
 };