Added parallel wait prevention, it’s working for duration waits but not well for triggerAndWait yet

Author: matt-aitken

internal-packages/run-engine/src/engine/index.ts

@@ -2857,7 +2857,7 @@ export class RunEngine {
         const retriableError = shouldRetryError(taskRunErrorEnhancer(completion.error));
 
         const permanentlyFailRun = async (run?: { status: TaskRunStatus; spanId: string }) => {
-          // Emit and event so we can complete any spans of stalled executions
+          // Emit an event so we can complete any spans of stalled executions
           if (forceRequeue && run) {
             this.eventBus.emit("runAttemptFailed", {
               time: failedAt,

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

@@ -10,6 +10,7 @@ import {
 } from "../schemas/index.js";
 import { ExecutorToWorkerProcessConnection } from "../zodIpc.js";
 import { RuntimeManager } from "./manager.js";
+import { preventMultipleWaits } from "./preventMultipleWaits.js";
 
 type Resolver = (value: CompletedWaitpoint) => void;
 
@@ -19,6 +20,8 @@ export class ManagedRuntimeManager implements RuntimeManager {
   // Maps a waitpoint ID to a wait ID
   private readonly resolversByWaitpoint: Map<string, string> = new Map();
 
+  private _preventMultipleWaits = preventMultipleWaits();
+
   constructor(
     private ipc: ExecutorToWorkerProcessConnection,
     private showLogs: boolean
@@ -36,40 +39,44 @@ export class ManagedRuntimeManager implements RuntimeManager {
   }
 
   async waitForTask(params: { id: string; ctx: TaskRunContext }): Promise<TaskRunExecutionResult> {
-    const promise = new Promise<CompletedWaitpoint>((resolve) => {
-      this.resolversByWaitId.set(params.id, resolve);
-    });
+    return this._preventMultipleWaits(async () => {
+      const promise = new Promise<CompletedWaitpoint>((resolve) => {
+        this.resolversByWaitId.set(params.id, resolve);
+      });
 
-    const waitpoint = await promise;
-    const result = this.waitpointToTaskRunExecutionResult(waitpoint);
+      const waitpoint = await promise;
+      const result = this.waitpointToTaskRunExecutionResult(waitpoint);
 
-    return result;
+      return result;
+    });
   }
 
   async waitForBatch(params: {
     id: string;
     runCount: number;
     ctx: TaskRunContext;
   }): Promise<BatchTaskRunExecutionResult> {
-    if (!params.runCount) {
-      return Promise.resolve({ id: params.id, items: [] });
-    }
+    return this._preventMultipleWaits(async () => {
+      if (!params.runCount) {
+        return Promise.resolve({ id: params.id, items: [] });
+      }
+
+      const promise = Promise.all(
+        Array.from({ length: params.runCount }, (_, index) => {
+          const resolverId = `${params.id}_${index}`;
+          return new Promise<CompletedWaitpoint>((resolve, reject) => {
+            this.resolversByWaitId.set(resolverId, resolve);
+          });
+        })
+      );
+
+      const waitpoints = await promise;
 
-    const promise = Promise.all(
-      Array.from({ length: params.runCount }, (_, index) => {
-        const resolverId = `${params.id}_${index}`;
-        return new Promise<CompletedWaitpoint>((resolve, reject) => {
-          this.resolversByWaitId.set(resolverId, resolve);
-        });
-      })
-    );
-
-    const waitpoints = await promise;
-
-    return {
-      id: params.id,
-      items: waitpoints.map(this.waitpointToTaskRunExecutionResult),
-    };
+      return {
+        id: params.id,
+        items: waitpoints.map(this.waitpointToTaskRunExecutionResult),
+      };
+    });
   }
 
   async waitForWaitpoint({
@@ -79,17 +86,19 @@ export class ManagedRuntimeManager implements RuntimeManager {
     waitpointFriendlyId: string;
     finishDate?: Date;
   }): Promise<WaitpointTokenResult> {
-    const promise = new Promise<CompletedWaitpoint>((resolve) => {
-      this.resolversByWaitId.set(waitpointFriendlyId, resolve);
-    });
+    return this._preventMultipleWaits(async () => {
+      const promise = new Promise<CompletedWaitpoint>((resolve) => {
+        this.resolversByWaitId.set(waitpointFriendlyId, resolve);
+      });
 
-    const waitpoint = await promise;
+      const waitpoint = await promise;
 
-    return {
-      ok: !waitpoint.outputIsError,
-      output: waitpoint.output,
-      outputType: waitpoint.outputType,
-    };
+      return {
+        ok: !waitpoint.outputIsError,
+        output: waitpoint.output,
+        outputType: waitpoint.outputType,
+      };
+    });
   }
 
   associateWaitWithWaitpoint(waitId: string, waitpointId: string) {

packages/core/src/v3/workers/taskExecutor.ts

@@ -540,17 +540,17 @@ export class TaskExecutor {
       return { status: "noop" };
     }
 
+    if (isInternalError(error) && error.skipRetrying) {
+      return { status: "skipped", error };
+    }
+
     if (
       error instanceof Error &&
       (error.name === "AbortTaskRunError" || error.name === "TaskPayloadParsedError")
     ) {
       return { status: "skipped" };
     }
 
-    if (isInternalError(error) && error.skipRetrying) {
-      return { status: "skipped", error };
-    }
-
     if (execution.run.maxAttempts) {
       retry.maxAttempts = Math.max(execution.run.maxAttempts, 1);
     }

references/hello-world/src/trigger/parallel-waits.ts

@@ -6,12 +6,14 @@ import { childTask } from "./example.js";
  */
 export const parallelWaits = task({
   id: "parallel-waits",
-  run: async (payload: any, { ctx }) => {
+  run: async ({ skipDuration = false }: { skipDuration?: boolean }) => {
     //parallel wait for 5/10 seconds
-    await Promise.all([
-      wait.for({ seconds: 5 }),
-      wait.until({ date: new Date(Date.now() + 10_000) }),
-    ]);
+    if (!skipDuration) {
+      await Promise.all([
+        wait.for({ seconds: 5 }),
+        wait.until({ date: new Date(Date.now() + 10_000) }),
+      ]);
+    }
 
     //parallel task call
     await Promise.all([