Track Owner for Server Components in DEV (#28753)

This implements the concept of a DEV-only "owner" for Server Components.
The owner concept isn't really super useful. We barely use it anymore,
but we do have it as a concept in DevTools in a couple of cases so this
adds it for parity. However, this is mainly interesting because it could
be used to wire up future owner-based stacks.

I do this by outlining the DebugInfo for a Server Component
(ReactComponentInfo). Then I just rely on Flight deduping to refer to
that. I refer to the same thing by referential equality so that we can
associate a Server Component parent in DebugInfo with an owner.

If you suspend and replay a Server Component, we have to restore the
same owner. To do that, I did a little ugly hack and stashed it on the
thenable state object. Felt unnecessarily complicated to add a stateful
wrapper for this one dev-only case.

The owner could really be anything since it could be coming from a
different implementation. Because this is the first time we have an
owner other than Fiber, I have to fix up a bunch of places that assumes
Fiber. I mainly did the `typeof owner.tag === 'number'` to assume it's a
Fiber for now.

This also doesn't actually add it to DevTools / RN Inspector yet. I just
ignore them there for now.

Because Server Components can be async the owner isn't tracked after an
await. We need per-component AsyncLocalStorage for that. This can be
done in a follow up.

packages/react-client/src/ReactFlightClient.js

@@ -484,6 +484,7 @@ function createElement(
   type: mixed,
   key: mixed,
   props: mixed,
+  owner: null | ReactComponentInfo, // DEV-only
 ): React$Element<any> {
   let element: any;
   if (__DEV__ && enableRefAsProp) {
@@ -493,7 +494,7 @@ function createElement(
       type,
       key,
       props,
-      _owner: null,
+      _owner: owner,
     }: any);
     Object.defineProperty(element, 'ref', {
       enumerable: false,
@@ -520,7 +521,7 @@ function createElement(
       props,
 
       // Record the component responsible for creating this element.
-      _owner: null,
+      _owner: owner,
     }: any);
   }
 
@@ -854,7 +855,12 @@ function parseModelTuple(
   if (tuple[0] === REACT_ELEMENT_TYPE) {
     // TODO: Consider having React just directly accept these arrays as elements.
     // Or even change the ReactElement type to be an array.
-    return createElement(tuple[1], tuple[2], tuple[3]);
+    return createElement(
+      tuple[1],
+      tuple[2],
+      tuple[3],
+      __DEV__ ? (tuple: any)[4] : null,
+    );
   }
   return value;
 }
@@ -1132,12 +1138,14 @@ function resolveConsoleEntry(
     );
   }
 
-  const payload: [string, string, string, mixed] = parseModel(response, value);
+  const payload: [string, string, null | ReactComponentInfo, string, mixed] =
+    parseModel(response, value);
   const methodName = payload[0];
   // TODO: Restore the fake stack before logging.
   // const stackTrace = payload[1];
-  const env = payload[2];
-  const args = payload.slice(3);
+  // const owner = payload[2];
+  const env = payload[3];
+  const args = payload.slice(4);
   printToConsole(methodName, args, env);
 }
 
@@ -1286,7 +1294,10 @@ function processFullRow(
     }
     case 68 /* "D" */: {
       if (__DEV__) {
-        const debugInfo = JSON.parse(row);
+        const debugInfo: ReactComponentInfo | ReactAsyncInfo = parseModel(
+          response,
+          row,
+        );
         resolveDebugInfo(response, id, debugInfo);
         return;
       }

packages/react-client/src/__tests__/ReactFlight-test.js

@@ -214,7 +214,7 @@ describe('ReactFlight', () => {
       const rootModel = await ReactNoopFlightClient.read(transport);
       const greeting = rootModel.greeting;
       expect(greeting._debugInfo).toEqual(
-        __DEV__ ? [{name: 'Greeting', env: 'Server'}] : undefined,
+        __DEV__ ? [{name: 'Greeting', env: 'Server', owner: null}] : undefined,
       );
       ReactNoop.render(greeting);
     });
@@ -241,7 +241,7 @@ describe('ReactFlight', () => {
     await act(async () => {
       const promise = ReactNoopFlightClient.read(transport);
       expect(promise._debugInfo).toEqual(
-        __DEV__ ? [{name: 'Greeting', env: 'Server'}] : undefined,
+        __DEV__ ? [{name: 'Greeting', env: 'Server', owner: null}] : undefined,
       );
       ReactNoop.render(await promise);
     });
@@ -2072,19 +2072,21 @@ describe('ReactFlight', () => {
     await act(async () => {
       const promise = ReactNoopFlightClient.read(transport);
       expect(promise._debugInfo).toEqual(
-        __DEV__ ? [{name: 'ServerComponent', env: 'Server'}] : undefined,
+        __DEV__
+          ? [{name: 'ServerComponent', env: 'Server', owner: null}]
+          : undefined,
       );
       const result = await promise;
       const thirdPartyChildren = await result.props.children[1];
       // We expect the debug info to be transferred from the inner stream to the outer.
       expect(thirdPartyChildren[0]._debugInfo).toEqual(
         __DEV__
-          ? [{name: 'ThirdPartyComponent', env: 'third-party'}]
+          ? [{name: 'ThirdPartyComponent', env: 'third-party', owner: null}]
           : undefined,
       );
       expect(thirdPartyChildren[1]._debugInfo).toEqual(
         __DEV__
-          ? [{name: 'ThirdPartyLazyComponent', env: 'third-party'}]
+          ? [{name: 'ThirdPartyLazyComponent', env: 'third-party', owner: null}]
           : undefined,
       );
       ReactNoop.render(result);
@@ -2145,4 +2147,50 @@ describe('ReactFlight', () => {
     expect(loggedFn).not.toBe(foo);
     expect(loggedFn.toString()).toBe(foo.toString());
   });
+
+  it('uses the server component debug info as the element owner in DEV', async () => {
+    function Container({children}) {
+      return children;
+    }
+
+    function Greeting({firstName}) {
+      // We can't use JSX here because it'll use the Client React.
+      return ReactServer.createElement(
+        Container,
+        null,
+        ReactServer.createElement('span', null, 'Hello, ', firstName),
+      );
+    }
+
+    const model = {
+      greeting: ReactServer.createElement(Greeting, {firstName: 'Seb'}),
+    };
+
+    const transport = ReactNoopFlightServer.render(model);
+
+    await act(async () => {
+      const rootModel = await ReactNoopFlightClient.read(transport);
+      const greeting = rootModel.greeting;
+      // We've rendered down to the span.
+      expect(greeting.type).toBe('span');
+      if (__DEV__) {
+        const greetInfo = {name: 'Greeting', env: 'Server', owner: null};
+        expect(greeting._debugInfo).toEqual([
+          greetInfo,
+          {name: 'Container', env: 'Server', owner: greetInfo},
+        ]);
+        // The owner that created the span was the outer server component.
+        // We expect the debug info to be referentially equal to the owner.
+        expect(greeting._owner).toBe(greeting._debugInfo[0]);
+      } else {
+        expect(greeting._debugInfo).toBe(undefined);
+        expect(greeting._owner).toBe(
+          gate(flags => flags.disableStringRefs) ? undefined : null,
+        );
+      }
+      ReactNoop.render(greeting);
+    });
+
+    expect(ReactNoop).toMatchRenderedOutput(<span>Hello, Seb</span>);
+  });
 });

packages/react-devtools-shared/src/__tests__/componentStacks-test.js

@@ -101,6 +101,7 @@ describe('component stack', () => {
       {
         name: 'ServerComponent',
         env: 'Server',
+        owner: null,
       },
     ];
     const Parent = () => ChildPromise;

packages/react-devtools-shared/src/backend/DevToolsComponentStackFrame.js

@@ -33,10 +33,7 @@ import {
 import {disableLogs, reenableLogs} from './DevToolsConsolePatching';
 
 let prefix;
-export function describeBuiltInComponentFrame(
-  name: string,
-  ownerFn: void | null | Function,
-): string {
+export function describeBuiltInComponentFrame(name: string): string {
   if (prefix === undefined) {
     // Extract the VM specific prefix used by each line.
     try {
@@ -51,10 +48,7 @@ export function describeBuiltInComponentFrame(
 }
 
 export function describeDebugInfoFrame(name: string, env: ?string): string {
-  return describeBuiltInComponentFrame(
-    name + (env ? ' (' + env + ')' : ''),
-    null,
-  );
+  return describeBuiltInComponentFrame(name + (env ? ' (' + env + ')' : ''));
 }
 
 let reentry = false;
@@ -292,15 +286,13 @@ export function describeNativeComponentFrame(
 
 export function describeClassComponentFrame(
   ctor: Function,
-  ownerFn: void | null | Function,
   currentDispatcherRef: CurrentDispatcherRef,
 ): string {
   return describeNativeComponentFrame(ctor, true, currentDispatcherRef);
 }
 
 export function describeFunctionComponentFrame(
   fn: Function,
-  ownerFn: void | null | Function,
   currentDispatcherRef: CurrentDispatcherRef,
 ): string {
   return describeNativeComponentFrame(fn, false, currentDispatcherRef);
@@ -313,7 +305,6 @@ function shouldConstruct(Component: Function) {
 
 export function describeUnknownElementTypeFrameInDEV(
   type: any,
-  ownerFn: void | null | Function,
   currentDispatcherRef: CurrentDispatcherRef,
 ): string {
   if (!__DEV__) {
@@ -330,31 +321,29 @@ export function describeUnknownElementTypeFrameInDEV(
     );
   }
   if (typeof type === 'string') {
-    return describeBuiltInComponentFrame(type, ownerFn);
+    return describeBuiltInComponentFrame(type);
   }
   switch (type) {
     case SUSPENSE_NUMBER:
     case SUSPENSE_SYMBOL_STRING:
-      return describeBuiltInComponentFrame('Suspense', ownerFn);
+      return describeBuiltInComponentFrame('Suspense');
     case SUSPENSE_LIST_NUMBER:
     case SUSPENSE_LIST_SYMBOL_STRING:
-      return describeBuiltInComponentFrame('SuspenseList', ownerFn);
+      return describeBuiltInComponentFrame('SuspenseList');
   }
   if (typeof type === 'object') {
     switch (type.$$typeof) {
       case FORWARD_REF_NUMBER:
       case FORWARD_REF_SYMBOL_STRING:
         return describeFunctionComponentFrame(
           type.render,
-          ownerFn,
           currentDispatcherRef,
         );
       case MEMO_NUMBER:
       case MEMO_SYMBOL_STRING:
         // Memo may contain any component type so we recursively resolve it.
         return describeUnknownElementTypeFrameInDEV(
           type.type,
-          ownerFn,
           currentDispatcherRef,
         );
       case LAZY_NUMBER:
@@ -366,7 +355,6 @@ export function describeUnknownElementTypeFrameInDEV(
           // Lazy may contain any component type so we recursively resolve it.
           return describeUnknownElementTypeFrameInDEV(
             init(payload),
-            ownerFn,
             currentDispatcherRef,
           );
         } catch (x) {}

packages/react-devtools-shared/src/backend/DevToolsFiberComponentStack.js

@@ -39,38 +39,30 @@ export function describeFiber(
     ClassComponent,
   } = workTagMap;
 
-  const owner: null | Function = __DEV__
-    ? workInProgress._debugOwner
-      ? workInProgress._debugOwner.type
-      : null
-    : null;
   switch (workInProgress.tag) {
     case HostComponent:
-      return describeBuiltInComponentFrame(workInProgress.type, owner);
+      return describeBuiltInComponentFrame(workInProgress.type);
     case LazyComponent:
-      return describeBuiltInComponentFrame('Lazy', owner);
+      return describeBuiltInComponentFrame('Lazy');
     case SuspenseComponent:
-      return describeBuiltInComponentFrame('Suspense', owner);
+      return describeBuiltInComponentFrame('Suspense');
     case SuspenseListComponent:
-      return describeBuiltInComponentFrame('SuspenseList', owner);
+      return describeBuiltInComponentFrame('SuspenseList');
     case FunctionComponent:
     case IndeterminateComponent:
     case SimpleMemoComponent:
       return describeFunctionComponentFrame(
         workInProgress.type,
-        owner,
         currentDispatcherRef,
       );
     case ForwardRef:
       return describeFunctionComponentFrame(
         workInProgress.type.render,
-        owner,
         currentDispatcherRef,
       );
     case ClassComponent:
       return describeClassComponentFrame(
         workInProgress.type,
-        owner,
         currentDispatcherRef,
       );
     default:

packages/react-devtools-shared/src/backend/renderer.js

@@ -1952,15 +1952,24 @@ export function attach(
       const {key} = fiber;
       const displayName = getDisplayNameForFiber(fiber);
       const elementType = getElementTypeForFiber(fiber);
-      const {_debugOwner} = fiber;
+      const debugOwner = fiber._debugOwner;
 
       // Ideally we should call getFiberIDThrows() for _debugOwner,
       // since owners are almost always higher in the tree (and so have already been processed),
       // but in some (rare) instances reported in open source, a descendant mounts before an owner.
       // Since this is a DEV only field it's probably okay to also just lazily generate and ID here if needed.
       // See https://github.com/facebook/react/issues/21445
-      const ownerID =
-        _debugOwner != null ? getOrGenerateFiberID(_debugOwner) : 0;
+      let ownerID: number;
+      if (debugOwner != null) {
+        if (typeof debugOwner.tag === 'number') {
+          ownerID = getOrGenerateFiberID((debugOwner: any));
+        } else {
+          // TODO: Track Server Component Owners.
+          ownerID = 0;
+        }
+      } else {
+        ownerID = 0;
+      }
       const parentID = parentFiber ? getFiberIDThrows(parentFiber) : 0;
 
       const displayNameStringID = getStringID(displayName);
@@ -3104,15 +3113,17 @@ export function attach(
       return null;
     }
 
-    const {_debugOwner} = fiber;
-
     const owners: Array<SerializedElement> = [fiberToSerializedElement(fiber)];
 
-    if (_debugOwner) {
-      let owner: null | Fiber = _debugOwner;
-      while (owner !== null) {
-        owners.unshift(fiberToSerializedElement(owner));
-        owner = owner._debugOwner || null;
+    let owner = fiber._debugOwner;
+    while (owner != null) {
+      if (typeof owner.tag === 'number') {
+        const ownerFiber: Fiber = (owner: any); // Refined
+        owners.unshift(fiberToSerializedElement(ownerFiber));
+        owner = ownerFiber._debugOwner;
+      } else {
+        // TODO: Track Server Component Owners.
+        break;
       }
     }
 
@@ -3173,7 +3184,7 @@ export function attach(
     }
 
     const {
-      _debugOwner,
+      _debugOwner: debugOwner,
       stateNode,
       key,
       memoizedProps,
@@ -3300,13 +3311,19 @@ export function attach(
       context = {value: context};
     }
 
-    let owners = null;
-    if (_debugOwner) {
-      owners = ([]: Array<SerializedElement>);
-      let owner: null | Fiber = _debugOwner;
-      while (owner !== null) {
-        owners.push(fiberToSerializedElement(owner));
-        owner = owner._debugOwner || null;
+    let owners: null | Array<SerializedElement> = null;
+    let owner = debugOwner;
+    while (owner != null) {
+      if (typeof owner.tag === 'number') {
+        const ownerFiber: Fiber = (owner: any); // Refined
+        if (owners === null) {
+          owners = [];
+        }
+        owners.push(fiberToSerializedElement(ownerFiber));
+        owner = ownerFiber._debugOwner;
+      } else {
+        // TODO: Track Server Component Owners.
+        break;
       }
     }