refactor: avoid unnecessary after wrapper

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,

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<T>(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

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<RequestStore, RequestContext> = <
   { 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<RequestStore, RequestContext> = <
 
     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: <Result>(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 = <Result>(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
 }