Reuse active cache

Author: eps1lon

packages/react-reconciler/src/ReactInternalTypes.js

@@ -457,7 +457,7 @@ export interface AsyncCache {
 }
 
 export type AsyncDispatcher = {
-  getActiveCache: () => AsyncCache,
+  getActiveCache: () => AsyncCache | null,
   // DEV-only (or !disableStringRefs)
   getOwner: () => null | Fiber | ReactComponentInfo | ComponentStackNode,
 };

packages/react-server/src/ReactFizzAsyncDispatcher.js

@@ -17,7 +17,7 @@ import {disableStringRefs} from 'shared/ReactFeatureFlags';
 
 import {currentTaskInDEV} from './ReactFizzCurrentTask';
 
-function getActiveCache(): AsyncCache {
+function getActiveCache(): AsyncCache | null {
   throw new Error('Not implemented.');
 }
 

packages/react-server/src/ReactFlightServer.js

@@ -10,10 +10,12 @@
 import type {Chunk, BinaryChunk, Destination} from './ReactServerStreamConfig';
 
 import type {Postpone} from 'react/src/ReactPostpone';
+import type {AsyncCache} from 'react-reconciler/src/ReactInternalTypes';
 
 import type {TemporaryReferenceSet} from './ReactFlightServerTemporaryReferences';
 
 import {
+  disableStringRefs,
   enableBinaryFlight,
   enablePostpone,
   enableHalt,
@@ -450,6 +452,24 @@ function RequestInstance(
   onAllReady: () => void,
   onFatalError: (error: mixed) => void,
 ) {
+  const previousAsyncDispatcher = ReactSharedInternals.A;
+
+  if (previousAsyncDispatcher !== null) {
+    const previousActiveCache = previousAsyncDispatcher.getActiveCache();
+    if (previousActiveCache !== null) {
+      ReactSharedInternals.A = ({
+        getActiveCache(): AsyncCache | null {
+          return previousActiveCache;
+        },
+      }: any);
+      if (__DEV__ || !disableStringRefs) {
+        ReactSharedInternals.A.getOwner = previousAsyncDispatcher.getOwner;
+      }
+    }
+  } else {
+    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

@@ -15,54 +15,17 @@ import type {
 } 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> {
+function getActiveCache(): AsyncCache | null {
   const request = resolveRequest();
   if (request) {
     return getCache(request);
   }
-  return new Map();
-}
-
-function getActiveCache(): AsyncCache {
-  const outerCache: AsyncCache | null =
-    previousAsyncDispatcher !== null
-      ? previousAsyncDispatcher.getActiveCache()
-      : null;
-
-  const innerCache = resolveCache();
-
-  if (outerCache === null) {
-    return innerCache;
-  }
-
-  // If both caches are active, reads will prefer the outer cache
-  // Writes will go into both caches.
-  const chainedCache: AsyncCache = {
-    get(resourceType: Function) {
-      const outerEntry = outerCache.get(resourceType);
-      if (outerEntry !== undefined) {
-        return outerEntry;
-      }
-      return innerCache.get(resourceType);
-    },
-    set(resourceType: Function, value: mixed) {
-      if (outerCache !== null) {
-        outerCache.set(resourceType, value);
-      }
-      innerCache.set(resourceType, value);
-      return chainedCache;
-    },
-  };
-
-  return chainedCache;
+  return null;
 }
 
 export const DefaultAsyncDispatcher: AsyncDispatcher = ({

packages/react/src/ReactCacheImpl.js

@@ -61,9 +61,10 @@ export function cache<A: Iterable<mixed>, T>(fn: (...A) => T): (...A) => T {
       return fn.apply(null, arguments);
     }
     const activeCache = dispatcher.getActiveCache();
-    let fnMap: WeakMap<any, CacheNode<T>> | void = (activeCache.get(
-      createCacheRoot,
-    ): any);
+    let fnMap: Map<any, CacheNode<T>> | void =
+      activeCache !== null
+        ? (activeCache.get(createCacheRoot): any)
+        : undefined;
     if (fnMap === undefined) {
       fnMap = createCacheRoot();
       // TODO: Warn if undefined?

packages/react/src/ReactHooks.js

@@ -54,7 +54,8 @@ export function getCacheForType<T>(resourceType: () => T): T {
     return resourceType();
   }
   const activeCache = dispatcher.getActiveCache();
-  let entry: T | void = (activeCache.get(resourceType): any);
+  let entry: T | void =
+    activeCache !== null ? (activeCache.get(resourceType): any) : undefined;
   if (entry === undefined) {
     entry = resourceType();
     // TODO: Warn if undefined?