Move Current Owner (and Cache) to an Async Dispatcher

packages/react-dom/src/__tests__/ReactCompositeComponent-test.js

@@ -537,16 +537,24 @@ describe('ReactCompositeComponent', () => {
   });
 
   it('should cleanup even if render() fatals', async () => {
+    const dispatcherEnabled =
+      __DEV__ ||
+      !gate(flags => flags.disableStringRefs) ||
+      gate(flags => flags.enableCache);
+    const ownerEnabled = __DEV__ || !gate(flags => flags.disableStringRefs);
+
+    let stashedDispatcher;
     class BadComponent extends React.Component {
       render() {
+        // Stash the dispatcher that was available in render so we can check
+        // that its internals also reset.
+        stashedDispatcher = ReactSharedInternals.A;
         throw new Error();
       }
     }
 
     const instance = <BadComponent />;
-    expect(ReactSharedInternals.owner).toBe(
-      __DEV__ || !gate(flags => flags.disableStringRefs) ? null : undefined,
-    );
+    expect(ReactSharedInternals.A).toBe(dispatcherEnabled ? null : undefined);
 
     const root = ReactDOMClient.createRoot(document.createElement('div'));
     await expect(async () => {
@@ -555,9 +563,16 @@ describe('ReactCompositeComponent', () => {
       });
     }).rejects.toThrow();
 
-    expect(ReactSharedInternals.owner).toBe(
-      __DEV__ || !gate(flags => flags.disableStringRefs) ? null : undefined,
-    );
+    expect(ReactSharedInternals.A).toBe(dispatcherEnabled ? null : undefined);
+    if (dispatcherEnabled) {
+      if (ownerEnabled) {
+        expect(stashedDispatcher.getOwner()).toBe(null);
+      } else {
+        expect(stashedDispatcher.getOwner).toBe(undefined);
+      }
+    } else {
+      expect(stashedDispatcher).toBe(undefined);
+    }
   });
 
   it('should call componentWillUnmount before unmounting', async () => {

packages/react-dom/src/client/ReactDOMRootFB.js

@@ -61,7 +61,7 @@ import {LegacyRoot} from 'react-reconciler/src/ReactRootTags';
 import getComponentNameFromType from 'shared/getComponentNameFromType';
 import {has as hasInstance} from 'shared/ReactInstanceMap';
 
-import ReactSharedInternals from 'shared/ReactSharedInternals';
+import {currentOwner} from 'react-reconciler/src/ReactFiberCurrentOwner';
 
 import assign from 'shared/assign';
 
@@ -342,7 +342,7 @@ export function findDOMNode(
   componentOrElement: Element | ?React$Component<any, any>,
 ): null | Element | Text {
   if (__DEV__) {
-    const owner = (ReactSharedInternals.owner: any);
+    const owner = currentOwner;
     if (owner !== null && owner.stateNode !== null) {
       const warnedAboutRefsInRender = owner.stateNode._warnedAboutRefsInRender;
       if (!warnedAboutRefsInRender) {

packages/react-native-renderer/src/ReactNativePublicCompat.js

@@ -24,14 +24,14 @@ import {
   findHostInstanceWithWarning,
 } from 'react-reconciler/src/ReactFiberReconciler';
 import {doesFiberContain} from 'react-reconciler/src/ReactFiberTreeReflection';
-import ReactSharedInternals from 'shared/ReactSharedInternals';
 import getComponentNameFromType from 'shared/getComponentNameFromType';
+import {currentOwner} from 'react-reconciler/src/ReactFiberCurrentOwner';
 
 export function findHostInstance_DEPRECATED<TElementType: ElementType>(
   componentOrHandle: ?(ElementRef<TElementType> | number),
 ): ?ElementRef<HostComponent<mixed>> {
   if (__DEV__) {
-    const owner = ReactSharedInternals.owner;
+    const owner = currentOwner;
     if (owner !== null && owner.stateNode !== null) {
       if (!owner.stateNode._warnedAboutRefsInRender) {
         console.error(
@@ -86,7 +86,7 @@ export function findHostInstance_DEPRECATED<TElementType: ElementType>(
 
 export function findNodeHandle(componentOrHandle: any): ?number {
   if (__DEV__) {
-    const owner = ReactSharedInternals.owner;
+    const owner = currentOwner;
     if (owner !== null && owner.stateNode !== null) {
       if (!owner.stateNode._warnedAboutRefsInRender) {
         console.error(

packages/react-reconciler/src/ReactFiberCache.js → packages/react-reconciler/src/ReactFiberAsyncDispatcher.js

@@ -7,13 +7,17 @@
  * @flow
  */
 
-import type {CacheDispatcher} from './ReactInternalTypes';
+import type {AsyncDispatcher, Fiber} from './ReactInternalTypes';
 import type {Cache} from './ReactFiberCacheComponent';
 
 import {enableCache} from 'shared/ReactFeatureFlags';
 import {readContext} from './ReactFiberNewContext';
 import {CacheContext} from './ReactFiberCacheComponent';
 
+import {disableStringRefs} from 'shared/ReactFeatureFlags';
+
+import {currentOwner} from './ReactFiberCurrentOwner';
+
 function getCacheForType<T>(resourceType: () => T): T {
   if (!enableCache) {
     throw new Error('Not implemented.');
@@ -27,6 +31,12 @@ function getCacheForType<T>(resourceType: () => T): T {
   return cacheForType;
 }
 
-export const DefaultCacheDispatcher: CacheDispatcher = {
+export const DefaultAsyncDispatcher: AsyncDispatcher = ({
   getCacheForType,
-};
+}: any);
+
+if (__DEV__ || !disableStringRefs) {
+  DefaultAsyncDispatcher.getOwner = (): null | Fiber => {
+    return currentOwner;
+  };
+}

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -91,7 +91,6 @@ import {
   Passive,
   DidDefer,
 } from './ReactFiberFlags';
-import ReactSharedInternals from 'shared/ReactSharedInternals';
 import {
   debugRenderPhaseSideEffectsForStrictMode,
   disableLegacyContext,
@@ -297,6 +296,7 @@ import {
   pushRootMarkerInstance,
   TransitionTracingMarker,
 } from './ReactFiberTracingMarkerComponent';
+import {setCurrentOwner} from './ReactFiberCurrentOwner';
 
 // A special exception that's used to unwind the stack when an update flows
 // into a dehydrated boundary.
@@ -432,7 +432,7 @@ function updateForwardRef(
     markComponentRenderStarted(workInProgress);
   }
   if (__DEV__) {
-    ReactSharedInternals.owner = workInProgress;
+    setCurrentOwner(workInProgress);
     setIsRendering(true);
     nextChildren = renderWithHooks(
       current,
@@ -1150,7 +1150,7 @@ function updateFunctionComponent(
     markComponentRenderStarted(workInProgress);
   }
   if (__DEV__) {
-    ReactSharedInternals.owner = workInProgress;
+    setCurrentOwner(workInProgress);
     setIsRendering(true);
     nextChildren = renderWithHooks(
       current,
@@ -1373,7 +1373,7 @@ function finishClassComponent(
 
   // Rerender
   if (__DEV__ || !disableStringRefs) {
-    ReactSharedInternals.owner = workInProgress;
+    setCurrentOwner(workInProgress);
   }
   let nextChildren;
   if (
@@ -3419,7 +3419,7 @@ function updateContextConsumer(
   }
   let newChildren;
   if (__DEV__) {
-    ReactSharedInternals.owner = workInProgress;
+    setCurrentOwner(workInProgress);
     setIsRendering(true);
     newChildren = render(newValue);
     setIsRendering(false);

packages/react-reconciler/src/ReactFiberCurrentOwner.js

@@ -0,0 +1,16 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import type {Fiber} from './ReactInternalTypes';
+
+export let currentOwner: Fiber | null = null;
+
+export function setCurrentOwner(fiber: null | Fiber) {
+  currentOwner = fiber;
+}

packages/react-reconciler/src/ReactFiberTreeReflection.js

@@ -12,7 +12,6 @@ import type {Container, SuspenseInstance} from './ReactFiberConfig';
 import type {SuspenseState} from './ReactFiberSuspenseComponent';
 
 import {get as getInstance} from 'shared/ReactInstanceMap';
-import ReactSharedInternals from 'shared/ReactSharedInternals';
 import getComponentNameFromFiber from 'react-reconciler/src/getComponentNameFromFiber';
 import {
   ClassComponent,
@@ -25,6 +24,7 @@ import {
   SuspenseComponent,
 } from './ReactWorkTags';
 import {NoFlags, Placement, Hydrating} from './ReactFiberFlags';
+import {currentOwner} from './ReactFiberCurrentOwner';
 
 export function getNearestMountedFiber(fiber: Fiber): null | Fiber {
   let node = fiber;
@@ -89,7 +89,7 @@ export function isFiberMounted(fiber: Fiber): boolean {
 
 export function isMounted(component: React$Component<any, any>): boolean {
   if (__DEV__) {
-    const owner = (ReactSharedInternals.owner: any);
+    const owner = currentOwner;
     if (owner !== null && owner.tag === ClassComponent) {
       const ownerFiber: Fiber = owner;
       const instance = ownerFiber.stateNode;

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -203,7 +203,8 @@ import {
   resetHooksOnUnwind,
   ContextOnlyDispatcher,
 } from './ReactFiberHooks';
-import {DefaultCacheDispatcher} from './ReactFiberCache';
+import {DefaultAsyncDispatcher} from './ReactFiberAsyncDispatcher';
+import {setCurrentOwner} from './ReactFiberCurrentOwner';
 import {
   createCapturedValueAtFiber,
   type CapturedValue,
@@ -1684,7 +1685,7 @@ function handleThrow(root: FiberRoot, thrownValue: any): void {
   resetHooksAfterThrow();
   resetCurrentDebugFiberInDEV();
   if (__DEV__ || !disableStringRefs) {
-    ReactSharedInternals.owner = null;
+    setCurrentOwner(null);
   }
 
   if (thrownValue === SuspenseException) {
@@ -1874,19 +1875,19 @@ function popDispatcher(prevDispatcher: any) {
   ReactSharedInternals.H = prevDispatcher;
 }
 
-function pushCacheDispatcher() {
-  if (enableCache) {
-    const prevCacheDispatcher = ReactSharedInternals.C;
-    ReactSharedInternals.C = DefaultCacheDispatcher;
-    return prevCacheDispatcher;
+function pushAsyncDispatcher() {
+  if (enableCache || __DEV__ || !disableStringRefs) {
+    const prevAsyncDispatcher = ReactSharedInternals.A;
+    ReactSharedInternals.A = DefaultAsyncDispatcher;
+    return prevAsyncDispatcher;
   } else {
     return null;
   }
 }
 
-function popCacheDispatcher(prevCacheDispatcher: any) {
-  if (enableCache) {
-    ReactSharedInternals.C = prevCacheDispatcher;
+function popAsyncDispatcher(prevAsyncDispatcher: any) {
+  if (enableCache || __DEV__ || !disableStringRefs) {
+    ReactSharedInternals.A = prevAsyncDispatcher;
   }
 }
 
@@ -1963,7 +1964,7 @@ function renderRootSync(root: FiberRoot, lanes: Lanes) {
   const prevExecutionContext = executionContext;
   executionContext |= RenderContext;
   const prevDispatcher = pushDispatcher(root.containerInfo);
-  const prevCacheDispatcher = pushCacheDispatcher();
+  const prevAsyncDispatcher = pushAsyncDispatcher();
 
   // If the root or lanes have changed, throw out the existing stack
   // and prepare a fresh one. Otherwise we'll continue where we left off.
@@ -2061,7 +2062,7 @@ function renderRootSync(root: FiberRoot, lanes: Lanes) {
 
   executionContext = prevExecutionContext;
   popDispatcher(prevDispatcher);
-  popCacheDispatcher(prevCacheDispatcher);
+  popAsyncDispatcher(prevAsyncDispatcher);
 
   if (workInProgress !== null) {
     // This is a sync render, so we should have finished the whole tree.
@@ -2104,7 +2105,7 @@ function renderRootConcurrent(root: FiberRoot, lanes: Lanes) {
   const prevExecutionContext = executionContext;
   executionContext |= RenderContext;
   const prevDispatcher = pushDispatcher(root.containerInfo);
-  const prevCacheDispatcher = pushCacheDispatcher();
+  const prevAsyncDispatcher = pushAsyncDispatcher();
 
   // If the root or lanes have changed, throw out the existing stack
   // and prepare a fresh one. Otherwise we'll continue where we left off.
@@ -2317,7 +2318,7 @@ function renderRootConcurrent(root: FiberRoot, lanes: Lanes) {
   resetContextDependencies();
 
   popDispatcher(prevDispatcher);
-  popCacheDispatcher(prevCacheDispatcher);
+  popAsyncDispatcher(prevAsyncDispatcher);
   executionContext = prevExecutionContext;
 
   if (__DEV__) {
@@ -2386,7 +2387,7 @@ function performUnitOfWork(unitOfWork: Fiber): void {
   }
 
   if (__DEV__ || !disableStringRefs) {
-    ReactSharedInternals.owner = null;
+    setCurrentOwner(null);
   }
 }
 
@@ -2501,7 +2502,7 @@ function replaySuspendedUnitOfWork(unitOfWork: Fiber): void {
   }
 
   if (__DEV__ || !disableStringRefs) {
-    ReactSharedInternals.owner = null;
+    setCurrentOwner(null);
   }
 }
 
@@ -2894,7 +2895,7 @@ function commitRootImpl(
 
     // Reset this to null before calling lifecycles
     if (__DEV__ || !disableStringRefs) {
-      ReactSharedInternals.owner = null;
+      setCurrentOwner(null);
     }
 
     // The commit phase is broken into several sub-phases. We do a separate pass

packages/react-reconciler/src/ReactInternalTypes.js

@@ -434,6 +434,8 @@ export type Dispatcher = {
   ) => [Awaited<S>, (P) => void, boolean],
 };
 
-export type CacheDispatcher = {
+export type AsyncDispatcher = {
   getCacheForType: <T>(resourceType: () => T) => T,
+  // DEV-only (or !disableStringRefs)
+  getOwner: () => null | Fiber | ReactComponentInfo,
 };

packages/react-server/src/ReactFizzAsyncDispatcher.js

@@ -0,0 +1,27 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import type {AsyncDispatcher} from 'react-reconciler/src/ReactInternalTypes';
+
+import {disableStringRefs} from 'shared/ReactFeatureFlags';
+
+function getCacheForType<T>(resourceType: () => T): T {
+  throw new Error('Not implemented.');
+}
+
+export const DefaultAsyncDispatcher: AsyncDispatcher = ({
+  getCacheForType,
+}: any);
+
+if (__DEV__ || !disableStringRefs) {
+  // Fizz never tracks owner but the JSX runtime looks for this.
+  DefaultAsyncDispatcher.getOwner = (): null => {
+    return null;
+  };
+}