Reset deployed run worker and also get the warm start stuff to work properly

Author: ericallam

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

@@ -250,6 +250,7 @@ let _tracingSDK: TracingSDK | undefined;
 let _executionMeasurement: UsageMeasurement | undefined;
 let _cancelController = new AbortController();
 let _lastFlushPromise: Promise<void> | undefined;
+let _sharedWorkerRuntime: SharedRuntimeManager | undefined;
 
 function resetExecutionEnvironment() {
   _execution = undefined;
@@ -265,7 +266,7 @@ function resetExecutionEnvironment() {
   usageTimeoutManager.reset();
   runMetadataManager.reset();
   waitUntilManager.reset();
-  sharedWorkerRuntime.reset();
+  _sharedWorkerRuntime?.reset();
   durableClock.reset();
   taskContext.disable();
 
@@ -478,6 +479,8 @@ const zodIpc = new ZodIpcConnection({
             });
           }
         } finally {
+          _execution = undefined;
+          _isRunning = false;
           log(`[${new Date().toISOString()}] Task run completed`);
         }
       } catch (err) {
@@ -515,7 +518,7 @@ const zodIpc = new ZodIpcConnection({
       await flushAll(timeoutInMs);
     },
     RESOLVE_WAITPOINT: async ({ waitpoint }) => {
-      sharedWorkerRuntime.resolveWaitpoints([waitpoint]);
+      _sharedWorkerRuntime?.resolveWaitpoints([waitpoint]);
     },
   },
 });
@@ -606,8 +609,8 @@ async function flushMetadata(timeoutInMs: number = 10_000) {
   };
 }
 
-const sharedWorkerRuntime = new SharedRuntimeManager(zodIpc, showInternalLogs);
-runtime.setGlobalRuntimeManager(sharedWorkerRuntime);
+_sharedWorkerRuntime = new SharedRuntimeManager(zodIpc, showInternalLogs);
+runtime.setGlobalRuntimeManager(_sharedWorkerRuntime);
 
 const heartbeatInterval = parseInt(heartbeatIntervalMs ?? "30000", 10);
 

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

@@ -18,6 +18,7 @@ import {
   runMetadata,
   runtime,
   runTimelineMetrics,
+  taskContext,
   TaskRunErrorCodes,
   TaskRunExecution,
   timeout,
@@ -58,6 +59,7 @@ import sourceMapSupport from "source-map-support";
 import { env } from "std-env";
 import { normalizeImportPath } from "../utilities/normalizeImportPath.js";
 import { VERSION } from "../version.js";
+import { promiseWithResolvers } from "@trigger.dev/core/utils";
 
 sourceMapSupport.install({
   handleUncaughtExceptions: false,
@@ -110,15 +112,18 @@ lifecycleHooks.setGlobalLifecycleHooksManager(standardLifecycleHooksManager);
 const standardRunTimelineMetricsManager = new StandardRunTimelineMetricsManager();
 runTimelineMetrics.setGlobalManager(standardRunTimelineMetricsManager);
 
-resourceCatalog.setGlobalResourceCatalog(new StandardResourceCatalog());
+const standardResourceCatalog = new StandardResourceCatalog();
+resourceCatalog.setGlobalResourceCatalog(standardResourceCatalog);
 
 const durableClock = new DurableClock();
 clock.setGlobalClock(durableClock);
+
 const runMetadataManager = new StandardMetadataManager(
   apiClientManager.clientOrThrow(),
   getEnvVar("TRIGGER_STREAM_URL", getEnvVar("TRIGGER_API_URL")) ?? "https://api.trigger.dev"
 );
 runMetadata.setGlobalManager(runMetadataManager);
+
 const waitUntilManager = new StandardWaitUntilManager();
 waitUntil.setGlobalManager(waitUntilManager);
 // Wait for all streams to finish before completing the run
@@ -236,7 +241,30 @@ let _isRunning = false;
 let _isCancelled = false;
 let _tracingSDK: TracingSDK | undefined;
 let _executionMeasurement: UsageMeasurement | undefined;
-const cancelController = new AbortController();
+let _cancelController = new AbortController();
+let _lastFlushPromise: Promise<void> | undefined;
+let _sharedWorkerRuntime: SharedRuntimeManager | undefined;
+
+function resetExecutionEnvironment() {
+  _execution = undefined;
+  _isRunning = false;
+  _isCancelled = false;
+  _executionMeasurement = undefined;
+  _cancelController = new AbortController();
+
+  standardLocalsManager.reset();
+  standardLifecycleHooksManager.reset();
+  standardRunTimelineMetricsManager.reset();
+  usage.reset();
+  timeout.reset();
+  runMetadataManager.reset();
+  waitUntilManager.reset();
+  _sharedWorkerRuntime?.reset();
+  durableClock.reset();
+  taskContext.disable();
+
+  console.log(`[${new Date().toISOString()}] Reset execution environment`);
+}
 
 const zodIpc = new ZodIpcConnection({
   listenSchema: WorkerToExecutorMessageCatalog,
@@ -253,6 +281,22 @@ const zodIpc = new ZodIpcConnection({
         });
       }
 
+      console.log(
+        `[${new Date().toISOString()}] Received EXECUTE_TASK_RUN isWarmStart ${String(isWarmStart)}`
+      );
+
+      if (_lastFlushPromise) {
+        const now = performance.now();
+
+        await _lastFlushPromise;
+
+        const duration = performance.now() - now;
+
+        console.log(`[${new Date().toISOString()}] Awaited last flush in ${duration}ms`);
+      }
+
+      resetExecutionEnvironment();
+
       initializeUsageManager({
         usageIntervalMs: getEnvVar("USAGE_HEARTBEAT_INTERVAL_MS"),
         usageEventUrl: getEnvVar("USAGE_EVENT_URL"),
@@ -421,7 +465,7 @@ const zodIpc = new ZodIpcConnection({
 
           const timeoutController = timeout.abortAfterTimeout(execution.run.maxDuration);
 
-          const signal = AbortSignal.any([cancelController.signal, timeoutController.signal]);
+          const signal = AbortSignal.any([_cancelController.signal, timeoutController.signal]);
 
           const { result } = await executor.execute(execution, metadata, traceContext, signal);
 
@@ -442,6 +486,8 @@ const zodIpc = new ZodIpcConnection({
         } finally {
           _execution = undefined;
           _isRunning = false;
+
+          console.log(`[${new Date().toISOString()}] Task run completed`);
         }
       } catch (err) {
         console.error("Failed to execute task", err);
@@ -467,7 +513,7 @@ const zodIpc = new ZodIpcConnection({
     },
     CANCEL: async ({ timeoutInMs }) => {
       _isCancelled = true;
-      cancelController.abort("run cancelled");
+      _cancelController.abort("run cancelled");
       await callCancelHooks(timeoutInMs);
       if (_executionMeasurement) {
         usage.stop(_executionMeasurement);
@@ -478,7 +524,7 @@ const zodIpc = new ZodIpcConnection({
       await flushAll(timeoutInMs);
     },
     RESOLVE_WAITPOINT: async ({ waitpoint }) => {
-      sharedWorkerRuntime.resolveWaitpoints([waitpoint]);
+      _sharedWorkerRuntime?.resolveWaitpoints([waitpoint]);
     },
   },
 });
@@ -498,6 +544,10 @@ async function callCancelHooks(timeoutInMs: number = 10_000) {
 async function flushAll(timeoutInMs: number = 10_000) {
   const now = performance.now();
 
+  const { promise, resolve } = promiseWithResolvers<void>();
+
+  _lastFlushPromise = promise;
+
   const results = await Promise.allSettled([
     flushUsage(timeoutInMs),
     flushTracingSDK(timeoutInMs),
@@ -530,6 +580,9 @@ async function flushAll(timeoutInMs: number = 10_000) {
   const duration = performance.now() - now;
 
   console.log(`Flushed all in ${duration}ms`);
+
+  // Resolve the last flush promise
+  resolve();
 }
 
 async function flushUsage(timeoutInMs: number = 10_000) {
@@ -597,9 +650,8 @@ function initializeUsageManager({
   timeout.setGlobalManager(new UsageTimeoutManager(devUsageManager));
 }
 
-const sharedWorkerRuntime = new SharedRuntimeManager(zodIpc, true);
-
-runtime.setGlobalRuntimeManager(sharedWorkerRuntime);
+_sharedWorkerRuntime = new SharedRuntimeManager(zodIpc, true);
+runtime.setGlobalRuntimeManager(_sharedWorkerRuntime);
 
 process.title = "trigger-managed-worker";
 

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

@@ -318,7 +318,9 @@ export class ManagedRunController {
           logger: this.logger,
           supervisorSocket: this.socket,
           taskRunProcessProvider: this.taskRunProcessProvider,
-        }).prepareForExecution({
+        });
+
+        await this.currentExecution.prepareForExecution({
           taskRunEnv: previousTaskRunEnv,
         });
       }

packages/cli-v3/src/entryPoints/managed/execution.ts

@@ -128,7 +128,7 @@ export class RunExecution {
    * Prepares the execution with task run environment variables.
    * This should be called before executing, typically after a successful run to prepare for the next one.
    */
-  public prepareForExecution(opts: RunExecutionPrepareOptions): this {
+  public async prepareForExecution(opts: RunExecutionPrepareOptions) {
     if (this.isShuttingDown) {
       throw new Error("prepareForExecution called after execution shut down");
     }
@@ -137,11 +137,10 @@ export class RunExecution {
       throw new Error("prepareForExecution called after process was already created");
     }
 
-    // Store the environment for later use, don't create process yet
-    // The process will be created when needed in executeRun
-    this.currentTaskRunEnv = opts.taskRunEnv;
-
-    return this;
+    this.taskRunProcess = await this.taskRunProcessProvider.getProcess({
+      taskRunEnv: opts.taskRunEnv,
+      isWarmStart: true,
+    });
   }
 
   private attachTaskRunProcessHandlers(taskRunProcess: TaskRunProcess): void {
@@ -183,6 +182,10 @@ export class RunExecution {
    * Returns true if no run has been started yet and we're prepared for the next run.
    */
   get canExecute(): boolean {
+    if (this.taskRunProcessProvider.hasPersistentProcess) {
+      return true;
+    }
+
     // If we've ever had a run ID, this execution can't be reused
     if (this._runFriendlyId) {
       return false;
@@ -359,7 +362,7 @@ export class RunExecution {
       throw new Error("Cannot start attempt: missing run or snapshot manager");
     }
 
-    this.sendDebugLog("starting attempt");
+    this.sendDebugLog("starting attempt", { isWarmStart: String(isWarmStart) });
 
     const attemptStartedAt = Date.now();
 
@@ -404,7 +407,7 @@ export class RunExecution {
       podScheduledAt: this.podScheduledAt?.getTime(),
     });
 
-    this.sendDebugLog("started attempt");
+    this.sendDebugLog("started attempt", { start: start.data });
 
     return { ...start.data, metrics };
   }
@@ -464,7 +467,9 @@ export class RunExecution {
       return;
     }
 
-    const [executeError] = await tryCatch(this.executeRunWrapper(start));
+    const [executeError] = await tryCatch(
+      this.executeRunWrapper({ ...start, isWarmStart: runOpts.isWarmStart })
+    );
 
     if (executeError) {
       this.sendDebugLog("failed to execute run", { error: executeError.message });

packages/cli-v3/src/entryPoints/managed/taskRunProcessProvider.ts

@@ -35,6 +35,10 @@ export class TaskRunProcessProvider {
     this.processKeepAliveMaxExecutionCount = opts.processKeepAliveMaxExecutionCount;
   }
 
+  get hasPersistentProcess(): boolean {
+    return !!this.persistentProcess;
+  }
+
   async handleImmediateRetry(): Promise<void> {
     if (!this.processKeepAliveEnabled) {
       // For immediate retries, we need to ensure we have a clean process

packages/core/src/v3/timeout/api.ts

@@ -7,6 +7,8 @@ class NoopTimeoutManager implements TimeoutManager {
   abortAfterTimeout(timeoutInSeconds?: number): AbortController {
     return new AbortController();
   }
+
+  reset() {}
 }
 
 const NOOP_TIMEOUT_MANAGER = new NoopTimeoutManager();
@@ -40,6 +42,11 @@ export class TimeoutAPI implements TimeoutManager {
     unregisterGlobal(API_NAME);
   }
 
+  public reset() {
+    this.#getManager().reset();
+    this.disable();
+  }
+
   #getManager(): TimeoutManager {
     return getGlobal(API_NAME) ?? NOOP_TIMEOUT_MANAGER;
   }

packages/core/src/v3/timeout/types.ts

@@ -1,6 +1,7 @@
 export interface TimeoutManager {
   abortAfterTimeout: (timeoutInSeconds?: number) => AbortController;
   signal?: AbortSignal;
+  reset: () => void;
 }
 
 export class TaskRunExceededMaxDuration extends Error {

packages/core/src/v3/usage/api.ts

@@ -48,6 +48,11 @@ export class UsageAPI implements UsageManager {
     return this.#getUsageManager().flush();
   }
 
+  public reset() {
+    this.#getUsageManager().reset();
+    this.disable();
+  }
+
   #getUsageManager(): UsageManager {
     return getGlobal(API_NAME) ?? NOOP_USAGE_MANAGER;
   }

packages/core/src/v3/usage/noopUsageManager.ts

@@ -26,4 +26,8 @@ export class NoopUsageManager implements UsageManager {
   sample(): UsageSample | undefined {
     return undefined;
   }
+
+  reset(): void {
+    // Noop
+  }
 }

packages/core/src/v3/usage/prodUsageManager.ts

@@ -27,6 +27,14 @@ export class ProdUsageManager implements UsageManager {
     return typeof this._usageClient !== "undefined";
   }
 
+  reset(): void {
+    this.delegageUsageManager.reset();
+    this._abortController = new AbortController();
+    this._usageClient = undefined;
+    this._measurement = undefined;
+    this._lastSample = undefined;
+  }
+
   disable(): void {
     this.delegageUsageManager.disable();
     this._abortController?.abort();