[Flight] Track Owner on AsyncInfo and IOInfo

packages/react-client/src/ReactFlightClient.js

@@ -2616,14 +2616,15 @@ function resolveDebugInfo(
     initializeFakeTask(response, componentInfoOrAsyncInfo, env);
   }
   if (debugInfo.owner === null && response._debugRootOwner != null) {
-    // $FlowFixMe[prop-missing] By narrowing `owner` to `null`, we narrowed `debugInfo` to `ReactComponentInfo`
-    const componentInfo: ReactComponentInfo = debugInfo;
+    const componentInfoOrAsyncInfo: ReactComponentInfo | ReactAsyncInfo =
+      // $FlowFixMe: By narrowing `owner` to `null`, we narrowed `debugInfo` to `ReactComponentInfo`
+      debugInfo;
     // $FlowFixMe[cannot-write]
-    componentInfo.owner = response._debugRootOwner;
+    componentInfoOrAsyncInfo.owner = response._debugRootOwner;
     // We override the stack if we override the owner since the stack where the root JSX
     // was created on the server isn't very useful but where the request was made is.
     // $FlowFixMe[cannot-write]
-    componentInfo.debugStack = response._debugRootStack;
+    componentInfoOrAsyncInfo.debugStack = response._debugRootStack;
   } else if (debugInfo.stack !== undefined) {
     const componentInfoOrAsyncInfo: ReactComponentInfo | ReactAsyncInfo =
       // $FlowFixMe[incompatible-type]
@@ -2764,7 +2765,6 @@ function initializeIOInfo(response: Response, ioInfo: ReactIOInfo): void {
     initializeFakeTask(response, ioInfo, env);
     initializeFakeStack(response, ioInfo);
   }
-  // TODO: Initialize owner.
   // Adjust the time to the current environment's time space.
   // $FlowFixMe[cannot-write]
   ioInfo.start += response._timeOrigin;

packages/react-server/src/ReactFlightAsyncSequence.js

@@ -7,12 +7,15 @@
  * @flow
  */
 
+import type {ReactComponentInfo} from 'shared/ReactTypes';
+
 export const IO_NODE = 0;
 export const PROMISE_NODE = 1;
 export const AWAIT_NODE = 2;
 
 export type IONode = {
   tag: 0,
+  owner: null | ReactComponentInfo,
   stack: Error, // callsite that spawned the I/O
   start: number, // start time when the first part of the I/O sequence started
   end: number, // we typically don't use this. only when there's no promise intermediate.
@@ -22,6 +25,7 @@ export type IONode = {
 
 export type PromiseNode = {
   tag: 1,
+  owner: null | ReactComponentInfo,
   stack: Error, // callsite that created the Promise
   start: number, // start time when the Promise was created
   end: number, // end time when the Promise was resolved.
@@ -31,6 +35,7 @@ export type PromiseNode = {
 
 export type AwaitNode = {
   tag: 2,
+  owner: null | ReactComponentInfo,
   stack: Error, // callsite that awaited (using await, .then(), Promise.all(), ...)
   start: -1.1, // not used. We use the timing of the awaited promise.
   end: -1.1, // not used.

packages/react-server/src/ReactFlightServer.js

@@ -1934,6 +1934,7 @@ function visitAsyncNode(
           request.pendingChunks++;
           emitDebugChunk(request, task.id, {
             awaited: ((ioNode: any): ReactIOInfo), // This is deduped by this reference.
+            owner: node.owner,
             stack: stack,
           });
         }
@@ -3523,6 +3524,7 @@ function emitIOInfoChunk(
   name: string,
   start: number,
   end: number,
+  owner: ?ReactComponentInfo,
   stack: ?ReactStackTrace,
 ): void {
   if (!__DEV__) {
@@ -3560,8 +3562,15 @@ function emitIOInfoChunk(
     name: name,
     start: relativeStartTimestamp,
     end: relativeEndTimestamp,
-    stack: stack,
   };
+  if (stack != null) {
+    // $FlowFixMe[cannot-write]
+    debugIOInfo.stack = stack;
+  }
+  if (owner != null) {
+    // $FlowFixMe[cannot-write]
+    debugIOInfo.owner = owner;
+  }
   // $FlowFixMe[incompatible-type] stringify can return null
   const json: string = stringify(debugIOInfo, replacer);
   const row = id.toString(16) + ':J' + json + '\n';
@@ -3577,12 +3586,18 @@ function outlineIOInfo(request: Request, ioInfo: ReactIOInfo): void {
   // We can't serialize the ConsoleTask/Error objects so we need to omit them before serializing.
   request.pendingChunks++;
   const id = request.nextChunkId++;
+  const owner = ioInfo.owner;
+  // Ensure the owner is already outlined.
+  if (owner != null) {
+    outlineComponentInfo(request, owner);
+  }
   emitIOInfoChunk(
     request,
     id,
     ioInfo.name,
     ioInfo.start,
     ioInfo.end,
+    owner,
     ioInfo.stack,
   );
   request.writtenObjects.set(ioInfo, serializeByValueID(id));
@@ -3612,10 +3627,15 @@ function serializeIONode(
       name = name.slice(7);
     }
   }
+  const owner = ioNode.owner;
+  // Ensure the owner is already outlined.
+  if (owner != null) {
+    outlineComponentInfo(request, owner);
+  }
 
   request.pendingChunks++;
   const id = request.nextChunkId++;
-  emitIOInfoChunk(request, id, name, ioNode.start, ioNode.end, stack);
+  emitIOInfoChunk(request, id, name, ioNode.start, ioNode.end, owner, stack);
   const ref = serializeByValueID(id);
   request.writtenObjects.set(ioNode, ref);
   return ref;
@@ -4141,6 +4161,7 @@ function forwardDebugInfo(
         const debugAsyncInfo: Omit<ReactAsyncInfo, 'debugTask' | 'debugStack'> =
           {
             awaited: ioInfo,
+            owner: debugInfo[i].owner,
             stack: debugInfo[i].stack,
           };
         emitDebugChunk(request, id, debugAsyncInfo);

packages/react-server/src/ReactFlightServerConfigDebugNode.js

@@ -15,6 +15,7 @@ import type {
 } from './ReactFlightAsyncSequence';
 
 import {IO_NODE, PROMISE_NODE, AWAIT_NODE} from './ReactFlightAsyncSequence';
+import {resolveOwner} from './flight/ReactFlightCurrentOwner';
 import {createHook, executionAsyncId} from 'async_hooks';
 import {enableAsyncDebugInfo} from 'shared/ReactFeatureFlags';
 
@@ -46,6 +47,7 @@ export function initAsyncDebugInfo(): void {
             // so that we can later pick the best stack trace in user space.
             node = ({
               tag: AWAIT_NODE,
+              owner: resolveOwner(),
               stack: new Error(),
               start: -1.1,
               end: -1.1,
@@ -55,6 +57,7 @@ export function initAsyncDebugInfo(): void {
           } else {
             node = ({
               tag: PROMISE_NODE,
+              owner: resolveOwner(),
               stack: new Error(),
               start: performance.now(),
               end: -1.1, // Set when we resolve.
@@ -74,6 +77,7 @@ export function initAsyncDebugInfo(): void {
             // We have begun a new I/O sequence.
             node = ({
               tag: IO_NODE,
+              owner: resolveOwner(),
               stack: new Error(), // This is only used if no native promises are used.
               start: performance.now(),
               end: -1.1, // Only set when pinged.
@@ -84,6 +88,7 @@ export function initAsyncDebugInfo(): void {
             // We have begun a new I/O sequence after the await.
             node = ({
               tag: IO_NODE,
+              owner: resolveOwner(),
               stack: new Error(),
               start: performance.now(),
               end: -1.1, // Only set when pinged.