Added preventMultipleWaits to prod

Author: matt-aitken

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

@@ -6,6 +6,7 @@ import {
 } from "../schemas/index.js";
 import { ExecutorToWorkerProcessConnection } from "../zodIpc.js";
 import { RuntimeManager } from "./manager.js";
+import { preventMultipleWaits } from "./preventMultipleWaits.js";
 
 export type ProdRuntimeManagerOptions = {
   waitThresholdInMs?: number;
@@ -21,6 +22,8 @@ export class ProdRuntimeManager implements RuntimeManager {
 
   _waitForDuration: { resolve: (value: void) => void; reject: (err?: any) => void } | undefined;
 
+  _preventMultipleWaits = preventMultipleWaits();
+
   constructor(
     private ipc: ExecutorToWorkerProcessConnection,
     private options: ProdRuntimeManagerOptions = {}
@@ -31,19 +34,21 @@ export class ProdRuntimeManager implements RuntimeManager {
   }
 
   async waitForDuration(ms: number): Promise<void> {
-    const now = Date.now();
+    return this._preventMultipleWaits(async () => {
+      const now = Date.now();
 
-    const resume = new Promise<void>((resolve, reject) => {
-      this._waitForDuration = { resolve, reject };
-    });
+      const resume = new Promise<void>((resolve, reject) => {
+        this._waitForDuration = { resolve, reject };
+      });
 
-    await this.ipc.send("WAIT_FOR_DURATION", {
-      ms,
-      now,
-      waitThresholdInMs: this.waitThresholdInMs,
-    });
+      await this.ipc.send("WAIT_FOR_DURATION", {
+        ms,
+        now,
+        waitThresholdInMs: this.waitThresholdInMs,
+      });
 
-    await resume;
+      await resume;
+    });
   }
 
   resumeAfterDuration(): void {
@@ -63,51 +68,55 @@ export class ProdRuntimeManager implements RuntimeManager {
   }
 
   async waitForTask(params: { id: string; ctx: TaskRunContext }): Promise<TaskRunExecutionResult> {
-    const promise = new Promise<TaskRunExecutionResult>((resolve) => {
-      this._taskWaits.set(params.id, { resolve });
-    });
+    return this._preventMultipleWaits(async () => {
+      const promise = new Promise<TaskRunExecutionResult>((resolve) => {
+        this._taskWaits.set(params.id, { resolve });
+      });
 
-    await this.ipc.send("WAIT_FOR_TASK", {
-      friendlyId: params.id,
-    });
+      await this.ipc.send("WAIT_FOR_TASK", {
+        friendlyId: params.id,
+      });
 
-    const result = await promise;
+      const result = await promise;
 
-    clock.reset();
+      clock.reset();
 
-    return result;
+      return result;
+    });
   }
 
   async waitForBatch(params: {
     id: string;
     runs: string[];
     ctx: TaskRunContext;
   }): Promise<BatchTaskRunExecutionResult> {
-    if (!params.runs.length) {
-      return Promise.resolve({ id: params.id, items: [] });
-    }
-
-    const promise = Promise.all(
-      params.runs.map((runId) => {
-        return new Promise<TaskRunExecutionResult>((resolve, reject) => {
-          this._taskWaits.set(runId, { resolve });
-        });
-      })
-    );
-
-    await this.ipc.send("WAIT_FOR_BATCH", {
-      batchFriendlyId: params.id,
-      runFriendlyIds: params.runs,
+    return this._preventMultipleWaits(async () => {
+      if (!params.runs.length) {
+        return Promise.resolve({ id: params.id, items: [] });
+      }
+
+      const promise = Promise.all(
+        params.runs.map((runId) => {
+          return new Promise<TaskRunExecutionResult>((resolve, reject) => {
+            this._taskWaits.set(runId, { resolve });
+          });
+        })
+      );
+
+      await this.ipc.send("WAIT_FOR_BATCH", {
+        batchFriendlyId: params.id,
+        runFriendlyIds: params.runs,
+      });
+
+      const results = await promise;
+
+      clock.reset();
+
+      return {
+        id: params.id,
+        items: results,
+      };
     });
-
-    const results = await promise;
-
-    clock.reset();
-
-    return {
-      id: params.id,
-      items: results,
-    };
   }
 
   resumeTask(completion: TaskRunExecutionResult): void {