reimplement server dispatcher

packages/react-dom-bindings/src/server/ReactDOMFloatServer.js

@@ -8,6 +8,8 @@
  */
 
 import ReactDOMSharedInternals from 'shared/ReactDOMSharedInternals';
+const ReactDOMCurrentDispatcher = ReactDOMSharedInternals.Dispatcher;
+
 import {
   validatePreloadResourceDifference,
   validateStyleResourceDifference,
@@ -99,34 +101,20 @@ export function mergeBoundaryResources(
   source.forEach(resource => target.add(resource));
 }
 
-let currentResources: null | Resources = null;
-
-let previousDispatcher = null;
-export function prepareToRenderResources(resources: Resources) {
-  currentResources = resources;
-
-  previousDispatcher = ReactDOMSharedInternals.Dispatcher.current;
-  ReactDOMSharedInternals.Dispatcher.current = Dispatcher;
-}
-
 export function setCurrentlyRenderingBoundaryResourcesTarget(
   resources: Resources,
   boundaryResources: null | BoundaryResources,
 ) {
   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};
-function preload(href: string, options: PreloadOptions) {
-  if (!currentResources) {
+export function preload(href: string, options: PreloadOptions) {
+  const currentDispatcher = ReactDOMCurrentDispatcher.current;
+  if (!currentDispatcher || !currentDispatcher.currentResources) {
     // While we expect that preload calls are primarily going to be observed
     // during render because effects and events don't run on the server it is
     // still possible that these get called in module scope. This is valid on
@@ -135,6 +123,7 @@ function preload(href: string, options: PreloadOptions) {
     // simply return and do not warn.
     return;
   }
+  const currentResources: Resources = (currentDispatcher.currentResources: any);
   if (__DEV__) {
     validatePreloadArguments(href, options);
   }
@@ -176,8 +165,9 @@ type PreinitOptions = {
   precedence?: string,
   crossOrigin?: string,
 };
-function preinit(href: string, options: PreinitOptions) {
-  if (!currentResources) {
+export function preinit(href: string, options: PreinitOptions) {
+  const currentDispatcher = ReactDOMCurrentDispatcher.current;
+  if (!currentDispatcher || !currentDispatcher.currentResources) {
     // While we expect that preinit calls are primarily going to be observed
     // during render because effects and events don't run on the server it is
     // still possible that these get called in module scope. This is valid on
@@ -186,6 +176,7 @@ function preinit(href: string, options: PreinitOptions) {
     // simply return and do not warn.
     return;
   }
+  const currentResources: Resources = (currentDispatcher.currentResources: any);
   if (__DEV__) {
     validatePreinitArguments(href, options);
   }
@@ -435,11 +426,13 @@ function captureImplicitPreloadResourceDependency(
 
 // Construct a resource from link props.
 export function resourcesFromLink(props: Props): boolean {
-  if (!currentResources) {
+  const currentDispatcher = ReactDOMCurrentDispatcher.current;
+  if (!currentDispatcher || !currentDispatcher.currentResources) {
     throw new Error(
       '"currentResources" was expected to exist. This is a bug in React.',
     );
   }
+  const currentResources: Resources = (currentDispatcher.currentResources: any);
   const {rel, href} = props;
   if (!href || typeof href !== 'string') {
     return false;
@@ -574,8 +567,3 @@ export function hoistResourcesToRoot(
   });
   boundaryResources.clear();
 }
-
-const Dispatcher = {
-  preload,
-  preinit,
-};

packages/react-dom-bindings/src/server/ReactDOMServerFormatConfig.js

@@ -11,6 +11,8 @@ import type {ReactNodeList} from 'shared/ReactTypes';
 import type {Resources, BoundaryResources} from './ReactDOMFloatServer';
 export type {Resources, BoundaryResources};
 
+import ReactDOMSharedInternals from 'shared/ReactDOMSharedInternals';
+const ReactDOMCurrentDispatcher = ReactDOMSharedInternals.Dispatcher;
 import {
   checkHtmlStringCoercion,
   checkCSSPropertyStringCoercion,
@@ -60,11 +62,7 @@ import hasOwnProperty from 'shared/hasOwnProperty';
 import sanitizeURL from '../shared/sanitizeURL';
 import isArray from 'shared/isArray';
 
-import {
-  prepareToRenderResources,
-  finishRenderingResources,
-  resourcesFromLink,
-} from './ReactDOMFloatServer';
+import {resourcesFromLink, preload, preinit} from './ReactDOMFloatServer';
 export {
   createResources,
   createBoundaryResources,
@@ -2723,10 +2721,25 @@ function writeStyleResourceAttribute(
   );
 }
 
-export function prepareToRender(resources: Resources) {
-  prepareToRenderResources(resources);
+export opaque type HostDispatcher = mixed;
+export function createHostDispatcher(resources: Resources): HostDispatcher {
+  return enableFloat
+    ? {
+        preinit,
+        preload,
+        currentResources: resources,
+      }
+    : {};
+}
+
+export function setCurrentHostDispatcher(
+  hostDispatcher: HostDispatcher,
+): mixed {
+  const previousDispatcher = ReactDOMCurrentDispatcher.current;
+  ReactDOMCurrentDispatcher.current = hostDispatcher;
+  return previousDispatcher;
 }
 
-export function cleanupAfterRender() {
-  finishRenderingResources();
+export function restoreHostDispatcher(previousDispatcher: mixed): void {
+  ReactDOMCurrentDispatcher.current = previousDispatcher;
 }

packages/react-dom-bindings/src/server/ReactDOMServerLegacyFormatConfig.js

@@ -80,6 +80,7 @@ export type {
   BoundaryResources,
   FormatContext,
   SuspenseBoundaryID,
+  HostDispatcher,
 } from './ReactDOMServerFormatConfig';
 
 export {
@@ -107,8 +108,9 @@ export {
   hoistResources,
   hoistResourcesToRoot,
   setCurrentlyRenderingBoundaryResourcesTarget,
-  prepareToRender,
-  cleanupAfterRender,
+  createHostDispatcher,
+  setCurrentHostDispatcher,
+  restoreHostDispatcher,
 } from './ReactDOMServerFormatConfig';
 
 import {stringToChunk} from 'react-server/src/ReactServerStreamConfig';

packages/react-dom-bindings/src/shared/ReactDOMFloat.js

@@ -1,8 +1,9 @@
 import ReactDOMSharedInternals from 'shared/ReactDOMSharedInternals';
+const ReactDOMCurrentDispatcher = ReactDOMSharedInternals.Dispatcher;
 
 export function preinit() {
-  const dispatcher = ReactDOMSharedInternals.Dispatcher.current;
-  if (dispatcher) {
+  const dispatcher = ReactDOMCurrentDispatcher.current;
+  if (dispatcher && dispatcher.preinit) {
     dispatcher.preinit.apply(this, arguments);
   }
   // We don't error because preinit needs to be resilient to being called in a variety of scopes
@@ -11,8 +12,8 @@ export function preinit() {
 }
 
 export function preload() {
-  const dispatcher = ReactDOMSharedInternals.Dispatcher.current;
-  if (dispatcher) {
+  const dispatcher = ReactDOMCurrentDispatcher.current;
+  if (dispatcher && dispatcher.preload) {
     dispatcher.preload.apply(this, arguments);
   }
   // We don't error because preload needs to be resilient to being called in a variety of scopes

packages/react-native-renderer/src/server/ReactNativeServerFormatConfig.js

@@ -341,8 +341,13 @@ export function hoistResourcesToRoot(
   boundaryResources: BoundaryResources,
 ) {}
 
-export function prepareToRender(resources: Resources) {}
-export function cleanupAfterRender() {}
+export opaque type HostDispatcher = void;
+export function createHostDispatcher(resources: Resources): HostDispatcher {}
+export function setCurrentHostDispatcher(
+  hostDispatcher: HostDispatcher,
+): mixed {}
+export function restoreHostDispatcher(previousDispatcher: mixed) {}
+
 export function createResources() {}
 export function createBoundaryResources() {}
 export function setCurrentlyRenderingBoundaryResourcesTarget(

packages/react-noop-renderer/src/ReactNoopServer.js

@@ -280,8 +280,9 @@ const ReactNoopServer = ReactFizzServer({
 
   setCurrentlyRenderingBoundaryResourcesTarget(resources: BoundaryResources) {},
 
-  prepareToRender() {},
-  cleanupAfterRender() {},
+  createHostDispatcher() {},
+  setCurrentHostDispatcher() {},
+  restoreHostDispatcher() {},
 });
 
 type Options = {

packages/react-server/src/ReactFizzServer.js

@@ -26,6 +26,7 @@ import type {
   FormatContext,
   Resources,
   BoundaryResources,
+  HostDispatcher,
 } from './ReactServerFormatConfig';
 import type {ContextSnapshot} from './ReactFizzNewContext';
 import type {ComponentStackNode} from './ReactFizzComponentStack';
@@ -69,11 +70,12 @@ import {
   writeImmediateResources,
   hoistResources,
   hoistResourcesToRoot,
-  prepareToRender,
-  cleanupAfterRender,
   setCurrentlyRenderingBoundaryResourcesTarget,
   createResources,
   createBoundaryResources,
+  setCurrentHostDispatcher,
+  restoreHostDispatcher,
+  createHostDispatcher,
 } from './ReactServerFormatConfig';
 import {
   constructClassInstance,
@@ -204,6 +206,7 @@ const CLOSED = 2;
 
 export opaque type Request = {
   destination: null | Destination,
+  hostDispatcher: HostDispatcher,
   +responseState: ResponseState,
   +progressiveChunkSize: number,
   status: 0 | 1 | 2,
@@ -276,8 +279,10 @@ export function createRequest(
   const pingedTasks = [];
   const abortSet: Set<Task> = new Set();
   const resources: Resources = createResources();
+  const hostDispatcher: HostDispatcher = createHostDispatcher(resources);
   const request = {
     destination: null,
+    hostDispatcher,
     responseState,
     progressiveChunkSize:
       progressiveChunkSize === undefined
@@ -1827,7 +1832,6 @@ function finishedTask(
 }
 
 function retryTask(request: Request, task: Task): void {
-  prepareToRender(request.resources);
   if (enableFloat) {
     const blockedBoundary = task.blockedBoundary;
     setCurrentlyRenderingBoundaryResourcesTarget(
@@ -1892,7 +1896,6 @@ function retryTask(request: Request, task: Task): void {
     if (__DEV__) {
       currentTaskInDEV = prevTaskInDEV;
     }
-    cleanupAfterRender();
   }
 }
 
@@ -1903,6 +1906,9 @@ export function performWork(request: Request): void {
   const prevContext = getActiveContext();
   const prevDispatcher = ReactCurrentDispatcher.current;
   ReactCurrentDispatcher.current = Dispatcher;
+  const previousHostDispatcher = setCurrentHostDispatcher(
+    request.hostDispatcher,
+  );
   let prevGetCurrentStackImpl;
   if (__DEV__) {
     prevGetCurrentStackImpl = ReactDebugCurrentFrame.getCurrentStack;
@@ -1927,6 +1933,7 @@ export function performWork(request: Request): void {
   } finally {
     setCurrentResponseState(prevResponseState);
     ReactCurrentDispatcher.current = prevDispatcher;
+    restoreHostDispatcher(previousHostDispatcher);
     if (__DEV__) {
       ReactDebugCurrentFrame.getCurrentStack = prevGetCurrentStackImpl;
     }

packages/react-server/src/forks/ReactServerFormatConfig.custom.js

@@ -30,6 +30,7 @@ export opaque type Resources = mixed;
 export opaque type BoundaryResources = mixed;
 export opaque type FormatContext = mixed;
 export opaque type SuspenseBoundaryID = mixed;
+export opaque type HostDispatcher = mixed;
 
 export const isPrimaryRenderer = false;
 
@@ -68,6 +69,9 @@ export const writeCompletedBoundaryInstruction =
   $$$hostConfig.writeCompletedBoundaryInstruction;
 export const writeClientRenderBoundaryInstruction =
   $$$hostConfig.writeClientRenderBoundaryInstruction;
+export const createHostDispatcher = $$$hostConfig.createHostDispatcher;
+export const setCurrentHostDispatcher = $$$hostConfig.setCurrentHostDispatcher;
+export const restoreHostDispatcher = $$$hostConfig.restoreHostDispatcher;
 
 // -------------------------
 //     Resources
@@ -76,8 +80,6 @@ export const writeInitialResources = $$$hostConfig.writeInitialResources;
 export const writeImmediateResources = $$$hostConfig.writeImmediateResources;
 export const hoistResources = $$$hostConfig.hoistResources;
 export const hoistResourcesToRoot = $$$hostConfig.hoistResourcesToRoot;
-export const prepareToRender = $$$hostConfig.prepareToRender;
-export const cleanupAfterRender = $$$hostConfig.cleanupAfterRender;
 export const createResources = $$$hostConfig.createResources;
 export const createBoundaryResources = $$$hostConfig.createBoundaryResources;
 export const setCurrentlyRenderingBoundaryResourcesTarget =