check edge cases for wairForStatus - prevents user from waiting indefinitely

Author: pombosilva

packages/vitest-pool-workers/src/worker/workflows.ts

@@ -1,3 +1,8 @@
+import {
+	instanceStatusName,
+	InstanceStatus as InstanceStatusNumber,
+} from "@cloudflare/workflows-shared/src/instance";
+import { WorkflowIntrospectorError } from "@cloudflare/workflows-shared/src/modifier";
 import { WORKFLOW_ENGINE_BINDING } from "../shared/workflows";
 import { internalEnv } from "./env";
 import type { InstanceModifier } from "@cloudflare/workflows-shared/src/modifier";
@@ -80,16 +85,31 @@ class WorkflowInstanceIntrospectorHandle
 		return stepResult;
 	}
 
-	async waitForStatus(status: string): Promise<void> {
+	async waitForStatus(status: InstanceStatus["status"]): Promise<void> {
 		console.log("[Vitest-Workflows] waiting for status");
+
+		if (
+			status === instanceStatusName(InstanceStatusNumber.Terminated) ||
+			status === instanceStatusName(InstanceStatusNumber.Paused)
+		) {
+			throw new WorkflowIntrospectorError(
+				`InstanceStatus '${status}' is not implemented yet and cannot be waited.`
+			);
+		}
+
+		if (status === instanceStatusName(InstanceStatusNumber.Queued)) {
+			// we currently don't have a queue mechanism, but it would happen before it
+			// starts running, so waiting for it to be queued should always return
+			return;
+		}
 		// @ts-expect-error waitForStatus not exposed
 		await this.engineStub.waitForStatus(status);
 
 		console.log("[Vitest-Workflows] status awaited");
 	}
 
 	async cleanUp(): Promise<void> {
-		// works with isolatedStorage = false
+		// this cleans state with isolatedStorage = false
 		try {
 			// @ts-expect-error DO binding created in runner worker start
 			await this.engineStub.unsafeAbort("user called delete");

packages/workflows-shared/src/engine.ts

@@ -16,7 +16,7 @@ import {
 	startGracePeriod,
 } from "./lib/gracePeriodSemaphore";
 import { TimePriorityQueue } from "./lib/timePriorityQueue";
-import { InstanceModifier } from "./modifier";
+import { InstanceModifier, WorkflowIntrospectorError } from "./modifier";
 import type { Event } from "./context";
 import type { InstanceMetadata, RawInstanceLog } from "./instance";
 import type { StepSelector } from "./modifier";
@@ -226,14 +226,16 @@ export class Engine extends DurableObject<Env> {
 	): Promise<void> {
 		await this.ctx.storage.put(ENGINE_STATUS_KEY, status);
 
-		// Check if anyone is waiting for this new status
+		// check if anyone is waiting for this status
 		this.handleStatusWaiter(status);
 
 		console.log("[Engine] Changed my status to", instanceStatusName(status));
 	}
 
-	private statusWaiters: Map<InstanceStatus, { resolve: () => void }> =
-		new Map();
+	private statusWaiters: Map<
+		InstanceStatus,
+		{ resolve: () => void; reject: (e: unknown) => void }
+	> = new Map();
 	async waitForStatus(status: string): Promise<void> {
 		const targetStatus = toInstanceStatus(status);
 		const currentStatus =
@@ -245,16 +247,74 @@ export class Engine extends DurableObject<Env> {
 		}
 
 		// if it hasn't reached the desired state, create a new promise and add its resolver to the waiters map
-		return new Promise((resolve) => {
-			this.statusWaiters.set(targetStatus, { resolve });
+		return new Promise((resolve, reject) => {
+			this.statusWaiters.set(targetStatus, { resolve, reject });
 		});
 	}
 
-	handleStatusWaiter(status: InstanceStatus) {
+	handleStatusWaiter(status: InstanceStatus): void {
 		const waiter = this.statusWaiters.get(status);
+
+		// resolve if it reached the desired status
 		if (waiter) {
 			waiter.resolve();
 			this.statusWaiters.delete(status);
+			return;
+		}
+
+		switch (status) {
+			case InstanceStatus.Errored: {
+				// if it reaches final status "errored", then it can't be waiting for it to complete or terminate
+				const unreachableStatuses = [
+					InstanceStatus.Complete,
+					InstanceStatus.Terminated,
+				];
+
+				this.rejectUnreachableStatus(status, unreachableStatuses);
+				break;
+			}
+			case InstanceStatus.Terminated: {
+				// if it reaches final status "terminated", then it can't be waiting for it to complete or error
+				const unreachableStatuses = [
+					InstanceStatus.Complete,
+					InstanceStatus.Errored,
+				];
+
+				this.rejectUnreachableStatus(status, unreachableStatuses);
+				break;
+			}
+			case InstanceStatus.Complete: {
+				// if it reaches final status "complete", then it can't be waiting for it to terminate or error
+				const unreachableStatuses = [
+					InstanceStatus.Terminated,
+					InstanceStatus.Errored,
+				];
+
+				this.rejectUnreachableStatus(status, unreachableStatuses);
+				break;
+			}
+			default:
+				break;
+		}
+	}
+
+	rejectUnreachableStatus(
+		reachedStatus: number,
+		unreachableStatuses: number[]
+	): void {
+		if (unreachableStatuses) {
+			for (const unreachableStatus of unreachableStatuses) {
+				const waiter = this.statusWaiters.get(unreachableStatus);
+				if (waiter) {
+					waiter.reject(
+						new WorkflowIntrospectorError(
+							`The Wokflow instance ${this.instanceId} has reached status '${instanceStatusName(reachedStatus)}'. This is a final state that prevents it from ever reaching the expected status of '${instanceStatusName(unreachableStatus)}'.`
+						)
+					);
+					this.statusWaiters.delete(unreachableStatus);
+					return;
+				}
+			}
 		}
 	}
 
@@ -287,7 +347,7 @@ export class Engine extends DurableObject<Env> {
 				return parsed?.result;
 			}
 			if (event === InstanceEvent.STEP_FAILURE) {
-				throw parsed?.error ?? parsed ?? new Error("Step failed");
+				throw parsed?.error ?? parsed;
 			}
 		}
 
@@ -311,8 +371,8 @@ export class Engine extends DurableObject<Env> {
 			waiter.resolve(result);
 			this.stepResultWaiters.delete(group);
 		} else if (event === InstanceEvent.STEP_FAILURE) {
-			const errorLike = metadata?.error ?? metadata;
-			waiter.reject(errorLike);
+			const error = metadata?.error ?? new Error("Step failed");
+			waiter.reject(error);
 			this.stepResultWaiters.delete(group);
 		}
 	}

packages/workflows-shared/src/instance.ts

@@ -83,7 +83,9 @@ export function toInstanceStatus(status: string): InstanceStatus {
 		case "unknown":
 			throw new Error("unknown cannot be parsed into a InstanceStatus");
 		default:
-			throw new Error(`${status} was not handled`);
+			throw new Error(
+				`${status} was not handled because it's not a valid InstanceStatus`
+			);
 	}
 }