Reset execution environment between runs

Author: ericallam

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

@@ -19,6 +19,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,
@@ -112,8 +114,12 @@ runTimelineMetrics.setGlobalManager(standardRunTimelineMetricsManager);
 
 const devUsageManager = new DevUsageManager();
 usage.setGlobalUsageManager(devUsageManager);
-timeout.setGlobalManager(new UsageTimeoutManager(devUsageManager));
-resourceCatalog.setGlobalResourceCatalog(new StandardResourceCatalog());
+
+const usageTimeoutManager = new UsageTimeoutManager(devUsageManager);
+timeout.setGlobalManager(usageTimeoutManager);
+
+const standardResourceCatalog = new StandardResourceCatalog();
+resourceCatalog.setGlobalResourceCatalog(standardResourceCatalog);
 
 const durableClock = new DurableClock();
 clock.setGlobalClock(durableClock);
@@ -238,7 +244,29 @@ 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;
+
+function resetExecutionEnvironment() {
+  _execution = undefined;
+  _isRunning = false;
+  _isCancelled = false;
+  _executionMeasurement = undefined;
+  _cancelController = new AbortController();
+
+  standardLocalsManager.reset();
+  standardLifecycleHooksManager.reset();
+  standardRunTimelineMetricsManager.reset();
+  devUsageManager.reset();
+  usageTimeoutManager.reset();
+  runMetadataManager.reset();
+  waitUntilManager.reset();
+  sharedWorkerRuntime.reset();
+  durableClock.reset();
+  taskContext.disable();
+
+  log(`[${new Date().toISOString()}] Reset execution environment`);
+}
 
 const zodIpc = new ZodIpcConnection({
   listenSchema: WorkerToExecutorMessageCatalog,
@@ -254,6 +282,18 @@ const zodIpc = new ZodIpcConnection({
 
       log(`[${new Date().toISOString()}] Received EXECUTE_TASK_RUN`, execution);
 
+      if (_lastFlushPromise) {
+        const now = performance.now();
+
+        await _lastFlushPromise;
+
+        const duration = performance.now() - now;
+
+        log(`[${new Date().toISOString()}] Awaited last flush in ${duration}ms`);
+      }
+
+      resetExecutionEnvironment();
+
       standardRunTimelineMetricsManager.registerMetricsFromExecution(metrics);
 
       if (_isRunning) {
@@ -365,8 +405,6 @@ const zodIpc = new ZodIpcConnection({
           return;
         }
 
-        process.title = `trigger-dev-worker: ${execution.task.id} ${execution.run.id}`;
-
         // Import the task module
         const task = resourceCatalog.getTask(execution.task.id);
 
@@ -413,7 +451,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);
 
@@ -432,8 +470,7 @@ const zodIpc = new ZodIpcConnection({
             });
           }
         } finally {
-          _execution = undefined;
-          _isRunning = false;
+          log(`[${new Date().toISOString()}] Task run completed`);
         }
       } catch (err) {
         logError("Failed to execute task", err);
@@ -459,7 +496,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);
@@ -490,6 +527,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([
     flushTracingSDK(timeoutInMs),
     flushMetadata(timeoutInMs),
@@ -522,6 +563,9 @@ async function flushAll(timeoutInMs: number = 10_000) {
   const duration = performance.now() - now;
 
   log(`Flushed all in ${duration}ms`);
+
+  // Resolve the last flush promise
+  resolve();
 }
 
 async function flushTracingSDK(timeoutInMs: number = 10_000) {

packages/core/src/v3/lifecycleHooks/manager.ts

@@ -68,6 +68,23 @@ export class StandardLifecycleHooksManager implements LifecycleHooksManager {
   private taskCancelHooks: Map<string, RegisteredHookFunction<AnyOnCancelHookFunction>> = new Map();
   private onCancelHookListeners: (() => Promise<void>)[] = [];
 
+  reset(): void {
+    this.taskInitHooks.clear();
+    this.taskStartHooks.clear();
+    this.taskFailureHooks.clear();
+    this.taskSuccessHooks.clear();
+    this.taskCompleteHooks.clear();
+    this.taskResumeHooks.clear();
+    this.taskCatchErrorHooks.clear();
+    this.taskMiddlewareHooks.clear();
+    this.taskCleanupHooks.clear();
+    this.taskWaitHooks.clear();
+    this.taskCancelHooks.clear();
+    this.onCancelHookListeners = [];
+    this.onWaitHookListeners = [];
+    this.onResumeHookListeners = [];
+  }
+
   registerOnCancelHookListener(listener: () => Promise<void>): void {
     this.onCancelHookListeners.push(listener);
   }

packages/core/src/v3/locals/manager.ts

@@ -33,5 +33,8 @@ export class StandardLocalsManager implements LocalsManager {
   setLocal<T>(key: LocalsKey<T>, value: T): void {
     this.store.set(key.__type, value);
   }
+
+  reset(): void {
+    this.store.clear();
+  }
 }
-0;

packages/core/src/v3/runMetadata/manager.ts

@@ -30,6 +30,17 @@ export class StandardMetadataManager implements RunMetadataManager {
     private streamsVersion: "v1" | "v2" = "v1"
   ) {}
 
+  reset(): void {
+    this.queuedOperations.clear();
+    this.queuedParentOperations.clear();
+    this.queuedRootOperations.clear();
+    this.activeStreams.clear();
+    this.store = undefined;
+    this.runId = undefined;
+    this.flushTimeoutId = null;
+    this.isFlushing = false;
+  }
+
   get parent(): RunMetadataUpdater {
     // Store a reference to 'this' to ensure proper context
     const self = this;

packages/core/src/v3/runTimelineMetrics/runTimelineMetricsManager.ts

@@ -30,6 +30,11 @@ export class StandardRunTimelineMetricsManager implements RunTimelineMetricsMana
     }
   }
 
+  reset(): void {
+    this._metrics = [];
+  }
+
+  // TODO: handle this when processKeepAlive is enabled
   #seedMetricsFromEnvironment() {
     const forkStartTime = getEnvVar("TRIGGER_PROCESS_FORK_START_TIME");
     const warmStart = getEnvVar("TRIGGER_WARM_START");

packages/core/src/v3/runtime/sharedRuntimeManager.ts

@@ -41,6 +41,11 @@ export class SharedRuntimeManager implements RuntimeManager {
     }, 300_000);
   }
 
+  reset(): void {
+    this.resolversById.clear();
+    this.waitpointsByResolverId.clear();
+  }
+
   disable(): void {
     // do nothing
   }

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

@@ -14,6 +14,12 @@ export class UsageTimeoutManager implements TimeoutManager {
     return this._abortSignal;
   }
 
+  reset(): void {
+    this._abortController = new AbortController();
+    this._abortSignal = undefined;
+    this._intervalId = undefined;
+  }
+
   abortAfterTimeout(timeoutInSeconds?: number): AbortController {
     this._abortSignal = this._abortController.signal;
 

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

@@ -50,6 +50,12 @@ export class DevUsageManager implements UsageManager {
 
   async flush(): Promise<void> {}
 
+  reset(): void {
+    this._firstMeasurement = undefined;
+    this._currentMeasurements.clear();
+    this._pauses.clear();
+  }
+
   sample(): UsageSample | undefined {
     return this._firstMeasurement?.sample();
   }

packages/core/src/v3/waitUntil/manager.ts

@@ -3,6 +3,10 @@ import { MaybeDeferredPromise, WaitUntilManager } from "./types.js";
 export class StandardWaitUntilManager implements WaitUntilManager {
   private maybeDeferredPromises: Set<MaybeDeferredPromise> = new Set();
 
+  reset(): void {
+    this.maybeDeferredPromises.clear();
+  }
+
   register(promise: MaybeDeferredPromise): void {
     this.maybeDeferredPromises.add(promise);
   }