Improve internal waitUntil utility (#56720)

⚠️ This is an internal API and will be removed soon. Please do not use.

Refactors #56404 to have a better internal API used by both Edge SSR and Edge Route Handlers.

This new API can buffer non-synchronously created "waitUntil"s even after the response has been returned, just need to make sure that there's at least one "waitUntil" queued. E.g.:

```js
async function handler() {
  internal_runWithWaitUntil(async () => { // ← no await
    await taskA()

    internal_runWithWaitUntil(async () => { // ← no await
      await longRunningTaskB()
    })

    await taskC()
  }) 

  return Response(...)
}
```

Internally, the "waitUntil" promise will resolve after all tasks are finished.

cc @ijjk @cramforce @feedthejim as we've synced about some of the details here. Not using ALS because of some promise-related issues.

Author: shuding

packages/next/src/build/webpack/loaders/next-edge-ssr-loader/render.ts

@@ -16,15 +16,7 @@ import { SERVER_RUNTIME } from '../../../../lib/constants'
 import type { PrerenderManifest } from '../../..'
 import { normalizeAppPath } from '../../../../shared/lib/router/utils/app-paths'
 import type { SizeLimit } from '../../../../../types'
-
-const NEXT_PRIVATE_GLOBAL_WAIT_UNTIL = Symbol.for(
-  '__next_private_global_wait_until__'
-)
-
-// @ts-ignore
-globalThis[NEXT_PRIVATE_GLOBAL_WAIT_UNTIL] =
-  // @ts-ignore
-  globalThis[NEXT_PRIVATE_GLOBAL_WAIT_UNTIL] || []
+import { internal_getCurrentFunctionWaitUntil } from '../../../../server/web/internal-edge-wait-until'
 
 export function getRender({
   dev,
@@ -161,10 +153,10 @@ export function getRender({
     const result = await extendedRes.toResponse()
 
     if (event && event.waitUntil) {
-      event.waitUntil(
-        // @ts-ignore
-        Promise.all([...globalThis[NEXT_PRIVATE_GLOBAL_WAIT_UNTIL]])
-      )
+      const waitUntilPromise = internal_getCurrentFunctionWaitUntil()
+      if (waitUntilPromise) {
+        event.waitUntil(waitUntilPromise)
+      }
     }
 
     // fetchMetrics is attached to the web request that going through the server,

packages/next/src/server/web/edge-route-module-wrapper.ts

@@ -13,6 +13,7 @@ import { RouteMatcher } from '../future/route-matchers/route-matcher'
 import { removeTrailingSlash } from '../../shared/lib/router/utils/remove-trailing-slash'
 import { removePathPrefix } from '../../shared/lib/router/utils/remove-path-prefix'
 import type { NextFetchEvent } from './spec-extension/fetch-event'
+import { internal_getCurrentFunctionWaitUntil } from './internal-edge-wait-until'
 
 type WrapOptions = Partial<Pick<AdapterOptions, 'page'>>
 
@@ -114,9 +115,12 @@ export class EdgeRouteModuleWrapper {
     // Get the response from the handler.
     const res = await this.routeModule.handle(request, context)
 
+    const waitUntilPromises = [internal_getCurrentFunctionWaitUntil()]
     if (context.renderOpts.waitUntil) {
-      evt.waitUntil(context.renderOpts.waitUntil)
+      waitUntilPromises.push(context.renderOpts.waitUntil)
     }
+    evt.waitUntil(Promise.all(waitUntilPromises))
+
     return res
   }
 }

packages/next/src/server/web/internal-edge-wait-until.ts

@@ -0,0 +1,59 @@
+// An internal module to expose the "waitUntil" API to Edge SSR and Edge Route Handler functions.
+// This is highly experimental and subject to change.
+
+// We still need a global key to bypass Webpack's layering of modules.
+const GLOBAL_KEY = Symbol.for('__next_internal_waitUntil__')
+
+const state: {
+  waitUntilCounter: number
+  waitUntilResolve: () => void
+  waitUntilPromise: Promise<void> | null
+} =
+  // @ts-ignore
+  globalThis[GLOBAL_KEY] ||
+  // @ts-ignore
+  (globalThis[GLOBAL_KEY] = {
+    waitUntilCounter: 0,
+    waitUntilResolve: undefined,
+    waitUntilPromise: null,
+  })
+
+// No matter how many concurrent requests are being handled, we want to make sure
+// that the final promise is only resolved once all of the waitUntil promises have
+// settled.
+function resolveOnePromise() {
+  state.waitUntilCounter--
+  if (state.waitUntilCounter === 0) {
+    state.waitUntilResolve()
+    state.waitUntilPromise = null
+  }
+}
+
+export function internal_getCurrentFunctionWaitUntil() {
+  return state.waitUntilPromise
+}
+
+export function internal_runWithWaitUntil<T>(fn: () => T): T {
+  const result = fn()
+  if (
+    result &&
+    typeof result === 'object' &&
+    'then' in result &&
+    'finally' in result &&
+    typeof result.then === 'function' &&
+    typeof result.finally === 'function'
+  ) {
+    if (!state.waitUntilCounter) {
+      // Create the promise for the next batch of waitUntil calls.
+      state.waitUntilPromise = new Promise<void>((resolve) => {
+        state.waitUntilResolve = resolve
+      })
+    }
+    state.waitUntilCounter++
+    return result.finally(() => {
+      resolveOnePromise()
+    })
+  }
+
+  return result
+}