diff --git a/packages/next/src/build/webpack/loaders/next-edge-ssr-loader/render.ts b/packages/next/src/build/webpack/loaders/next-edge-ssr-loader/render.ts index 8e71beaf8aefd..1da4bcc8d16df 100644 --- a/packages/next/src/build/webpack/loaders/next-edge-ssr-loader/render.ts +++ b/packages/next/src/build/webpack/loaders/next-edge-ssr-loader/render.ts @@ -157,17 +157,27 @@ export function getRender({ request: NextRequestHint, event?: NextFetchEvent ) { + const isAfterEnabled = !!process.env.__NEXT_AFTER + const extendedReq = new WebNextRequest(request) const extendedRes = new WebNextResponse( undefined, // tracking onClose adds overhead, so only do it if `experimental.after` is on. - !!process.env.__NEXT_AFTER + isAfterEnabled ) handler(extendedReq, extendedRes) const result = await extendedRes.toResponse() if (event?.waitUntil) { + if (isAfterEnabled) { + // make sure that NextRequestHint's awaiter stays open long enough + // for late `waitUntil`s called during streaming to get picked up. + event.waitUntil( + new Promise((resolve) => extendedRes.onClose(resolve)) + ) + } + // TODO(after): // remove `internal_runWithWaitUntil` and the `internal-edge-wait-until` module // when consumers switch to `unstable_after`. diff --git a/packages/next/src/server/lib/awaiter.test.ts b/packages/next/src/server/lib/awaiter.test.ts new file mode 100644 index 0000000000000..89af27bb052b4 --- /dev/null +++ b/packages/next/src/server/lib/awaiter.test.ts @@ -0,0 +1,81 @@ +import { InvariantError } from '../../shared/lib/invariant-error' +import { AwaiterMulti, AwaiterOnce } from './awaiter' + +describe('AwaiterOnce/AwaiterMulti', () => { + describe.each([ + { name: 'AwaiterMulti', impl: AwaiterMulti }, + { name: 'AwaiterOnce', impl: AwaiterOnce }, + ])('$name', ({ impl: AwaiterImpl }) => { + it('awaits promises added by other promises', async () => { + const awaiter = new AwaiterImpl() + + const MAX_DEPTH = 5 + const promises: TrackedPromise[] = [] + + const waitUntil = (promise: Promise) => { + promises.push(trackPromiseSettled(promise)) + awaiter.waitUntil(promise) + } + + const makeNestedPromise = async () => { + if (promises.length >= MAX_DEPTH) { + return + } + await sleep(100) + waitUntil(makeNestedPromise()) + } + + waitUntil(makeNestedPromise()) + + await awaiter.awaiting() + + for (const promise of promises) { + expect(promise.isSettled).toBe(true) + } + }) + + it('calls onError for rejected promises', async () => { + const onError = jest.fn() + const awaiter = new AwaiterImpl({ onError }) + + awaiter.waitUntil(Promise.reject('error 1')) + awaiter.waitUntil( + sleep(100).then(() => awaiter.waitUntil(Promise.reject('error 2'))) + ) + + await awaiter.awaiting() + + expect(onError).toHaveBeenCalledWith('error 1') + expect(onError).toHaveBeenCalledWith('error 2') + }) + }) +}) + +describe('AwaiterOnce', () => { + it("does not allow calling waitUntil after it's been awaited", async () => { + const awaiter = new AwaiterOnce() + awaiter.waitUntil(Promise.resolve(1)) + await awaiter.awaiting() + expect(() => awaiter.waitUntil(Promise.resolve(2))).toThrow(InvariantError) + }) +}) + +type TrackedPromise = Promise & { isSettled: boolean } + +function trackPromiseSettled(promise: Promise): TrackedPromise { + const tracked = promise as TrackedPromise + tracked.isSettled = false + tracked.then( + () => { + tracked.isSettled = true + }, + () => { + tracked.isSettled = true + } + ) + return tracked +} + +function sleep(duration: number) { + return new Promise((resolve) => setTimeout(resolve, duration)) +} diff --git a/test/e2e/app-dir/next-after-app/utils/awaiter.ts b/packages/next/src/server/lib/awaiter.ts similarity index 96% rename from test/e2e/app-dir/next-after-app/utils/awaiter.ts rename to packages/next/src/server/lib/awaiter.ts index f4d3214b57868..7f5c32c441302 100644 --- a/test/e2e/app-dir/next-after-app/utils/awaiter.ts +++ b/packages/next/src/server/lib/awaiter.ts @@ -1,4 +1,4 @@ -import { InvariantError } from 'next/dist/shared/lib/invariant-error' +import { InvariantError } from '../../shared/lib/invariant-error' /** * Provides a `waitUntil` implementation which gathers promises to be awaited later (via {@link AwaiterMulti.awaiting}). diff --git a/packages/next/src/server/web/adapter.ts b/packages/next/src/server/web/adapter.ts index 7d5fde2d2b320..dcd388246432c 100644 --- a/packages/next/src/server/web/adapter.ts +++ b/packages/next/src/server/web/adapter.ts @@ -422,7 +422,7 @@ export async function adapter( return { response: finalResponse, - waitUntil: Promise.all(event[waitUntilSymbol]), + waitUntil: event[waitUntilSymbol](), fetchMetrics: request.fetchMetrics, } } diff --git a/packages/next/src/server/web/edge-route-module-wrapper.ts b/packages/next/src/server/web/edge-route-module-wrapper.ts index 244281e531765..cfc90c7ce3fa6 100644 --- a/packages/next/src/server/web/edge-route-module-wrapper.ts +++ b/packages/next/src/server/web/edge-route-module-wrapper.ts @@ -148,6 +148,13 @@ export class EdgeRouteModuleWrapper { const trackedBody = trackStreamConsumed(res.body, () => _closeController.dispatchClose() ) + + // make sure that NextRequestHint's awaiter stays open long enough + // for `waitUntil`s called late during streaming to get picked up. + evt.waitUntil( + new Promise((resolve) => _closeController.onClose(resolve)) + ) + res = new Response(trackedBody, { status: res.status, statusText: res.statusText, diff --git a/packages/next/src/server/web/spec-extension/fetch-event.ts b/packages/next/src/server/web/spec-extension/fetch-event.ts index 8f7776ac00dea..43d509bcdd3fa 100644 --- a/packages/next/src/server/web/spec-extension/fetch-event.ts +++ b/packages/next/src/server/web/spec-extension/fetch-event.ts @@ -1,14 +1,22 @@ +import { AwaiterOnce } from '../../lib/awaiter' import { PageSignatureError } from '../error' import type { NextRequest } from './request' const responseSymbol = Symbol('response') const passThroughSymbol = Symbol('passThrough') +const awaiterSymbol = Symbol('awaiter') + export const waitUntilSymbol = Symbol('waitUntil') class FetchEvent { - readonly [waitUntilSymbol]: Promise[] = []; [responseSymbol]?: Promise; - [passThroughSymbol] = false + [passThroughSymbol] = false; + + [awaiterSymbol] = new AwaiterOnce(); + + [waitUntilSymbol] = () => { + return this[awaiterSymbol].awaiting() + } // eslint-disable-next-line @typescript-eslint/no-useless-constructor constructor(_request: Request) {} @@ -24,7 +32,7 @@ class FetchEvent { } waitUntil(promise: Promise): void { - this[waitUntilSymbol].push(promise) + this[awaiterSymbol].waitUntil(promise) } } diff --git a/test/e2e/app-dir/next-after-app/index.test.ts b/test/e2e/app-dir/next-after-app/index.test.ts index 1df6dc29cb5e5..7cab04a5ad818 100644 --- a/test/e2e/app-dir/next-after-app/index.test.ts +++ b/test/e2e/app-dir/next-after-app/index.test.ts @@ -380,10 +380,7 @@ describe.each(runtimes)('unstable_after() in %s runtime', (runtimeValue) => { return cleanup } - /* eslint-disable jest/no-standalone-expect */ - const it_failingForEdge = runtimeValue === 'edge' ? it.failing : it - - it_failingForEdge('during render', async () => { + it('during render', async () => { const cleanup = await setup() try { const response = await next.fetch('/delay-deep') @@ -401,27 +398,23 @@ describe.each(runtimes)('unstable_after() in %s runtime', (runtimeValue) => { } }) - it_failingForEdge( - 'in a route handler that streams a response', - async () => { - const cleanup = await setup() - try { - const response = await next.fetch('/route-streaming') - expect(response.status).toBe(200) - await response.text() - await retry(() => { - expect(getLogs()).toContainEqual('simulated-invocation :: end') - }, 10_000) - - expect(getLogs()).toContainEqual({ - source: '[route handler] /route-streaming - after', - }) - } finally { - await cleanup() - } + it('in a route handler that streams a response', async () => { + const cleanup = await setup() + try { + const response = await next.fetch('/route-streaming') + expect(response.status).toBe(200) + await response.text() + await retry(() => { + expect(getLogs()).toContainEqual('simulated-invocation :: end') + }, 10_000) + + expect(getLogs()).toContainEqual({ + source: '[route handler] /route-streaming - after', + }) + } finally { + await cleanup() } - ) - /* eslint-enable jest/no-standalone-expect */ + }) }) } diff --git a/test/e2e/app-dir/next-after-app/utils/simulated-invocation.js b/test/e2e/app-dir/next-after-app/utils/simulated-invocation.js index 02e242fb6d97f..0afb788a57db1 100644 --- a/test/e2e/app-dir/next-after-app/utils/simulated-invocation.js +++ b/test/e2e/app-dir/next-after-app/utils/simulated-invocation.js @@ -1,5 +1,5 @@ import { requestAsyncStorage } from 'next/dist/client/components/request-async-storage.external' -import { AwaiterOnce } from './awaiter' +import { AwaiterOnce } from 'next/dist/server/lib/awaiter' import { cliLog } from './log' // replaced in tests