remove runInContext callbacks from stream-ops (#90609)

The `runInContext` pattern that's used all over `stream-ops` is weird
and feels very unnecessary -- we're passing in callbacks that just get
invoked immediately, and all they do is call `AsyncLocalStorage.run`. We
should just use the standard `AsyncLocalStorage.run` method outside
instead. I commented about getting rid of it somewhere on the original
node streams PR, but that got lost when it was was split up.

I believe the original reasoning for addding it at all was that some of
the node-streams implementations needed to preserve async context for
callbacks -- see eg
[here](https://github.com/vercel/next.js/blob/558f4248e5e5641d78b865a291c099791460d486/packages/next/src/server/app-render/node-stream-tee.ts#L244).
But if that's the motivation, then the current `runInContext` pattern
**is wrong anyway** -- sure, it preserves `workUnitAsyncStorage`, but
doesn't preserve `workAsyncStorage` and whetever other ALSes may be
present, so it's not semantically correct. If this is the goal, then the
relevant callbacks should use `bindSnapshot` (i.e.
`AsyncLocalStorage.bind`) instead, because it preserves the entire
context, not just one ALS.

(Note that technically the `runInContext` pattern is salvageable if we
just always pass `AsyncLocalStorage.snapshot()`, but then i see no
reason why the caller should have to do that. `AsyncLocalStorage.bind`
is cleaner)

Note that this'll require changes in #89859

packages/next/src/server/app-render/app-render.tsx

@@ -723,7 +723,9 @@ async function generateDynamicFlightRenderResult(
     options
   )
 
-  const flightStream = renderToFlightStream(
+  const flightStream = workUnitAsyncStorage.run(
+    requestStore,
+    renderToFlightStream,
     ctx.componentMod,
     rscPayload,
     clientModules,
@@ -732,8 +734,7 @@ async function generateDynamicFlightRenderResult(
       temporaryReferences: options?.temporaryReferences,
       filterStackFrame,
       debugChannel: debugChannel?.serverSide,
-    },
-    (fn) => workUnitAsyncStorage.run(requestStore, fn)
+    }
   )
 
   return new FlightRenderResult(
@@ -902,15 +903,16 @@ async function stagedRenderToReadableStreamWithoutCachesInDev(
   return await runInSequentialTasks(
     () => {
       stageController.advanceStage(RenderStage.Static)
-      return renderToFlightStream(
+      return workUnitAsyncStorage.run(
+        requestStore,
+        renderToFlightStream,
         ctx.componentMod,
         rscPayload,
         clientModules,
         {
           ...options,
           environmentName,
-        },
-        (fn) => workUnitAsyncStorage.run(requestStore, fn)
+        }
       )
     },
     () => {
@@ -2906,7 +2908,9 @@ async function renderToStream(
         }
 
         reactServerResult = new ReactServerResult(
-          renderToFlightStream(
+          workUnitAsyncStorage.run(
+            requestStore,
+            renderToFlightStream,
             ctx.componentMod,
             RSCPayload,
             clientModules,
@@ -2914,8 +2918,7 @@ async function renderToStream(
               filterStackFrame,
               onError: serverComponentsErrorHandler,
               debugChannel: debugChannel?.serverSide,
-            },
-            (fn) => workUnitAsyncStorage.run(requestStore, fn)
+            }
           )
         )
       }
@@ -2966,12 +2969,14 @@ async function renderToStream(
             tracingMetadata: tracingMetadata,
           })
 
-          const { stream: htmlStream, allReady } = await resumeToFizzStream(
-            resumeAppElement,
-            postponed,
-            { onError: htmlRendererErrorHandler, nonce },
-            (fn) => workUnitAsyncStorage.run(requestStore, fn)
-          )
+          const { stream: htmlStream, allReady } =
+            await workUnitAsyncStorage.run(
+              requestStore,
+              resumeToFizzStream,
+              resumeAppElement,
+              postponed,
+              { onError: htmlRendererErrorHandler, nonce }
+            )
 
           // End the render span only after React completed rendering (including anything inside Suspense boundaries)
           allReady.finally(() => {
@@ -3031,10 +3036,11 @@ async function renderToStream(
         formState,
       }
 
-      const { stream: htmlStream, allReady } = await renderToFizzStream(
+      const { stream: htmlStream, allReady } = await workUnitAsyncStorage.run(
+        requestStore,
+        renderToFizzStream,
         appElement,
-        fizzOptions,
-        (fn) => workUnitAsyncStorage.run(requestStore, fn)
+        fizzOptions
       )
 
       // End the render span only after React completed rendering (including anything inside Suspense boundaries)
@@ -3136,15 +3142,16 @@ async function renderToStream(
           errorType
         )
 
-        errorServerStream = renderToFlightStream(
+        errorServerStream = workUnitAsyncStorage.run(
+          requestStore,
+          renderToFlightStream,
           ctx.componentMod,
           errorRSCPayload,
           clientModules,
           {
             filterStackFrame,
             onError: serverComponentsErrorHandler,
-          },
-          (fn) => workUnitAsyncStorage.run(requestStore, fn)
+          }
         )
 
         if (reactServerResult === null) {
@@ -3163,7 +3170,9 @@ async function renderToStream(
           supportsDynamicResponse !== true || !!shouldWaitOnAllReady
 
         const { stream: errorHtmlStream, allReady: errorAllReady } =
-          await renderToFizzStream(
+          await workUnitAsyncStorage.run(
+            requestStore,
+            renderToFizzStream,
             <ErrorApp
               reactServerStream={errorServerStream}
               ServerInsertedHTMLProvider={ServerInsertedHTMLProvider}
@@ -3176,8 +3185,7 @@ async function renderToStream(
               bootstrapScriptContent,
               bootstrapScripts: [errorBootstrapScript],
               formState,
-            },
-            (fn) => workUnitAsyncStorage.run(requestStore, fn)
+            }
           )
 
         // End the render span only after React completed rendering (including anything inside Suspense boundaries)
@@ -3310,20 +3318,23 @@ async function renderWithRestartOnCacheMissInDev(
       initialStageController.advanceStage(RenderStage.EarlyStatic)
       startTime = performance.now() + performance.timeOrigin
 
-      const streamPair = renderToFlightStream(
-        ComponentMod,
-        initialRscPayload,
-        clientModules,
-        {
-          onError,
-          environmentName,
-          startTime,
-          filterStackFrame,
-          debugChannel: debugChannel?.serverSide,
-          signal: initialReactController.signal,
-        },
-        (fn) => workUnitAsyncStorage.run(requestStore, fn)
-      ).tee()
+      const streamPair = workUnitAsyncStorage
+        .run(
+          requestStore,
+          renderToFlightStream,
+          ComponentMod,
+          initialRscPayload,
+          clientModules,
+          {
+            onError,
+            environmentName,
+            startTime,
+            filterStackFrame,
+            debugChannel: debugChannel?.serverSide,
+            signal: initialReactController.signal,
+          }
+        )
+        .tee()
 
       // If we abort the render, we want to reject the stage-dependent promises as well.
       // Note that we want to install this listener after the render is started
@@ -3457,19 +3468,22 @@ async function renderWithRestartOnCacheMissInDev(
       finalStageController.advanceStage(RenderStage.EarlyStatic)
       startTime = performance.now() + performance.timeOrigin
 
-      const streamPair = renderToFlightStream(
-        ComponentMod,
-        finalRscPayload,
-        clientModules,
-        {
-          onError,
-          environmentName,
-          startTime,
-          filterStackFrame,
-          debugChannel: debugChannel?.serverSide,
-        },
-        (fn) => workUnitAsyncStorage.run(requestStore, fn)
-      ).tee()
+      const streamPair = workUnitAsyncStorage
+        .run(
+          requestStore,
+          renderToFlightStream,
+          ComponentMod,
+          finalRscPayload,
+          clientModules,
+          {
+            onError,
+            environmentName,
+            startTime,
+            filterStackFrame,
+            debugChannel: debugChannel?.serverSide,
+          }
+        )
+        .tee()
 
       return {
         stream: streamPair[0],
@@ -5782,7 +5796,9 @@ async function prerenderToStream(
           )
         )
 
-      const { stream: htmlStream } = await renderToFizzStream(
+      const { stream: htmlStream } = await workUnitAsyncStorage.run(
+        prerenderLegacyStore,
+        renderToFizzStream,
         // eslint-disable-next-line @next/internal/no-ambiguous-jsx
         <App
           reactServerStream={reactServerResult.asUnclosingStream()}
@@ -5797,8 +5813,7 @@ async function prerenderToStream(
           onError: htmlRendererErrorHandler,
           nonce,
           bootstrapScripts: [bootstrapScript],
-        },
-        (fn) => workUnitAsyncStorage.run(prerenderLegacyStore, fn)
+        }
       )
 
       if (shouldGenerateStaticFlightData(workStore)) {
@@ -5936,24 +5951,22 @@ async function prerenderToStream(
       errorType
     )
 
-    // Keep prerender-legacy async storage stable across node stream event
-    // callbacks and pipeable renderer hooks.
-    const runInLegacyContext = <T,>(fn: () => T): T =>
-      workUnitAsyncStorage.run(prerenderLegacyStore, fn)
-
-    const errorServerStream = renderToFlightStream(
+    const errorServerStream = workUnitAsyncStorage.run(
+      prerenderLegacyStore,
+      renderToFlightStream,
       ComponentMod,
       errorRSCPayload,
       clientModules,
       {
         filterStackFrame,
         onError: serverComponentsErrorHandler,
-      },
-      runInLegacyContext
+      }
     )
 
     try {
-      const { stream: errorHtmlStream } = await renderToFizzStream(
+      const { stream: errorHtmlStream } = await workUnitAsyncStorage.run(
+        prerenderLegacyStore,
+        renderToFizzStream,
         // eslint-disable-next-line @next/internal/no-ambiguous-jsx
         <ErrorApp
           reactServerStream={errorServerStream}
@@ -5966,8 +5979,7 @@ async function prerenderToStream(
           nonce,
           bootstrapScripts: [errorBootstrapScript],
           formState,
-        },
-        runInLegacyContext
+        }
       )
 
       if (shouldGenerateStaticFlightData(workStore)) {

packages/next/src/server/app-render/stream-ops.web.ts

@@ -147,13 +147,9 @@ export function renderToFlightStream(
   ComponentMod: FlightComponentMod,
   payload: FlightPayload,
   clientModules: FlightClientModules,
-  opts: FlightRenderOptions,
-  runInContext?: <T>(fn: () => T) => T
+  opts: FlightRenderOptions
 ): AnyStream {
-  const run: <T>(fn: () => T) => T = runInContext ?? ((fn) => fn())
-  return run(() =>
-    ComponentMod.renderToReadableStream(payload, clientModules, opts)
-  )
+  return ComponentMod.renderToReadableStream(payload, clientModules, opts)
 }
 
 export async function streamToString(stream: AnyStream): Promise<string> {
@@ -162,28 +158,22 @@ export async function streamToString(stream: AnyStream): Promise<string> {
 
 export async function renderToFizzStream(
   element: React.ReactElement,
-  streamOptions: any,
-  runInContext?: <T>(fn: () => T) => T
+  streamOptions: any
 ): Promise<FizzStreamResult> {
-  const run: <T>(fn: () => T) => T = runInContext ?? ((fn) => fn())
-  const stream = await run(() =>
-    renderToInitialFizzStream({
-      ReactDOMServer: { renderToReadableStream },
-      element,
-      streamOptions,
-    })
-  )
+  const stream = await renderToInitialFizzStream({
+    ReactDOMServer: { renderToReadableStream },
+    element,
+    streamOptions,
+  })
   return { stream, allReady: stream.allReady, abort: undefined }
 }
 
 export async function resumeToFizzStream(
   element: React.ReactElement,
   postponedState: PostponedState,
-  streamOptions: any,
-  runInContext?: <T>(fn: () => T) => T
+  streamOptions: any
 ): Promise<FizzStreamResult> {
-  const run: <T>(fn: () => T) => T = runInContext ?? ((fn) => fn())
-  const stream = await run(() => resume(element, postponedState, streamOptions))
+  const stream = await resume(element, postponedState, streamOptions)
   return { stream, allReady: stream.allReady, abort: undefined }
 }