From f7a088a344e99f504073dae92c1d0d44fc29784b Mon Sep 17 00:00:00 2001 From: Josh Story Date: Thu, 10 Oct 2024 08:45:15 -0700 Subject: [PATCH] Add RequestStore around renders Today we scope the entire work store with a RequestStore. This is because prerenders still have some dependency on this store existing. We want to get rid of this but to do so we need to ensure there is a RequestStore around the renders inside renderToStream. This change adds the necessary scoping inside renderToStream. Today this is redundant but it sets us upto be able to remove the outer RequestStore scope. In addition I removed the withRequestStore helper. We generally don't want to construct a new RequestStore for everything we will scope in a RequestStore scope so we can just create the store and then wrap the necessary functions with workUnitAsyncStorage.run... --- .../next/src/server/app-render/app-render.tsx | 73 +++++++++++++------ .../server/route-modules/app-route/module.ts | 24 ++++-- 2 files changed, 67 insertions(+), 30 deletions(-) diff --git a/packages/next/src/server/app-render/app-render.tsx b/packages/next/src/server/app-render/app-render.tsx index 6a3806e7f6c37..5c6cdf2fd35e5 100644 --- a/packages/next/src/server/app-render/app-render.tsx +++ b/packages/next/src/server/app-render/app-render.tsx @@ -1168,6 +1168,7 @@ async function renderToHTMLOrFlightImpl( const notFoundLoaderTree = createNotFoundLoaderTree(loaderTree) res.statusCode = 404 const stream = await renderToStreamWithTracing( + requestStore, req, res, ctx, @@ -1193,6 +1194,7 @@ async function renderToHTMLOrFlightImpl( } const stream = await renderToStreamWithTracing( + requestStore, req, res, ctx, @@ -1319,6 +1321,7 @@ export const renderToHTMLOrFlight: AppPageRender = ( } async function renderToStream( + requestStore: RequestStore, req: BaseNextRequest, res: BaseNextResponse, ctx: AppRenderContext, @@ -1409,9 +1412,17 @@ async function renderToStream( try { // This is a dynamic render. We don't do dynamic tracking because we're not prerendering - const RSCPayload = await getRSCPayload(tree, ctx, res.statusCode === 404) + const RSCPayload = await workUnitAsyncStorage.run( + requestStore, + getRSCPayload, + tree, + ctx, + res.statusCode === 404 + ) reactServerResult = new ReactServerResult( - ComponentMod.renderToReadableStream( + workUnitAsyncStorage.run( + requestStore, + ComponentMod.renderToReadableStream, RSCPayload, clientReferenceManifest.clientModules, { @@ -1449,7 +1460,9 @@ async function renderToStream( const resume = require('react-dom/server.edge') .resume as (typeof import('react-dom/server.edge'))['resume'] - const htmlStream = await resume( + const htmlStream = await workUnitAsyncStorage.run( + requestStore, + resume, - ), - streamOptions: { - nonce: ctx.nonce, - // Include hydration scripts in the HTML - bootstrapScripts: [errorBootstrapScript], - formState, - }, - }) + const fizzStream = await workUnitAsyncStorage.run( + requestStore, + renderToInitialFizzStream, + { + ReactDOMServer: require('react-dom/server.edge'), + element: ( + + ), + streamOptions: { + nonce: ctx.nonce, + // Include hydration scripts in the HTML + bootstrapScripts: [errorBootstrapScript], + formState, + }, + } + ) /** * Rules of Static & Dynamic HTML: diff --git a/packages/next/src/server/route-modules/app-route/module.ts b/packages/next/src/server/route-modules/app-route/module.ts index 61e6afa2394c5..fb69109dda946 100644 --- a/packages/next/src/server/route-modules/app-route/module.ts +++ b/packages/next/src/server/route-modules/app-route/module.ts @@ -41,7 +41,7 @@ import { } from '../../app-render/work-async-storage.external' import { workUnitAsyncStorage, - type WorkUnitStore, + type RequestStore, type PrerenderStore, } from '../../app-render/work-unit-async-storage.external' import { @@ -276,7 +276,10 @@ export class AppRouteRouteModule extends RouteModule< handler: AppRouteHandlerFn, actionStore: ActionStore, workStore: WorkStore, - workUnitStore: WorkUnitStore, + // @TODO refactor to not take this argument but instead construct the RequestStore + // inside this function. Right now we get passed a RequestStore even when + // we're going to do a prerender. We should probably just split do up into prexecute and execute + requestStore: RequestStore, implicitTags: string[], request: NextRequest, context: AppRouteRouteHandlerContext @@ -503,7 +506,12 @@ export class AppRouteRouteModule extends RouteModule< ) } } else { - res = await handler(request, handlerContext) + res = await workUnitAsyncStorage.run( + requestStore, + handler, + request, + handlerContext + ) } } catch (err) { if (isRedirectError(err)) { @@ -518,8 +526,10 @@ export class AppRouteRouteModule extends RouteModule< // Let's append any cookies that were added by the // cookie API. - if (workUnitStore.type === 'request') { - appendMutableCookies(headers, workUnitStore.mutableCookies) + // TODO leaving the gate here b/c it indicates that we we might not actually want to do this + // on every `do` call. During prerender there should be no mutableCookies because + if (requestStore.type === 'request') { + appendMutableCookies(headers, requestStore.mutableCookies) } // Return the redirect response. @@ -567,8 +577,8 @@ export class AppRouteRouteModule extends RouteModule< // here. const headers = new Headers(res.headers) if ( - workUnitStore.type === 'request' && - appendMutableCookies(headers, workUnitStore.mutableCookies) + requestStore.type === 'request' && + appendMutableCookies(headers, requestStore.mutableCookies) ) { return new Response(res.body, { status: res.status,