refactor: store WaitUntil in ALS

vicb · vicb · commit e89733487f30 · 2025-02-10T07:10:39.000+01:00
diff --git a/packages/open-next/src/adapters/edge-adapter.ts b/packages/open-next/src/adapters/edge-adapter.ts
@@ -1,6 +1,6 @@
 import type { ReadableStream } from "node:stream/web";
 
-import type { InternalEvent, InternalResult } from "types/open-next";
+import type { InternalEvent, InternalResult, WaitUntil } from "types/open-next";
 import { runWithOpenNextRequestContext } from "utils/promise";
 import { emptyReadableStream } from "utils/stream";
 
@@ -16,12 +16,13 @@ globalThis.__openNextAls = new AsyncLocalStorage();
 
 const defaultHandler = async (
   internalEvent: InternalEvent,
+  options?: { waitUntil?: WaitUntil },
 ): Promise<InternalResult> => {
   globalThis.isEdgeRuntime = true;
 
   // We run everything in the async local storage context so that it is available in edge runtime functions
   return runWithOpenNextRequestContext(
-    { isISRRevalidation: false },
+    { isISRRevalidation: false, waitUntil: options?.waitUntil },
     async () => {
       const host = internalEvent.headers.host
         ? `https://${internalEvent.headers.host}`
diff --git a/packages/open-next/src/adapters/image-optimization-adapter.ts b/packages/open-next/src/adapters/image-optimization-adapter.ts
@@ -22,6 +22,7 @@ import type {
   InternalEvent,
   InternalResult,
   StreamCreator,
+  WaitUntil,
 } from "types/open-next.js";
 import { emptyReadableStream, toReadableStream } from "utils/stream.js";
 
@@ -58,7 +59,7 @@ export const handler = await createGenericHandler({
 
 export async function defaultHandler(
   event: InternalEvent,
-  streamCreator?: StreamCreator,
+  options?: { streamCreator?: StreamCreator; waitUntil?: WaitUntil },
 ): Promise<InternalResult> {
   // Images are handled via header and query param information.
   debug("handler event", event);
@@ -99,9 +100,9 @@ export async function defaultHandler(
       downloadHandler,
     );
 
-    return buildSuccessResponse(result, streamCreator, etag);
+    return buildSuccessResponse(result, options?.streamCreator, etag);
   } catch (e: any) {
-    return buildFailureResponse(e, streamCreator);
+    return buildFailureResponse(e, options?.streamCreator);
   }
 }
 
diff --git a/packages/open-next/src/adapters/middleware.ts b/packages/open-next/src/adapters/middleware.ts
@@ -2,6 +2,7 @@ import type {
   InternalEvent,
   InternalResult,
   MiddlewareResult,
+  WaitUntil,
 } from "types/open-next";
 import { runWithOpenNextRequestContext } from "utils/promise";
 
@@ -24,6 +25,7 @@ globalThis.__openNextAls = new AsyncLocalStorage();
 
 const defaultHandler = async (
   internalEvent: InternalEvent,
+  options?: { waitUntil?: WaitUntil },
 ): Promise<InternalResult | MiddlewareResult> => {
   const originResolver = await resolveOriginResolver(
     globalThis.openNextConfig.middleware?.originResolver,
@@ -49,7 +51,10 @@ const defaultHandler = async (
 
   // We run everything in the async local storage context so that it is available in the external middleware
   return runWithOpenNextRequestContext(
-    { isISRRevalidation: internalEvent.headers["x-isr"] === "1" },
+    {
+      isISRRevalidation: internalEvent.headers["x-isr"] === "1",
+      waitUntil: options?.waitUntil,
+    },
     async () => {
       const result = await routingHandler(internalEvent);
       if ("internalEvent" in result) {
diff --git a/packages/open-next/src/core/requestHandler.ts b/packages/open-next/src/core/requestHandler.ts
@@ -8,6 +8,7 @@ import type {
   ResolvedRoute,
   RoutingResult,
   StreamCreator,
+  WaitUntil,
 } from "types/open-next";
 import { runWithOpenNextRequestContext } from "utils/promise";
 
@@ -29,12 +30,18 @@ patchAsyncStorage();
 
 export async function openNextHandler(
   internalEvent: InternalEvent,
-  responseStreaming?: StreamCreator,
+  options?: {
+    streamCreator?: StreamCreator;
+    waitUntil?: WaitUntil;
+  },
 ): Promise<InternalResult> {
   const initialHeaders = internalEvent.headers;
   // We run everything in the async local storage context so that it is available in the middleware as well as in NextServer
   return runWithOpenNextRequestContext(
-    { isISRRevalidation: initialHeaders["x-isr"] === "1" },
+    {
+      isISRRevalidation: initialHeaders["x-isr"] === "1",
+      waitUntil: options?.waitUntil,
+    },
     async () => {
       if (initialHeaders["x-forwarded-host"]) {
         initialHeaders.host = initialHeaders["x-forwarded-host"];
@@ -116,7 +123,7 @@ export async function openNextHandler(
 
       if ("type" in routingResult) {
         // response is used only in the streaming case
-        if (responseStreaming) {
+        if (options?.streamCreator) {
           const response = createServerResponse(
             {
               internalEvent,
@@ -127,7 +134,7 @@ export async function openNextHandler(
               initialPath: internalEvent.rawPath,
             },
             headers,
-            responseStreaming,
+            options.streamCreator,
           );
           response.statusCode = routingResult.statusCode;
           response.flushHeaders();
@@ -171,7 +178,7 @@ export async function openNextHandler(
       const res = createServerResponse(
         routingResult,
         overwrittenResponseHeaders,
-        responseStreaming,
+        options?.streamCreator,
       );
 
       await processRequest(req, res, preprocessedEvent);
diff --git a/packages/open-next/src/overrides/wrappers/aws-lambda-streaming.ts b/packages/open-next/src/overrides/wrappers/aws-lambda-streaming.ts
@@ -94,7 +94,7 @@ const handler: WrapperHandler = async (handler, converter) =>
         },
       };
 
-      const response = await handler(internalEvent, streamCreator);
+      const response = await handler(internalEvent, { streamCreator });
 
       const isUsingEdge = globalThis.isEdgeRuntime ?? false;
       if (isUsingEdge) {
diff --git a/packages/open-next/src/overrides/wrappers/aws-lambda.ts b/packages/open-next/src/overrides/wrappers/aws-lambda.ts
@@ -61,7 +61,9 @@ const handler: WrapperHandler =
       },
     };
 
-    const response = await handler(internalEvent, fakeStream);
+    const response = await handler(internalEvent, {
+      streamCreator: fakeStream,
+    });
 
     return converter.convertTo(response, event);
   };
diff --git a/packages/open-next/src/overrides/wrappers/cloudflare-edge.ts b/packages/open-next/src/overrides/wrappers/cloudflare-edge.ts
@@ -33,7 +33,6 @@ const handler: WrapperHandler<
     ctx: WorkerContext,
   ): Promise<Response> => {
     globalThis.process = process;
-    globalThis.openNextWaitUntil = ctx.waitUntil.bind(ctx);
 
     // Set the environment variables
     // Cloudflare suggests to not override the process.env object but instead apply the values to it
@@ -63,7 +62,9 @@ const handler: WrapperHandler<
       }
     }
 
-    const response = await handler(internalEvent);
+    const response = await handler(internalEvent, {
+      waitUntil: ctx.waitUntil.bind(ctx),
+    });
 
     const result: Response = await converter.convertTo(response);
 
diff --git a/packages/open-next/src/overrides/wrappers/cloudflare-node.ts b/packages/open-next/src/overrides/wrappers/cloudflare-node.ts
@@ -18,7 +18,6 @@ const handler: WrapperHandler<InternalEvent, InternalResult> =
     ctx: any,
   ): Promise<Response> => {
     globalThis.process = process;
-    globalThis.openNextWaitUntil = ctx.waitUntil.bind(ctx);
 
     // Set the environment variables
     // Cloudflare suggests to not override the process.env object but instead apply the values to it
@@ -75,7 +74,12 @@ const handler: WrapperHandler<InternalEvent, InternalResult> =
       },
     };
 
-    ctx.waitUntil(handler(internalEvent, streamCreator));
+    ctx.waitUntil(
+      handler(internalEvent, {
+        streamCreator,
+        waitUntil: ctx.waitUntil.bind(ctx),
+      }),
+    );
 
     return promiseResponse;
   };
diff --git a/packages/open-next/src/overrides/wrappers/dummy.ts b/packages/open-next/src/overrides/wrappers/dummy.ts
@@ -1,9 +1,12 @@
-import type { InternalEvent, StreamCreator } from "types/open-next";
+import type { InternalEvent, StreamCreator, WaitUntil } from "types/open-next";
 import type { Wrapper, WrapperHandler } from "types/overrides";
 
 const dummyWrapper: WrapperHandler = async (handler, converter) => {
-  return async (event: InternalEvent, responseStream?: StreamCreator) => {
-    return await handler(event, responseStream);
+  return async (
+    event: InternalEvent,
+    options?: { streamCreator?: StreamCreator; waitUntil?: WaitUntil },
+  ) => {
+    return await handler(event, options);
   };
 };
 
diff --git a/packages/open-next/src/overrides/wrappers/express-dev.ts b/packages/open-next/src/overrides/wrappers/express-dev.ts
@@ -13,25 +13,25 @@ const wrapper: WrapperHandler = async (handler, converter) => {
 
   app.all("/_next/image", async (req, res) => {
     const internalEvent = await converter.convertFrom(req);
-    const _res: StreamCreator = {
+    const streamCreator: StreamCreator = {
       writeHeaders: (prelude) => {
         res.writeHead(prelude.statusCode, prelude.headers);
         return res;
       },
     };
-    await imageHandler(internalEvent, _res);
+    await imageHandler(internalEvent, { streamCreator });
   });
 
   app.all("*paths", async (req, res) => {
     const internalEvent = await converter.convertFrom(req);
-    const _res: StreamCreator = {
+    const streamCreator: StreamCreator = {
       writeHeaders: (prelude) => {
         res.writeHead(prelude.statusCode, prelude.headers);
         return res;
       },
       onFinish: () => {},
     };
-    await handler(internalEvent, _res);
+    await handler(internalEvent, { streamCreator });
   });
 
   const server = app.listen(
diff --git a/packages/open-next/src/overrides/wrappers/node.ts b/packages/open-next/src/overrides/wrappers/node.ts
@@ -8,7 +8,7 @@ import { debug, error } from "../../adapters/logger";
 const wrapper: WrapperHandler = async (handler, converter) => {
   const server = createServer(async (req, res) => {
     const internalEvent = await converter.convertFrom(req);
-    const _res: StreamCreator = {
+    const streamCreator: StreamCreator = {
       writeHeaders: (prelude) => {
         res.setHeader("Set-Cookie", prelude.cookies);
         res.writeHead(prelude.statusCode, prelude.headers);
@@ -23,7 +23,7 @@ const wrapper: WrapperHandler = async (handler, converter) => {
       });
       res.end("OK");
     } else {
-      await handler(internalEvent, _res);
+      await handler(internalEvent, { streamCreator });
     }
   });
 
diff --git a/packages/open-next/src/types/global.ts b/packages/open-next/src/types/global.ts
@@ -10,7 +10,7 @@ import type {
 } from "types/overrides";
 
 import type { DetachedPromiseRunner } from "../utils/promise";
-import type { OpenNextConfig } from "./open-next";
+import type { OpenNextConfig, WaitUntil } from "./open-next";
 
 export interface RequestData {
   geo?: {
@@ -59,6 +59,7 @@ interface OpenNextRequestContext {
   pendingPromiseRunner: DetachedPromiseRunner;
   isISRRevalidation?: boolean;
   mergeHeadersPriority?: "middleware" | "handler";
+  waitUntil?: WaitUntil;
 }
 
 declare global {
@@ -152,13 +153,6 @@ declare global {
    */
   var __openNextAls: AsyncLocalStorage<OpenNextRequestContext>;
 
-  /**
-   * The function that is used to run background tasks even after the response has been sent.
-   * This one is defined by the wrapper function as most of them don't need or support this feature.
-   * If not present, all the awaiting promises will be resolved before sending the response.
-   */
-  var openNextWaitUntil: ((promise: Promise<void>) => void) | undefined;
-
   /**
    * The entries object that contains the functions that are available in the function.
    * Only available in edge runtime functions.
diff --git a/packages/open-next/src/types/open-next.ts b/packages/open-next/src/types/open-next.ts
@@ -48,6 +48,7 @@ export interface StreamCreator {
   onFinish?: (length: number) => void;
 }
 
+export type WaitUntil = (promise: Promise<void>) => void;
 export interface DangerousOptions {
   /**
    * The tag cache is used for revalidateTags and revalidatePath.
diff --git a/packages/open-next/src/types/overrides.ts b/packages/open-next/src/types/overrides.ts
@@ -10,6 +10,7 @@ import type {
   Origin,
   ResolvedRoute,
   StreamCreator,
+  WaitUntil,
 } from "./open-next";
 
 // Queue
@@ -125,7 +126,10 @@ export type Wrapper<
 export type OpenNextHandler<
   E extends BaseEventOrResult = InternalEvent,
   R extends BaseEventOrResult = InternalResult,
-> = (event: E, responseStream?: StreamCreator) => Promise<R>;
+> = (
+  event: E,
+  options?: { streamCreator?: StreamCreator; waitUntil?: WaitUntil },
+) => Promise<R>;
 
 export type Converter<
   E extends BaseEventOrResult = InternalEvent,
diff --git a/packages/open-next/src/utils/promise.ts b/packages/open-next/src/utils/promise.ts
@@ -1,3 +1,4 @@
+import type { WaitUntil } from "types/open-next";
 import { debug, error } from "../adapters/logger";
 
 /**
@@ -60,30 +61,30 @@ export class DetachedPromiseRunner {
 }
 
 async function awaitAllDetachedPromise() {
+  const store = globalThis.__openNextAls.getStore();
+
   const promisesToAwait =
-    globalThis.__openNextAls.getStore()?.pendingPromiseRunner.await() ??
-    Promise.resolve();
-  if (globalThis.openNextWaitUntil) {
-    globalThis.openNextWaitUntil(promisesToAwait);
+    store?.pendingPromiseRunner.await() ?? Promise.resolve();
+  if (store?.waitUntil) {
+    store.waitUntil(promisesToAwait);
     return;
   }
   await promisesToAwait;
 }
 
 function provideNextAfterProvider() {
-  /** This should be considered unstable until `unstable_after` is stablized. */
+  /** This should be considered unstable until `unstable_after` is stabilized. */
   const NEXT_REQUEST_CONTEXT_SYMBOL = Symbol.for("@next/request-context");
 
   // This is needed by some lib that relies on the vercel request context to properly await stuff.
   // Remove this when vercel builder is updated to provide '@next/request-context'.
   const VERCEL_REQUEST_CONTEXT_SYMBOL = Symbol.for("@vercel/request-context");
 
-  const openNextStoreContext = globalThis.__openNextAls.getStore();
+  const store = globalThis.__openNextAls.getStore();
 
   const waitUntil =
-    globalThis.openNextWaitUntil ??
-    ((promise: Promise<unknown>) =>
-      openNextStoreContext?.pendingPromiseRunner.add(promise));
+    store?.waitUntil ??
+    ((promise: Promise<unknown>) => store?.pendingPromiseRunner.add(promise));
 
   const nextAfterContext = {
     get: () => ({
@@ -102,14 +103,18 @@ function provideNextAfterProvider() {
 }
 
 export function runWithOpenNextRequestContext<T>(
-  { isISRRevalidation }: { isISRRevalidation: boolean },
+  {
+    isISRRevalidation,
+    waitUntil,
+  }: { isISRRevalidation: boolean; waitUntil?: WaitUntil },
   fn: () => Promise<T>,
 ): Promise<T> {
   return globalThis.__openNextAls.run(
     {
       requestId: Math.random().toString(36),
       pendingPromiseRunner: new DetachedPromiseRunner(),
       isISRRevalidation,
+      waitUntil,
     },
     async () => {
       provideNextAfterProvider();

Original file line number	Diff line number	Diff line change
`@@ -33,7 +33,6 @@ const handler: WrapperHandler<`
`33`	`33`	`ctx: WorkerContext,`
`34`	`34`	`): Promise<Response> => {`
`35`	`35`	`globalThis.process = process;`
`36`		`- globalThis.openNextWaitUntil = ctx.waitUntil.bind(ctx);`
`37`	`36`
`38`	`37`	`// Set the environment variables`
`39`	`38`	`// Cloudflare suggests to not override the process.env object but instead apply the values to it`
`@@ -63,7 +62,9 @@ const handler: WrapperHandler<`
`63`	`62`	`}`
`64`	`63`	`}`
`65`	`64`
`66`		`- const response = await handler(internalEvent);`
	`65`	`+ const response = await handler(internalEvent, {`
	`66`	`+ waitUntil: ctx.waitUntil.bind(ctx),`
	`67`	`+ });`
`67`	`68`
`68`	`69`	`const result: Response = await converter.convertTo(response);`
`69`	`70`