Set async dispatchers during module init

Author: eps1lon

packages/react-server/src/ReactFlightServer.js

@@ -395,6 +395,8 @@ const {
   TaintRegistryPendingRequests,
 } = ReactSharedInternals;
 
+ReactSharedInternals.A = DefaultAsyncDispatcher;
+
 function throwTaintViolation(message: string) {
   // eslint-disable-next-line react-internal/prod-error-codes
   throw new Error(message);
@@ -448,11 +450,6 @@ function RequestInstance(
   onAllReady: () => void,
   onFatalError: (error: mixed) => void,
 ) {
-  const previousAsyncDispatcher = ReactSharedInternals.A;
-  if (previousAsyncDispatcher === null) {
-    ReactSharedInternals.A = DefaultAsyncDispatcher;
-  }
-
   if (__DEV__) {
     // Unlike Fizz or Fiber, we don't reset this and just keep it on permanently.
     // This lets it act more like the AsyncDispatcher so that we can get the

packages/react-server/src/flight/ReactFlightAsyncDispatcher.js

@@ -12,11 +12,14 @@ import type {ReactComponentInfo} from 'shared/ReactTypes';
 import type {AsyncDispatcher} from 'react-reconciler/src/ReactInternalTypes';
 
 import {resolveRequest, getCache} from '../ReactFlightServer';
+import ReactSharedInternals from '../ReactSharedInternalsServer';
 
 import {disableStringRefs} from 'shared/ReactFeatureFlags';
 
 import {resolveOwner} from './ReactFlightCurrentOwner';
 
+const previousAsyncDispatcher = ReactSharedInternals.A;
+
 function resolveCache(): Map<Function, mixed> {
   const request = resolveRequest();
   if (request) {
@@ -30,9 +33,19 @@ export const DefaultAsyncDispatcher: AsyncDispatcher = ({
     const cache = resolveCache();
     let entry: T | void = (cache.get(resourceType): any);
     if (entry === undefined) {
-      entry = resourceType();
-      // TODO: Warn if undefined?
-      cache.set(resourceType, entry);
+      const chainedEntry: T | void =
+        previousAsyncDispatcher !== null
+          ? previousAsyncDispatcher.getCacheForType(resourceType)
+          : undefined;
+
+      if (chainedEntry !== undefined) {
+        entry = chainedEntry;
+        cache.set(resourceType, chainedEntry);
+      } else {
+        entry = resourceType();
+        // TODO: Warn if undefined?
+        cache.set(resourceType, entry);
+      }
     }
     return entry;
   },