Scope request store in render (vercel#71109)

Stacked on vercel#71107 Scopes the request store to the render callsites within `renderToStream`. This is currently duplicative because we have an outer requestStore scope but we want to get rid of this after we remove the dependency on it from the prerender side. This just sets us up to do that in the future.
sanity-io · Oct 11, 2024 · a40c1d6 · a40c1d6
1 parent 1b89b4b
commit a40c1d6
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 30 deletions.
diff --git a/packages/next/src/server/app-render/app-render.tsx b/packages/next/src/server/app-render/app-render.tsx
@@ -1194,6 +1194,7 @@ async function renderToHTMLOrFlightImpl(
           const notFoundLoaderTree = createNotFoundLoaderTree(loaderTree)
           res.statusCode = 404
           const stream = await renderToStreamWithTracing(
+            requestStore,
             req,
             res,
             ctx,
@@ -1219,6 +1220,7 @@ async function renderToHTMLOrFlightImpl(
     }
 
     const stream = await renderToStreamWithTracing(
+      requestStore,
       req,
       res,
       ctx,
@@ -1345,6 +1347,7 @@ export const renderToHTMLOrFlight: AppPageRender = (
 }
 
 async function renderToStream(
+  requestStore: RequestStore,
   req: BaseNextRequest,
   res: BaseNextResponse,
   ctx: AppRenderContext,
@@ -1435,9 +1438,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,
         {
@@ -1475,7 +1486,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,
           <App
             reactServerStream={reactServerResult.tee()}
             preinitScripts={preinitScripts}
@@ -1512,7 +1525,9 @@ async function renderToStream(
     const renderToReadableStream = require('react-dom/server.edge')
       .renderToReadableStream as (typeof import('react-dom/server.edge'))['renderToReadableStream']
 
-    const htmlStream = await renderToReadableStream(
+    const htmlStream = await workUnitAsyncStorage.run(
+      requestStore,
+      renderToReadableStream,
       <App
         reactServerStream={reactServerResult.tee()}
         preinitScripts={preinitScripts}
@@ -1634,9 +1649,17 @@ async function renderToStream(
       '/_not-found/page'
     )
 
-    const errorRSCPayload = await getErrorRSCPayload(tree, ctx, errorType)
+    const errorRSCPayload = await workUnitAsyncStorage.run(
+      requestStore,
+      getErrorRSCPayload,
+      tree,
+      ctx,
+      errorType
+    )
 
-    const errorServerStream = ComponentMod.renderToReadableStream(
+    const errorServerStream = workUnitAsyncStorage.run(
+      requestStore,
+      ComponentMod.renderToReadableStream,
       errorRSCPayload,
       clientReferenceManifest.clientModules,
       {
@@ -1651,23 +1674,27 @@ async function renderToStream(
     }
 
     try {
-      const fizzStream = await renderToInitialFizzStream({
-        ReactDOMServer: require('react-dom/server.edge'),
-        element: (
-          <AppWithoutContext
-            reactServerStream={errorServerStream}
-            preinitScripts={errorPreinitScripts}
-            clientReferenceManifest={clientReferenceManifest}
-            nonce={ctx.nonce}
-          />
-        ),
-        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: (
+            <AppWithoutContext
+              reactServerStream={errorServerStream}
+              preinitScripts={errorPreinitScripts}
+              clientReferenceManifest={clientReferenceManifest}
+              nonce={ctx.nonce}
+            />
+          ),
+          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
@@ -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,