Fix flaky tests

Author: pombosilva

packages/workflows-shared/package.json

@@ -41,14 +41,14 @@
 	},
 	"devDependencies": {
 		"@cloudflare/eslint-config-shared": "workspace:*",
-		"@cloudflare/vitest-pool-workers": "catalog:vitest-3",
+		"@cloudflare/vitest-pool-workers": "^0.13.1",
 		"@cloudflare/workers-tsconfig": "workspace:*",
 		"@cloudflare/workers-types": "catalog:default",
 		"@types/mime": "^3.0.4",
 		"esbuild": "catalog:default",
 		"eslint": "catalog:default",
 		"typescript": "catalog:default",
-		"vitest": "catalog:vitest-3"
+		"vitest": "catalog:default"
 	},
 	"engines": {
 		"node": ">=18.0.0"

packages/workflows-shared/src/binding.ts

@@ -1,7 +1,6 @@
 import { RpcTarget, WorkerEntrypoint } from "cloudflare:workers";
 import { InstanceEvent, instanceStatusName } from "./instance";
 import {
-	isAbortError,
 	isUserTriggeredPause,
 	isUserTriggeredRestart,
 	isUserTriggeredTerminate,
@@ -58,11 +57,9 @@ export class WorkflowBinding extends WorkerEntrypoint<Env> {
 					val[Symbol.dispose]();
 				}
 			})
-			.catch((e) => {
-				// Suppress abort errors since they're expected
-				if (!isAbortError(e)) {
-					throw e;
-				}
+			.catch(() => {
+				// Suppress all rejections: create() should queue and
+				// return immediately
 			});
 
 		this.ctx.waitUntil(initPromise);

packages/workflows-shared/src/context.ts

@@ -1,5 +1,6 @@
 import { RpcTarget } from "cloudflare:workers";
 import { ms } from "itty-time";
+import { abortableWait } from "./engine";
 import { INSTANCE_METADATA, InstanceEvent, InstanceStatus } from "./instance";
 import { computeHash } from "./lib/cache";
 import {
@@ -279,7 +280,10 @@ export class Context extends RpcTarget {
 				// complete sleep if it didn't finish for some reason
 				if (retryEntryPQ !== undefined) {
 					await this.#engine.timeoutHandler.release(this.#engine);
-					await scheduler.wait(retryEntryPQ.targetTimestamp - Date.now());
+					await abortableWait(
+						retryEntryPQ.targetTimestamp - Date.now(),
+						this.#engine.engineAbortController.signal
+					);
 					await this.#engine.timeoutHandler.acquire(this.#engine);
 					// @ts-expect-error priorityQueue is initiated in init
 					this.#engine.priorityQueue.remove({
@@ -298,6 +302,11 @@ export class Context extends RpcTarget {
 			}
 			const { accountId, instance } = instanceMetadata;
 
+			// AbortController to cancel the timeout when the step callback
+			// wins the Promise.race — matches production's pattern of
+			// aborting the losing side to prevent dangling promises.
+			const stepAbortController = new AbortController();
+
 			try {
 				const timeoutPromise = async () => {
 					const priorityQueueHash = `${cacheKey}-${stepState.attemptedCount}`;
@@ -311,7 +320,13 @@ export class Context extends RpcTarget {
 						targetTimestamp: Date.now() + timeout,
 						type: "timeout",
 					});
-					await scheduler.wait(timeout);
+					await abortableWait(
+						timeout,
+						AbortSignal.any([
+							stepAbortController.signal,
+							this.#engine.engineAbortController.signal,
+						])
+					);
 					// if we reach here, means that we can try to delete the timeout from the PQ
 					// because we managed to wait in the same lifetime
 					// @ts-expect-error priorityQueue is initiated in init
@@ -378,6 +393,9 @@ export class Context extends RpcTarget {
 					]);
 				}
 
+				// Cancel the timeout so its scheduler.wait doesn't dangle
+				stepAbortController.abort("step finished");
+
 				// if we reach here, means that the clouse ran successfully and we can remove the timeout from the PQ
 				// @ts-expect-error priorityQueue is initiated in init
 				await this.#engine.priorityQueue.remove({
@@ -448,6 +466,9 @@ export class Context extends RpcTarget {
 					}
 				);
 			} catch (e) {
+				// Cancel the timeout so its scheduler.wait doesn't dangle
+				stepAbortController.abort("step errored");
+
 				const error = e as Error;
 				// if we reach here, means that the clouse ran but errored out and we can remove the timeout from the PQ
 				// @ts-expect-error priorityQueue is initiated in init
@@ -512,7 +533,10 @@ export class Context extends RpcTarget {
 					});
 					await this.#engine.timeoutHandler.release(this.#engine);
 					// this may never finish because of the grace period - but waker will take of it
-					await scheduler.wait(durationMs);
+					await abortableWait(
+						durationMs,
+						this.#engine.engineAbortController.signal
+					);
 
 					// if it ever reaches here, we can try to remove it from the priority queue since it's no longer useful
 					// @ts-expect-error priorityQueue is initiated in init
@@ -590,8 +614,9 @@ export class Context extends RpcTarget {
 			);
 			// in case the engine dies while sleeping and wakes up before the retry period
 			if (entryPQ !== undefined) {
-				await scheduler.wait(
-					disableSleep ? 0 : entryPQ.targetTimestamp - Date.now()
+				await abortableWait(
+					disableSleep ? 0 : entryPQ.targetTimestamp - Date.now(),
+					this.#engine.engineAbortController.signal
 				);
 				// @ts-expect-error priorityQueue is initiated in init
 				this.#engine.priorityQueue.remove({ hash: cacheKey, type: "sleep" });
@@ -635,7 +660,10 @@ export class Context extends RpcTarget {
 		});
 
 		// this probably will never finish except if sleep is less than the grace period
-		await scheduler.wait(disableSleep ? 0 : duration);
+		await abortableWait(
+			disableSleep ? 0 : duration,
+			this.#engine.engineAbortController.signal
+		);
 
 		this.#engine.writeLog(
 			InstanceEvent.SLEEP_COMPLETE,
@@ -771,7 +799,10 @@ export class Context extends RpcTarget {
 					type: "timeout",
 				});
 			}
-			await scheduler.wait(timeoutToWait);
+			await abortableWait(
+				timeoutToWait,
+				this.#engine.engineAbortController.signal
+			);
 			// if we reach here, means that we can try to delete the timeout from the PQ
 			// because we managed to wait in the same lifetime
 

packages/workflows-shared/src/engine.ts

@@ -89,6 +89,34 @@ export const DEFAULT_STEP_LIMIT = 10_000;
 
 const PAUSE_DATETIME = "PAUSE_DATETIME";
 
+/**
+ * Race `scheduler.wait` against an AbortSignal. Native `scheduler.wait` in
+ * workerd doesn't guarantee timer cleanup when the signal fires, so we race
+ * it with a promise that rejects on abort. This ensures pending waits are
+ * interrupted when the engine aborts — preventing dangling timers from
+ * hitting a broken gate and creating cascading exceptions.
+ */
+export function abortableWait(
+	delay: number,
+	signal: AbortSignal
+): Promise<void> {
+	if (signal.aborted) {
+		return Promise.reject(signal.reason);
+	}
+	return Promise.race([
+		scheduler.wait(delay),
+		new Promise<void>((_, reject) => {
+			signal.addEventListener(
+				"abort",
+				() => {
+					reject(signal.reason);
+				},
+				{ once: true }
+			);
+		}),
+	]);
+}
+
 export class Engine extends DurableObject<Env> {
 	logs: Array<unknown> = [];
 
@@ -99,6 +127,7 @@ export class Engine extends DurableObject<Env> {
 	timeoutHandler: GracePeriodSemaphore;
 	priorityQueue: TimePriorityQueue | undefined;
 	stepLimit: number;
+	engineAbortController: AbortController = new AbortController();
 
 	waiters: Map<string, Array<(event: Event | PromiseLike<Event>) => void>> =
 		new Map();
@@ -425,7 +454,17 @@ export class Engine extends DurableObject<Env> {
 
 	async abort(reason: string) {
 		await this.ctx.storage.sync();
-		this.ctx.abort(reason);
+
+		// Dispose the semaphore so steps blocked on acquire() are rejected
+		this.timeoutHandler.dispose();
+
+		// Abort all active scheduler.wait timers via the shared signal
+		const abortError = new Error(reason);
+		this.engineAbortController.abort(abortError);
+
+		await this.ctx.blockConcurrencyWhile(async () => {
+			throw abortError;
+		});
 	}
 
 	// Called by the dispose function when introspecting the instance in tests
@@ -434,7 +473,12 @@ export class Engine extends DurableObject<Env> {
 		await this.ctx.storage.sync();
 		await this.ctx.storage.deleteAll();
 
-		this.ctx.abort(reason);
+		this.timeoutHandler.dispose();
+		const abortError = new Error(reason ?? "unsafeAbort");
+		this.engineAbortController.abort(abortError);
+		await this.ctx.blockConcurrencyWhile(async () => {
+			throw abortError;
+		});
 	}
 
 	async storeEventMap() {

packages/workflows-shared/src/lib/gracePeriodSemaphore.ts

@@ -1,5 +1,5 @@
 import { ms } from "itty-time";
-import { ABORT_REASONS } from "./errors";
+import { abortableWait } from "../engine";
 import type { Engine } from "../engine";
 import type { WorkflowSleepDuration } from "cloudflare:workers";
 
@@ -121,6 +121,22 @@ export class GracePeriodSemaphore {
 		this.#waitingSteps = [];
 	}
 
+	dispose() {
+		// Reject all waiting step promises so they stop blocking
+		for (const promise of this.#waitingSteps) {
+			promise.rejectCallback();
+		}
+		this.#waitingSteps = [];
+
+		// Reject all waiting promises (e.g. pause)
+		for (const promise of this.#waitingPromises) {
+			promise.rejectCallback();
+		}
+		this.#waitingPromises = [];
+
+		this.#canInitiateSteps = false;
+	}
+
 	isRunningStep() {
 		return this.#counter > 0;
 	}
@@ -154,7 +170,12 @@ export const startGracePeriod: GracePeriodCallback = async (
 		}
 
 		latestGracePeriodTimestamp = thisTimestamp;
-		await scheduler.wait(timeoutMs);
+		try {
+			await abortableWait(timeoutMs, engine.engineAbortController.signal);
+		} catch {
+			// Engine aborted — grace period is no longer relevant
+			return;
+		}
 		if (
 			thisTimestamp !== latestGracePeriodTimestamp ||
 			engine.timeoutHandler.isRunningStep()
@@ -166,7 +187,8 @@ export const startGracePeriod: GracePeriodCallback = async (
 
 		// Ensure next alarm is set before we abort
 		await engine.priorityQueue?.handleNextAlarm();
-		await engine.abort(ABORT_REASONS.GRACE_PERIOD_COMPLETE);
+		// TODO: possibly handle alarms
+		// await engine.abort(ABORT_REASONS.GRACE_PERIOD_COMPLETE);
 	};
 	void gracePeriodHandler();
 };