Fix nested waitUntil (#606)

* Fix nested `waitUntil` not being awaited

* Remove changeset

* Move `waitUntilAll()` to `standards/event.ts`

* Add nested `waitUntil` tests for `scheduled` and `queue` events

* Add nested `waitUntil` tests for unit testing environments

---------

Co-authored-by: bcoll <[email protected]>

Author: hansottowirtz

packages/core/src/standards/event.ts

@@ -32,6 +32,21 @@ const SUGGEST_GLOBAL_BINDING_MODULES =
   "handlers/Durable Object constructors, or `context.env` with " +
   "Pages Functions.";
 
+// Like `Promise.all()`, but also handles nested changes to the promises array
+export async function waitUntilAll<WaitUntil extends any[] = unknown[]>(
+  promises: Promise<unknown>[]
+): Promise<WaitUntil> {
+  let len = 0;
+  let last: WaitUntil = [] as unknown as WaitUntil;
+  // When the length of the array changes, there has been a nested call to
+  // `waitUntil` and we should await the promises again
+  while (len !== promises.length) {
+    len = promises.length;
+    last = (await Promise.all(promises)) as WaitUntil;
+  }
+  return last;
+}
+
 const kResponse = Symbol("kResponse");
 const kPassThrough = Symbol("kPassThrough");
 export const kWaitUntil = Symbol("kWaitUntil");
@@ -404,7 +419,7 @@ export class ServiceWorkerGlobalScope extends WorkerGlobalScope {
       }
 
       // noinspection ES6MissingAwait
-      const waitUntil = Promise.all(event[kWaitUntil]) as Promise<WaitUntil>;
+      const waitUntil = waitUntilAll<WaitUntil>(event[kWaitUntil]);
       return withWaitUntil(res, waitUntil);
     }
 
@@ -436,7 +451,7 @@ export class ServiceWorkerGlobalScope extends WorkerGlobalScope {
     }
 
     // noinspection ES6MissingAwait
-    const waitUntil = Promise.all(event[kWaitUntil]) as Promise<WaitUntil>;
+    const waitUntil = waitUntilAll<WaitUntil>(event[kWaitUntil]);
     return withWaitUntil(await fetch(request), waitUntil);
   }
 
@@ -449,15 +464,15 @@ export class ServiceWorkerGlobalScope extends WorkerGlobalScope {
       cron: cron ?? "",
     });
     super.dispatchEvent(event);
-    return (await Promise.all(event[kWaitUntil])) as WaitUntil;
+    return waitUntilAll<WaitUntil>(event[kWaitUntil]);
   }
 
   async [kDispatchQueue]<WaitUntil extends any[] = any[]>(
     batch: MessageBatch
   ): Promise<WaitUntil> {
     const event = new QueueEvent("queue", { batch });
     super.dispatchEvent(event);
-    return (await Promise.all(event[kWaitUntil])) as WaitUntil;
+    return waitUntilAll<WaitUntil>(event[kWaitUntil]);
   }
 
   // eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types

packages/core/test/index.spec.ts

@@ -1384,3 +1384,48 @@ test("MiniflareCore: dispose: cleans up watcher", async (t) => {
   await setTimeout(100); // Shouldn't reload here
   t.is(log.logsAtLevel(LogLevel.DEBUG).length, 0);
 });
+test("MiniflareCore: dispatchFetch: awaits nested waitUntil", async (t) => {
+  const mf = useMiniflare(
+    {},
+    {
+      script: `
+        async function waitAgain(ctx) {
+          await scheduler.wait(100);
+          ctx.waitUntil(scheduler.wait(100).then(() => 2));
+          return 1;
+        }
+      
+        export default {
+          async fetch(req, env, ctx) {
+            ctx.waitUntil(waitAgain(ctx));
+            return new Response();
+          },
+          async scheduled(controller, env, ctx) {
+            ctx.waitUntil(waitAgain(ctx));
+            return 3;
+          },
+          async queue(batch, env, ctx) {
+            ctx.waitUntil(waitAgain(ctx));
+            return 4;
+          }
+        }
+      `,
+      modules: true,
+    },
+    log
+  );
+
+  const res = await mf.dispatchFetch("https://test");
+  let waitUntil = await res.waitUntil();
+  t.deepEqual(waitUntil, [1, 2]);
+
+  waitUntil = await mf.dispatchScheduled();
+  t.deepEqual(waitUntil, [1, 3, 2]);
+
+  waitUntil = await mf.dispatchQueue({
+    queue: "queue",
+    messages: [],
+    retryAll() {},
+  });
+  t.deepEqual(waitUntil, [1, 4, 2]);
+});

packages/jest-environment-miniflare/test/fixtures/waitUntil.worker.spec.js

@@ -71,3 +71,16 @@ test("ExecutionContext: waitUntil: await resolving promise array using getMinifl
   expect(cachedResponse).toBeTruthy();
   expect(await cachedResponse.text()).toBe("example cache");
 });
+
+test("getMiniflareWaitUntil: supports nested wait until", async () => {
+  const ctx = new ExecutionContext();
+  ctx.waitUntil(
+    (async () => {
+      await scheduler.wait(100);
+      ctx.waitUntil(scheduler.wait(100).then(() => 2));
+      return 1;
+    })()
+  );
+  const waitUntilList = await getMiniflareWaitUntil(ctx);
+  expect(waitUntilList).toEqual([1, 2]);
+});

packages/shared-test-environment/src/globals.ts

@@ -3,6 +3,7 @@ import {
   MiniflareCore,
   ScheduledEvent,
   kWaitUntil,
+  waitUntilAll,
 } from "@miniflare/core";
 import {
   DurableObject,
@@ -88,7 +89,7 @@ export async function createMiniflareEnvironmentUtilities(
     getMiniflareWaitUntil<WaitUntil extends any[] = unknown[]>(
       event: FetchEvent | ScheduledEvent | ExecutionContext
     ): Promise<WaitUntil> {
-      return Promise.all(event[kWaitUntil]) as Promise<WaitUntil>;
+      return waitUntilAll<WaitUntil>(event[kWaitUntil]);
     },
     async flushMiniflareDurableObjectAlarms(ids?: DurableObjectId[]) {
       const plugin = (await mf.getPlugins()).DurableObjectsPlugin;

packages/vitest-environment-miniflare/test/fixtures/service-worker/waitUntil.worker.spec.js

@@ -74,3 +74,16 @@ test("ExecutionContext: waitUntil: await resolving promise array using getMinifl
   expect(cachedResponse).toBeTruthy();
   expect(await cachedResponse.text()).toBe("example cache");
 });
+
+test("getMiniflareWaitUntil: supports nested wait until", async () => {
+  const ctx = new ExecutionContext();
+  ctx.waitUntil(
+    (async () => {
+      await scheduler.wait(100);
+      ctx.waitUntil(scheduler.wait(100).then(() => 2));
+      return 1;
+    })()
+  );
+  const waitUntilList = await getMiniflareWaitUntil(ctx);
+  expect(waitUntilList).toEqual([1, 2]);
+});