diff --git a/packages/next/src/server/after/after-context.test.ts b/packages/next/src/server/after/after-context.test.ts index c171b307701b0..290c3979f66c0 100644 --- a/packages/next/src/server/after/after-context.test.ts +++ b/packages/next/src/server/after/after-context.test.ts @@ -4,7 +4,7 @@ import { AsyncLocalStorage } from 'async_hooks' import type { RequestStore } from '../../client/components/request-async-storage.external' import type { AfterContext } from './after-context' -describe('createAfterContext', () => { +describe('AfterContext', () => { // 'async-local-storage.ts' needs `AsyncLocalStorage` on `globalThis` at import time, // so we have to do some contortions here to set it up before running anything else type RASMod = @@ -13,7 +13,7 @@ describe('createAfterContext', () => { type AfterContextMod = typeof import('./after-context') let requestAsyncStorage: RASMod['requestAsyncStorage'] - let createAfterContext: AfterContextMod['createAfterContext'] + let AfterContext: AfterContextMod['AfterContext'] let after: AfterMod['unstable_after'] beforeAll(async () => { @@ -26,7 +26,7 @@ describe('createAfterContext', () => { requestAsyncStorage = RASMod.requestAsyncStorage const AfterContextMod = await import('./after-context') - createAfterContext = AfterContextMod.createAfterContext + AfterContext = AfterContextMod.AfterContext const AfterMod = await import('./after') after = AfterMod.unstable_after @@ -49,7 +49,7 @@ describe('createAfterContext', () => { onCloseCallback = cb }) - const afterContext = createAfterContext({ + const afterContext = new AfterContext({ waitUntil, onClose, cacheScope: undefined, @@ -116,7 +116,7 @@ describe('createAfterContext', () => { onCloseCallback = cb }) - const afterContext = createAfterContext({ + const afterContext = new AfterContext({ waitUntil, onClose, cacheScope: undefined, @@ -164,7 +164,7 @@ describe('createAfterContext', () => { onCloseCallback = cb }) - const afterContext = createAfterContext({ + const afterContext = new AfterContext({ waitUntil, onClose, cacheScope: undefined, @@ -255,7 +255,7 @@ describe('createAfterContext', () => { onCloseCallback = cb }) - const afterContext = createAfterContext({ + const afterContext = new AfterContext({ waitUntil, onClose, cacheScope: undefined, @@ -315,7 +315,7 @@ describe('createAfterContext', () => { throw new Error('onClose is broken for some reason') }) - const afterContext = createAfterContext({ + const afterContext = new AfterContext({ waitUntil, onClose, cacheScope: undefined, @@ -353,7 +353,7 @@ describe('createAfterContext', () => { onCloseCallback = cb }) - const afterContext = createAfterContext({ + const afterContext = new AfterContext({ waitUntil, onClose, cacheScope: undefined, @@ -406,7 +406,7 @@ describe('createAfterContext', () => { const waitUntil = undefined const onClose = jest.fn() - const afterContext = createAfterContext({ + const afterContext = new AfterContext({ waitUntil, onClose, cacheScope: undefined, @@ -436,7 +436,7 @@ describe('createAfterContext', () => { const onClose = undefined - const afterContext = createAfterContext({ + const afterContext = new AfterContext({ waitUntil, onClose, cacheScope: undefined, diff --git a/packages/next/src/server/after/after-context.ts b/packages/next/src/server/after/after-context.ts index 2d75d280ce2c6..8e7bf1945b2e4 100644 --- a/packages/next/src/server/after/after-context.ts +++ b/packages/next/src/server/after/after-context.ts @@ -9,22 +9,13 @@ import type { RequestLifecycleOpts } from '../base-server' import type { AfterCallback, AfterTask } from './after' import { InvariantError } from '../../shared/lib/invariant-error' -export interface AfterContext { - run(requestStore: RequestStore, callback: () => T): T - after(task: AfterTask): void -} - export type AfterContextOpts = { waitUntil: RequestLifecycleOpts['waitUntil'] | undefined onClose: RequestLifecycleOpts['onClose'] | undefined cacheScope: CacheScope | undefined } -export function createAfterContext(opts: AfterContextOpts): AfterContext { - return new AfterContextImpl(opts) -} - -export class AfterContextImpl implements AfterContext { +export class AfterContext { private waitUntil: RequestLifecycleOpts['waitUntil'] | undefined private onClose: RequestLifecycleOpts['onClose'] | undefined private cacheScope: CacheScope | undefined diff --git a/packages/next/src/server/async-storage/with-request-store.ts b/packages/next/src/server/async-storage/with-request-store.ts index 571ff2210415f..596ffff7be7ba 100644 --- a/packages/next/src/server/async-storage/with-request-store.ts +++ b/packages/next/src/server/async-storage/with-request-store.ts @@ -20,7 +20,7 @@ import { import { ResponseCookies, RequestCookies } from '../web/spec-extension/cookies' import { DraftModeProvider } from './draft-mode-provider' import { splitCookiesString } from '../web/utils' -import { createAfterContext, type AfterContext } from '../after/after-context' +import { AfterContext } from '../after/after-context' import type { RequestLifecycleOpts } from '../base-server' function getHeaders(headers: Headers | IncomingHttpHeaders): ReadonlyHeaders { @@ -84,8 +84,6 @@ export const withRequestStore: WithStore = < { req, url, res, renderOpts }: RequestContext, callback: (store: RequestStore) => Result ): Result => { - const [wrapWithAfter, afterContext] = createAfterWrapper(renderOpts) - function defaultOnUpdateCookies(cookies: string[]) { if (res) { res.setHeader('Set-Cookie', cookies) @@ -172,33 +170,33 @@ export const withRequestStore: WithStore = < reactLoadableManifest: renderOpts?.reactLoadableManifest || {}, assetPrefix: renderOpts?.assetPrefix || '', - afterContext, + afterContext: createAfterContext(renderOpts), + } + + if (store.afterContext) { + return store.afterContext.run(store, () => + storage.run(store, callback, store) + ) } - return wrapWithAfter(store, () => storage.run(store, callback, store)) + + return storage.run(store, callback, store) } -function createAfterWrapper( +function createAfterContext( renderOpts: WrapperRenderOpts | undefined -): [ - wrap: (requestStore: RequestStore, callback: () => Result) => Result, - afterContext: AfterContext | undefined, -] { - const isAfterEnabled = renderOpts?.experimental?.after ?? false - if (!renderOpts || !isAfterEnabled) { - return [(_, callback) => callback(), undefined] +): AfterContext | undefined { + if (!isAfterEnabled(renderOpts)) { + return undefined } - const { waitUntil, onClose } = renderOpts - const cacheScope = renderOpts.ComponentMod?.createCacheScope() + const { waitUntil, onClose, ComponentMod } = renderOpts + const cacheScope = ComponentMod?.createCacheScope() - const afterContext = createAfterContext({ - waitUntil, - onClose, - cacheScope, - }) - - const wrap = (requestStore: RequestStore, callback: () => Result) => - afterContext.run(requestStore, callback) + return new AfterContext({ waitUntil, onClose, cacheScope }) +} - return [wrap, afterContext] +function isAfterEnabled( + renderOpts: WrapperRenderOpts | undefined +): renderOpts is WrapperRenderOpts { + return renderOpts?.experimental?.after ?? false }