From b845eff4794ef0059e8c7e5eb63769202788fb7b Mon Sep 17 00:00:00 2001
From: Josh Story <story@hey.com>
Date: Fri, 30 Sep 2022 13:33:59 -0700
Subject: [PATCH] use a stack for currentResources

---
 .../src/server/ReactDOMFloatServer.js         | 22 +++++++---------
 .../src/server/ReactDOMServerFormatConfig.js  | 25 +++++++++++++------
 packages/react-server/src/ReactFizzServer.js  |  5 ++--
 3 files changed, 29 insertions(+), 23 deletions(-)

diff --git a/packages/react-dom-bindings/src/server/ReactDOMFloatServer.js b/packages/react-dom-bindings/src/server/ReactDOMFloatServer.js
index 210093e5e4bbf..3fa4b8fc0fe9a 100644
--- a/packages/react-dom-bindings/src/server/ReactDOMFloatServer.js
+++ b/packages/react-dom-bindings/src/server/ReactDOMFloatServer.js
@@ -100,13 +100,16 @@ export function mergeBoundaryResources(
 }
 
 let currentResources: null | Resources = null;
+let currentResourcesStack = [];
 
 let previousDispatcher = null;
 export function prepareToRenderResources(resources: Resources) {
+  currentResourcesStack.push(currentResources);
   currentResources = resources;
+}
 
-  previousDispatcher = ReactDOMSharedInternals.Dispatcher.current;
-  ReactDOMSharedInternals.Dispatcher.current = Dispatcher;
+export function finishRenderingResources() {
+  currentResources = currentResourcesStack.pop();
 }
 
 export function setCurrentlyRenderingBoundaryResourcesTarget(
@@ -116,12 +119,10 @@ export function setCurrentlyRenderingBoundaryResourcesTarget(
   resources.boundaryResources = boundaryResources;
 }
 
-export function finishRenderingResources() {
-  currentResources = null;
-
-  ReactDOMSharedInternals.Dispatcher.current = previousDispatcher;
-  previousDispatcher = null;
-}
+export const ReactDOMServerDispatcher = {
+  preload,
+  preinit,
+};
 
 type PreloadAs = ResourceType;
 type PreloadOptions = {as: PreloadAs, crossOrigin?: string};
@@ -574,8 +575,3 @@ export function hoistResourcesToRoot(
   });
   boundaryResources.clear();
 }
-
-const Dispatcher = {
-  preload,
-  preinit,
-};
diff --git a/packages/react-dom-bindings/src/server/ReactDOMServerFormatConfig.js b/packages/react-dom-bindings/src/server/ReactDOMServerFormatConfig.js
index b1b3705a593a3..7c236326e1837 100644
--- a/packages/react-dom-bindings/src/server/ReactDOMServerFormatConfig.js
+++ b/packages/react-dom-bindings/src/server/ReactDOMServerFormatConfig.js
@@ -64,6 +64,7 @@ import {
   prepareToRenderResources,
   finishRenderingResources,
   resourcesFromLink,
+  ReactDOMServerDispatcher,
 } from './ReactDOMFloatServer';
 export {
   createResources,
@@ -73,6 +74,22 @@ export {
   hoistResourcesToRoot,
 } from './ReactDOMFloatServer';
 
+import ReactDOMSharedInternals from 'shared/ReactDOMSharedInternals';
+const ReactDOMCurrentDispatcher = ReactDOMSharedInternals.Dispatcher;
+
+export function prepareToRender(resources: Resources): mixed {
+  prepareToRenderResources(resources);
+
+  let previousHostDispatcher = ReactDOMCurrentDispatcher.current;
+  ReactDOMCurrentDispatcher.current = ReactDOMServerDispatcher;
+  return previousHostDispatcher;
+}
+
+export function cleanupAfterRender(previousDispatcher: mixed) {
+  finishRenderingResources();
+  ReactDOMCurrentDispatcher.current = previousDispatcher;
+}
+
 // Used to distinguish these contexts from ones used in other renderers.
 // E.g. this can be used to distinguish legacy renderers from this modern one.
 export const isPrimaryRenderer = true;
@@ -2722,11 +2739,3 @@ function writeStyleResourceAttribute(
     stringToChunk(escapeJSObjectForInstructionScripts(attributeValue)),
   );
 }
-
-export function prepareToRender(resources: Resources) {
-  prepareToRenderResources(resources);
-}
-
-export function cleanupAfterRender() {
-  finishRenderingResources();
-}
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
index dfc742bbc0c37..ab9415d88d369 100644
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -1827,7 +1827,6 @@ function finishedTask(
 }
 
 function retryTask(request: Request, task: Task): void {
-  prepareToRender(request.resources);
   if (enableFloat) {
     const blockedBoundary = task.blockedBoundary;
     setCurrentlyRenderingBoundaryResourcesTarget(
@@ -1892,7 +1891,6 @@ function retryTask(request: Request, task: Task): void {
     if (__DEV__) {
       currentTaskInDEV = prevTaskInDEV;
     }
-    cleanupAfterRender();
   }
 }
 
@@ -1903,6 +1901,7 @@ export function performWork(request: Request): void {
   const prevContext = getActiveContext();
   const prevDispatcher = ReactCurrentDispatcher.current;
   ReactCurrentDispatcher.current = Dispatcher;
+  const previousHostDispatcher = prepareToRender(request.resources);
   let prevGetCurrentStackImpl;
   if (__DEV__) {
     prevGetCurrentStackImpl = ReactDebugCurrentFrame.getCurrentStack;
@@ -1927,6 +1926,8 @@ export function performWork(request: Request): void {
   } finally {
     setCurrentResponseState(prevResponseState);
     ReactCurrentDispatcher.current = prevDispatcher;
+    cleanupAfterRender(previousHostDispatcher);
+
     if (__DEV__) {
       ReactDebugCurrentFrame.getCurrentStack = prevGetCurrentStackImpl;
     }