Completed batch waitpoints when we completed the BatchTaskRun

Author: matt-aitken

apps/webapp/app/runEngine/services/batchTrigger.server.ts

@@ -571,17 +571,6 @@ export class RunEngineBatchTriggerService extends WithRunEngine {
 
     //triggered all the runs
     if (updatedBatch.runIds.length === updatedBatch.runCount) {
-      //unblock the parent run from the batch
-      //this prevents the parent continuing before all the runs are created
-      if (parentRunId && resumeParentOnCompletion) {
-        await this._engine.unblockRunForCreatedBatch({
-          runId: RunId.fromFriendlyId(parentRunId),
-          batchId: batch.id,
-          environmentId: environment.id,
-          projectId: environment.projectId,
-        });
-      }
-
       //if all the runs were idempotent, it's possible the batch is already completed
       await this._engine.tryCompleteBatch({ batchId: batch.id });
     }

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

@@ -290,6 +290,7 @@ export class RunEngine {
 
     this.batchSystem = new BatchSystem({
       resources,
+      waitpointSystem: this.waitpointSystem,
     });
 
     this.runAttemptSystem = new RunAttemptSystem({
@@ -905,43 +906,6 @@ export class RunEngine {
     }
   }
 
-  /**
-   * This is called when all the runs for a batch have been created.
-   * This does NOT mean that all the runs for the batch are completed.
-   */
-  async unblockRunForCreatedBatch({
-    runId,
-    batchId,
-    tx,
-  }: {
-    runId: string;
-    batchId: string;
-    environmentId: string;
-    projectId: string;
-    tx?: PrismaClientOrTransaction;
-  }): Promise<void> {
-    const prisma = tx ?? this.prisma;
-
-    const waitpoint = await prisma.waitpoint.findFirst({
-      where: {
-        completedByBatchId: batchId,
-      },
-    });
-
-    if (!waitpoint) {
-      this.logger.error("RunEngine.unblockRunForBatch(): Waitpoint not found", {
-        runId,
-        batchId,
-      });
-      throw new ServiceValidationError("Waitpoint not found for batch", 404);
-    }
-
-    await this.completeWaitpoint({
-      id: waitpoint.id,
-      output: { value: "Batch waitpoint completed", isError: false },
-    });
-  }
-
   async tryCompleteBatch({ batchId }: { batchId: string }): Promise<void> {
     return this.batchSystem.scheduleCompleteBatch({ batchId });
   }

internal-packages/run-engine/src/engine/systems/batchSystem.ts

@@ -1,16 +1,20 @@
 import { startSpan } from "@internal/tracing";
 import { isFinalRunStatus } from "../statuses.js";
 import { SystemResources } from "./systems.js";
+import { WaitpointSystem } from "./waitpointSystem.js";
 
 export type BatchSystemOptions = {
   resources: SystemResources;
+  waitpointSystem: WaitpointSystem;
 };
 
 export class BatchSystem {
   private readonly $: SystemResources;
+  private readonly waitpointSystem: WaitpointSystem;
 
   constructor(private readonly options: BatchSystemOptions) {
     this.$ = options.resources;
+    this.waitpointSystem = options.waitpointSystem;
   }
 
   public async scheduleCompleteBatch({ batchId }: { batchId: string }): Promise<void> {
@@ -75,6 +79,28 @@ export class BatchSystem {
             status: "COMPLETED",
           },
         });
+
+        //get waitpoint (if there is one)
+        const waitpoint = await this.$.prisma.waitpoint.findFirst({
+          where: {
+            completedByBatchId: batchId,
+          },
+        });
+
+        if (!waitpoint) {
+          this.$.logger.debug(
+            "RunEngine.unblockRunForBatch(): Waitpoint not found. This is ok, because only batchTriggerAndWait has waitpoints",
+            {
+              batchId,
+            }
+          );
+          return;
+        }
+
+        await this.waitpointSystem.completeWaitpoint({
+          id: waitpoint.id,
+          output: { value: "Batch waitpoint completed", isError: false },
+        });
       } else {
         this.$.logger.debug("#tryCompleteBatch: Not all runs are completed", { batchId });
       }