[miniflare] fix: ensure Mutex doesn't report itself as drained if l…

…ocked (#4321)

* fix: ensure `Mutex` doesn't report itself as drained if locked

* fix: ensure proxies poisoned synchronously

Previously, magic proxies were only poisoned once the runtime had
restarted. This meant it was possible for heap objects to be
freed while they were still in use, e.g.

1. `const caches = await Miniflare#getCaches()` _(create proxy)_
2. `await caches.default.put(...)` _(resolve `put()` function)_
3. `void Miniflare#setOptions()` _(in background)_
4. Either...
   a. `caches` proxy GCed, `FinalizationRegistry` callback called,
      not yet poisoned, so calls `dispatchFetch()`, but
      `dispatchFetch()` blocked until `setOptions()` completes
   b. `caches.default.put(...)`, calls `dispatchFetch()`, but
      `dispatchFetch()` blocked until `setOptions()` completes
5. `setOptions()` completes, poisoning proxies, but too late
6. `dispatchFetch()` unblocked, sends request to incorrect DO

The "blocked until `setOptions()` completes" step was fixed by the
previous commit, hence we're only seeing this bug now.

This change ensures we poison proxies synchronously as soon as
`setOptions()` is called, ensuring the problematic `dispatchFetch()`
calls never happen.

.changeset/rude-chicken-raise.md

@@ -0,0 +1,13 @@
+---
+"miniflare": patch
+---
+
+fix: ensure `Mutex` doesn't report itself as drained if locked
+
+Previously, Miniflare's `Mutex` implementation would report itself as drained
+if there were no waiters, regardless of the locked state. This bug meant that
+if you called but didn't `await` `Miniflare#setOptions()`, future calls to
+`Miniflare#dispatchFetch()` (or any other asynchronous `Miniflare` method)
+wouldn't wait for the options update to apply and the runtime to restart before
+sending requests. This change ensures we wait until the mutex is unlocked before
+reporting it as drained.

packages/miniflare/src/index.ts

@@ -1160,10 +1160,9 @@ export class Miniflare {
 				this.dispatchFetch
 			);
 		} else {
-			// The `ProxyServer` "heap" will have been destroyed when `workerd` was
-			// restarted, invalidating all existing native target references. Mark
-			// all proxies as invalid, noting the new runtime URL to send requests to.
-			this.#proxyClient.poisonProxies(this.#runtimeEntryURL);
+			// Update the proxy client with the new runtime URL to send requests to.
+			// Existing proxies will already have been poisoned in `setOptions()`.
+			this.#proxyClient.setRuntimeEntryURL(this.#runtimeEntryURL);
 		}
 
 		if (!this.#runtimeMutex.hasWaiting) {
@@ -1293,6 +1292,9 @@ export class Miniflare {
 
 	setOptions(opts: MiniflareOptions): Promise<void> {
 		this.#checkDisposed();
+		// The `ProxyServer` "heap" will be destroyed when `workerd` restarts,
+		// invalidating all existing native references. Mark all proxies as invalid.
+		this.#proxyClient?.poisonProxies();
 		// Wait for initial initialisation and other setOptions to complete before
 		// changing options
 		return this.#runtimeMutex.runWith(() => this.#setOptions(opts));
@@ -1496,6 +1498,11 @@ export class Miniflare {
 
 	async dispose(): Promise<void> {
 		this.#disposeController.abort();
+		// The `ProxyServer` "heap" will be destroyed when `workerd` shuts down,
+		// invalidating all existing native references. Mark all proxies as invalid.
+		// Note `dispose()`ing the `#proxyClient` implicitly poison's proxies, but
+		// we'd like them to be poisoned synchronously here.
+		this.#proxyClient?.poisonProxies();
 		try {
 			await this.#waitForReady(/* disposing */ true);
 		} finally {

packages/miniflare/src/plugins/core/proxy/client.ts

@@ -81,13 +81,19 @@ export class ProxyClient {
 		return (this.#envProxy ??= this.#bridge.getProxy(TARGET_ENV));
 	}
 
-	poisonProxies(runtimeEntryURL?: URL): void {
-		this.#bridge.poisonProxies(runtimeEntryURL);
+	poisonProxies(): void {
+		this.#bridge.poisonProxies();
 		// Reset `#{global,env}Proxy` so they aren't poisoned on next access
 		this.#globalProxy = undefined;
 		this.#envProxy = undefined;
 	}
 
+	setRuntimeEntryURL(runtimeEntryURL: URL) {
+		// This function will be called whenever the runtime restarts. The URL may
+		// be different if the port has changed.
+		this.#bridge.url = runtimeEntryURL;
+	}
+
 	dispose(): Promise<void> {
 		// Intentionally not resetting `#{global,env}Proxy` to keep them poisoned.
 		// `workerd` won't be started again by this `Miniflare` instance after
@@ -188,13 +194,12 @@ class ProxyClientBridge {
 		return proxy;
 	}
 
-	poisonProxies(url?: URL): void {
+	poisonProxies(): void {
 		this.#version++;
-		// This function will be called whenever the runtime restarts. The URL may
-		// be different if the port has changed. We must also unregister all
-		// finalizers as the heap will be reset, and we don't want a new object
-		// added with the same address to be freed when it's still accessible.
-		if (url !== undefined) this.url = url;
+		// This function will be called whenever `setOptions()` or `dispose()` is
+		// called. We must also unregister all finalizers as the heap will be reset,
+		// and we don't want a new object added with the same address to be freed
+		// when it's still accessible.
 		this.#finalizationRegistry.unregister(this);
 	}
 
@@ -399,6 +404,8 @@ class ProxyStubHandler<T extends object> implements ProxyHandler<T> {
 	}
 
 	getOwnPropertyDescriptor(target: T, key: string | symbol) {
+		this.#assertSafe();
+
 		if (typeof key === "symbol") return undefined;
 
 		// Optimisation: assume constant prototypes of proxied objects, descriptors
@@ -424,6 +431,8 @@ class ProxyStubHandler<T extends object> implements ProxyHandler<T> {
 	}
 
 	ownKeys(_target: T) {
+		this.#assertSafe();
+
 		// Optimisation: assume constant prototypes of proxied objects, own keys
 		// should never change after we've fetched them
 		if (this.#knownOwnKeys !== undefined) return this.#knownOwnKeys;
@@ -442,6 +451,8 @@ class ProxyStubHandler<T extends object> implements ProxyHandler<T> {
 	}
 
 	getPrototypeOf(_target: T) {
+		this.#assertSafe();
+
 		// Return a `null` prototype, so users know this isn't a plain object
 		return null;
 	}

packages/miniflare/src/workers/shared/sync.ts

@@ -71,7 +71,7 @@ export class Mutex {
 	}
 
 	async drained(): Promise<void> {
-		if (this.resolveQueue.length === 0) return;
+		if (this.resolveQueue.length === 0 && !this.locked) return;
 		return new Promise((resolve) => this.drainQueue.push(resolve));
 	}
 }

packages/miniflare/test/shared/sync.spec.ts

@@ -46,6 +46,7 @@ test("Mutex: maintains separate drain queue", async (t) => {
 	void mutex.runWith(() => deferred1);
 	let drained = false;
 	mutex.drained().then(() => (drained = true));
+	await setTimeout();
 	t.false(drained);
 	deferred1.resolve();
 	await setTimeout();
@@ -64,6 +65,7 @@ test("Mutex: maintains separate drain queue", async (t) => {
 	});
 	drained = false;
 	mutex.drained().then(() => (drained = true));
+	await setTimeout();
 	t.false(drained);
 	deferred2.resolve();
 	await setTimeout();