[Flight] Track Debug Info from Synchronously Unwrapped Promises (#33485)

Stacked on #33482.

There's a flaw with getting information from the execution context of
the ping. For the soft-deprecated "throw a promise" technique, this is a
bit unreliable because you could in theory throw the same one multiple
times. Similarly, a more fundamental flaw with that API is that it
doesn't allow for tracking the information of Promises that are already
synchronously able to resolve.

This stops tracking the async debug info in the case of throwing a
Promise and only when you render a Promise. That means some loss of data
but we should just warn for throwing a Promise anyway.

Instead, this also adds support for tracking `use()`d thenables and
forwarding `_debugInfo` from then. This is done by extracting the info
from the Promise after the fact instead of in the resolve so that it
only happens once at the end after the pings are done.

This also supports passing the same Promise in multiple places and
tracking the debug info at each location, even if it was already
instrumented with a synchronous value by the time of the second use.

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

@@ -2991,6 +2991,64 @@ describe('ReactFlight', () => {
     );
   });
 
+  // @gate !__DEV__ || enableComponentPerformanceTrack
+  it('preserves debug info for server-to-server through use()', async () => {
+    function ThirdPartyComponent() {
+      return 'hi';
+    }
+
+    function ServerComponent({transport}) {
+      // This is a Server Component that receives other Server Components from a third party.
+      const text = ReactServer.use(ReactNoopFlightClient.read(transport));
+      return <div>{text.toUpperCase()}</div>;
+    }
+
+    const thirdPartyTransport = ReactNoopFlightServer.render(
+      <ThirdPartyComponent />,
+      {
+        environmentName: 'third-party',
+      },
+    );
+
+    const transport = ReactNoopFlightServer.render(
+      <ServerComponent transport={thirdPartyTransport} />,
+    );
+
+    await act(async () => {
+      const promise = ReactNoopFlightClient.read(transport);
+      expect(getDebugInfo(promise)).toEqual(
+        __DEV__
+          ? [
+              {time: 16},
+              {
+                name: 'ServerComponent',
+                env: 'Server',
+                key: null,
+                stack: '    in Object.<anonymous> (at **)',
+                props: {
+                  transport: expect.arrayContaining([]),
+                },
+              },
+              {time: 16},
+              {
+                name: 'ThirdPartyComponent',
+                env: 'third-party',
+                key: null,
+                stack: '    in Object.<anonymous> (at **)',
+                props: {},
+              },
+              {time: 16},
+              {time: 17},
+            ]
+          : undefined,
+      );
+      const result = await promise;
+      ReactNoop.render(result);
+    });
+
+    expect(ReactNoop).toMatchRenderedOutput(<div>HI</div>);
+  });
+
   it('preserves error stacks passed through server-to-server with source maps', async () => {
     async function ServerComponent({transport}) {
       // This is a Server Component that receives other Server Components from a third party.

packages/react-server-dom-turbopack/src/ReactFlightTurbopackReferences.js

@@ -161,6 +161,9 @@ const deepProxyHandlers = {
       // reference.
       case 'defaultProps':
         return undefined;
+      // React looks for debugInfo on thenables.
+      case '_debugInfo':
+        return undefined;
       // Avoid this attempting to be serialized.
       case 'toJSON':
         return undefined;
@@ -210,6 +213,9 @@ function getReference(target: Function, name: string | symbol): $FlowFixMe {
     // reference.
     case 'defaultProps':
       return undefined;
+    // React looks for debugInfo on thenables.
+    case '_debugInfo':
+      return undefined;
     // Avoid this attempting to be serialized.
     case 'toJSON':
       return undefined;

packages/react-server-dom-webpack/src/ReactFlightWebpackReferences.js

@@ -162,6 +162,9 @@ const deepProxyHandlers = {
       // reference.
       case 'defaultProps':
         return undefined;
+      // React looks for debugInfo on thenables.
+      case '_debugInfo':
+        return undefined;
       // Avoid this attempting to be serialized.
       case 'toJSON':
         return undefined;
@@ -211,6 +214,9 @@ function getReference(target: Function, name: string | symbol): $FlowFixMe {
     // reference.
     case 'defaultProps':
       return undefined;
+    // React looks for debugInfo on thenables.
+    case '_debugInfo':
+      return undefined;
     // Avoid this attempting to be serialized.
     case 'toJSON':
       return undefined;

packages/react-server/src/ReactFlightHooks.js

@@ -58,6 +58,12 @@ export function getThenableStateAfterSuspending(): ThenableState {
   return state;
 }
 
+export function getTrackedThenablesAfterRendering(): null | Array<
+  Thenable<any>,
+> {
+  return thenableState;
+}
+
 export const HooksDispatcher: Dispatcher = {
   readContext: (unsupportedContext: any),
 

packages/react-server/src/ReactFlightServer.js

@@ -91,6 +91,7 @@ import {
   initAsyncDebugInfo,
   markAsyncSequenceRootTask,
   getCurrentAsyncSequence,
+  getAsyncSequenceFromPromise,
   parseStackTrace,
   supportsComponentStorage,
   componentStorage,
@@ -106,6 +107,7 @@ import {
   prepareToUseHooksForRequest,
   prepareToUseHooksForComponent,
   getThenableStateAfterSuspending,
+  getTrackedThenablesAfterRendering,
   resetHooksForRequest,
 } from './ReactFlightHooks';
 import {DefaultAsyncDispatcher} from './flight/ReactFlightAsyncDispatcher';
@@ -690,26 +692,14 @@ function serializeThenable(
 
   switch (thenable.status) {
     case 'fulfilled': {
-      if (__DEV__) {
-        // If this came from Flight, forward any debug info into this new row.
-        const debugInfo: ?ReactDebugInfo = (thenable: any)._debugInfo;
-        if (debugInfo) {
-          forwardDebugInfo(request, newTask, debugInfo);
-        }
-      }
+      forwardDebugInfoFromThenable(request, newTask, thenable, null, null);
       // We have the resolved value, we can go ahead and schedule it for serialization.
       newTask.model = thenable.value;
       pingTask(request, newTask);
       return newTask.id;
     }
     case 'rejected': {
-      if (__DEV__) {
-        // If this came from Flight, forward any debug info into this new row.
-        const debugInfo: ?ReactDebugInfo = (thenable: any)._debugInfo;
-        if (debugInfo) {
-          forwardDebugInfo(request, newTask, debugInfo);
-        }
-      }
+      forwardDebugInfoFromThenable(request, newTask, thenable, null, null);
       const x = thenable.reason;
       erroredTask(request, newTask, x);
       return newTask.id;
@@ -758,24 +748,11 @@ function serializeThenable(
 
   thenable.then(
     value => {
-      if (__DEV__) {
-        // If this came from Flight, forward any debug info into this new row.
-        const debugInfo: ?ReactDebugInfo = (thenable: any)._debugInfo;
-        if (debugInfo) {
-          forwardDebugInfo(request, newTask, debugInfo);
-        }
-      }
+      forwardDebugInfoFromCurrentContext(request, newTask, thenable);
       newTask.model = value;
       pingTask(request, newTask);
     },
     reason => {
-      if (__DEV__) {
-        // If this came from Flight, forward any debug info into this new row.
-        const debugInfo: ?ReactDebugInfo = (thenable: any)._debugInfo;
-        if (debugInfo) {
-          forwardDebugInfo(request, newTask, debugInfo);
-        }
-      }
       if (newTask.status === PENDING) {
         if (enableProfilerTimer && enableComponentPerformanceTrack) {
           // If this is async we need to time when this task finishes.
@@ -1055,13 +1032,21 @@ function readThenable<T>(thenable: Thenable<T>): T {
   throw thenable;
 }
 
-function createLazyWrapperAroundWakeable(wakeable: Wakeable) {
+function createLazyWrapperAroundWakeable(
+  request: Request,
+  task: Task,
+  wakeable: Wakeable,
+) {
   // This is a temporary fork of the `use` implementation until we accept
   // promises everywhere.
   const thenable: Thenable<mixed> = (wakeable: any);
   switch (thenable.status) {
-    case 'fulfilled':
+    case 'fulfilled': {
+      forwardDebugInfoFromThenable(request, task, thenable, null, null);
+      return thenable.value;
+    }
     case 'rejected':
+      forwardDebugInfoFromThenable(request, task, thenable, null, null);
       break;
     default: {
       if (typeof thenable.status === 'string') {
@@ -1074,6 +1059,7 @@ function createLazyWrapperAroundWakeable(wakeable: Wakeable) {
       pendingThenable.status = 'pending';
       pendingThenable.then(
         fulfilledValue => {
+          forwardDebugInfoFromCurrentContext(request, task, thenable);
           if (thenable.status === 'pending') {
             const fulfilledThenable: FulfilledThenable<mixed> = (thenable: any);
             fulfilledThenable.status = 'fulfilled';
@@ -1081,6 +1067,7 @@ function createLazyWrapperAroundWakeable(wakeable: Wakeable) {
           }
         },
         (error: mixed) => {
+          forwardDebugInfoFromCurrentContext(request, task, thenable);
           if (thenable.status === 'pending') {
             const rejectedThenable: RejectedThenable<mixed> = (thenable: any);
             rejectedThenable.status = 'rejected';
@@ -1096,10 +1083,6 @@ function createLazyWrapperAroundWakeable(wakeable: Wakeable) {
     _payload: thenable,
     _init: readThenable,
   };
-  if (__DEV__) {
-    // If this came from React, transfer the debug info.
-    lazyType._debugInfo = (thenable: any)._debugInfo || [];
-  }
   return lazyType;
 }
 
@@ -1178,12 +1161,9 @@ function processServerComponentReturnValue(
         }
       }, voidHandler);
     }
-    if (thenable.status === 'fulfilled') {
-      return thenable.value;
-    }
     // TODO: Once we accept Promises as children on the client, we can just return
     // the thenable here.
-    return createLazyWrapperAroundWakeable(result);
+    return createLazyWrapperAroundWakeable(request, task, result);
   }
 
   if (__DEV__) {
@@ -1386,6 +1366,7 @@ function renderFunctionComponent<Props>(
       }
     }
   } else {
+    componentDebugInfo = (null: any);
     prepareToUseHooksForComponent(prevThenableState, null);
     // The secondArg is always undefined in Server Components since refs error early.
     const secondArg = undefined;
@@ -1408,6 +1389,34 @@ function renderFunctionComponent<Props>(
     throw null;
   }
 
+  if (
+    __DEV__ ||
+    (enableProfilerTimer &&
+      enableComponentPerformanceTrack &&
+      enableAsyncDebugInfo)
+  ) {
+    // Forward any debug information for any Promises that we use():ed during the render.
+    // We do this at the end so that we don't keep doing this for each retry.
+    const trackedThenables = getTrackedThenablesAfterRendering();
+    if (trackedThenables !== null) {
+      const stacks: Array<Error> =
+        __DEV__ && enableAsyncDebugInfo
+          ? (trackedThenables: any)._stacks ||
+            ((trackedThenables: any)._stacks = [])
+          : (null: any);
+      for (let i = 0; i < trackedThenables.length; i++) {
+        const stack = __DEV__ && enableAsyncDebugInfo ? stacks[i] : null;
+        forwardDebugInfoFromThenable(
+          request,
+          task,
+          trackedThenables[i],
+          __DEV__ ? componentDebugInfo : null,
+          stack,
+        );
+      }
+    }
+  }
+
   // Apply special cases.
   result = processServerComponentReturnValue(request, task, Component, result);
 
@@ -1884,7 +1893,7 @@ function visitAsyncNode(
   request: Request,
   task: Task,
   node: AsyncSequence,
-  visited: Set<AsyncSequence>,
+  visited: Set<AsyncSequence | ReactDebugInfo>,
   cutOff: number,
 ): null | PromiseNode | IONode {
   if (visited.has(node)) {
@@ -1943,7 +1952,8 @@ function visitAsyncNode(
       // We need to forward after we visit awaited nodes because what ever I/O we requested that's
       // the thing that generated this node and its virtual children.
       const debugInfo = node.debugInfo;
-      if (debugInfo !== null) {
+      if (debugInfo !== null && !visited.has(debugInfo)) {
+        visited.add(debugInfo);
         forwardDebugInfo(request, task, debugInfo);
       }
       return match;
@@ -2003,8 +2013,9 @@ function visitAsyncNode(
       }
       // We need to forward after we visit awaited nodes because what ever I/O we requested that's
       // the thing that generated this node and its virtual children.
-      const debugInfo: null | ReactDebugInfo = node.debugInfo;
-      if (debugInfo !== null) {
+      const debugInfo = node.debugInfo;
+      if (debugInfo !== null && !visited.has(debugInfo)) {
+        visited.add(debugInfo);
         forwardDebugInfo(request, task, debugInfo);
       }
       return match;
@@ -2020,8 +2031,14 @@ function emitAsyncSequence(
   request: Request,
   task: Task,
   node: AsyncSequence,
+  alreadyForwardedDebugInfo: ?ReactDebugInfo,
+  owner: null | ReactComponentInfo,
+  stack: null | Error,
 ): void {
-  const visited: Set<AsyncSequence> = new Set();
+  const visited: Set<AsyncSequence | ReactDebugInfo> = new Set();
+  if (__DEV__ && alreadyForwardedDebugInfo) {
+    visited.add(alreadyForwardedDebugInfo);
+  }
   const awaitedNode = visitAsyncNode(request, task, node, visited, task.time);
   if (awaitedNode !== null) {
     // Nothing in user space (unfiltered stack) awaited this.
@@ -2032,10 +2049,21 @@ function emitAsyncSequence(
     const env = (0, request.environmentName)();
     // If we don't have any thing awaited, the time we started awaiting was internal
     // when we yielded after rendering. The current task time is basically that.
-    emitDebugChunk(request, task.id, {
+    const debugInfo: ReactAsyncInfo = {
       awaited: ((awaitedNode: any): ReactIOInfo), // This is deduped by this reference.
       env: env,
-    });
+    };
+    if (__DEV__) {
+      if (owner != null) {
+        // $FlowFixMe[cannot-write]
+        debugInfo.owner = owner;
+      }
+      if (stack != null) {
+        // $FlowFixMe[cannot-write]
+        debugInfo.stack = filterStackTrace(request, parseStackTrace(stack, 1));
+      }
+    }
+    emitDebugChunk(request, task.id, debugInfo);
     markOperationEndTime(request, task, awaitedNode.end);
   }
 }
@@ -2044,12 +2072,6 @@ function pingTask(request: Request, task: Task): void {
   if (enableProfilerTimer && enableComponentPerformanceTrack) {
     // If this was async we need to emit the time when it completes.
     task.timed = true;
-    if (enableAsyncDebugInfo) {
-      const sequence = getCurrentAsyncSequence();
-      if (sequence !== null) {
-        emitAsyncSequence(request, task, sequence);
-      }
-    }
   }
   const pingedTasks = request.pingedTasks;
   pingedTasks.push(task);
@@ -4316,6 +4338,58 @@ function forwardDebugInfo(
   }
 }
 
+function forwardDebugInfoFromThenable(
+  request: Request,
+  task: Task,
+  thenable: Thenable<any>,
+  owner: null | ReactComponentInfo, // DEV-only
+  stack: null | Error, // DEV-only
+): void {
+  let debugInfo: ?ReactDebugInfo;
+  if (__DEV__) {
+    // If this came from Flight, forward any debug info into this new row.
+    debugInfo = thenable._debugInfo;
+    if (debugInfo) {
+      forwardDebugInfo(request, task, debugInfo);
+    }
+  }
+  if (
+    enableProfilerTimer &&
+    enableComponentPerformanceTrack &&
+    enableAsyncDebugInfo
+  ) {
+    const sequence = getAsyncSequenceFromPromise(thenable);
+    if (sequence !== null) {
+      emitAsyncSequence(request, task, sequence, debugInfo, owner, stack);
+    }
+  }
+}
+
+function forwardDebugInfoFromCurrentContext(
+  request: Request,
+  task: Task,
+  thenable: Thenable<any>,
+): void {
+  let debugInfo: ?ReactDebugInfo;
+  if (__DEV__) {
+    // If this came from Flight, forward any debug info into this new row.
+    debugInfo = thenable._debugInfo;
+    if (debugInfo) {
+      forwardDebugInfo(request, task, debugInfo);
+    }
+  }
+  if (
+    enableProfilerTimer &&
+    enableComponentPerformanceTrack &&
+    enableAsyncDebugInfo
+  ) {
+    const sequence = getCurrentAsyncSequence();
+    if (sequence !== null) {
+      emitAsyncSequence(request, task, sequence, debugInfo, null, null);
+    }
+  }
+}
+
 function emitTimingChunk(
   request: Request,
   id: number,

packages/react-server/src/ReactFlightServerConfigDebugNode.js

@@ -24,9 +24,12 @@ import {
   UNRESOLVED_AWAIT_NODE,
 } from './ReactFlightAsyncSequence';
 import {resolveOwner} from './flight/ReactFlightCurrentOwner';
-import {createHook, executionAsyncId} from 'async_hooks';
+import {createHook, executionAsyncId, AsyncResource} from 'async_hooks';
 import {enableAsyncDebugInfo} from 'shared/ReactFeatureFlags';
 
+// $FlowFixMe[method-unbinding]
+const getAsyncId = AsyncResource.prototype.asyncId;
+
 const pendingOperations: Map<number, AsyncSequence> =
   __DEV__ && enableAsyncDebugInfo ? new Map() : (null: any);
 
@@ -260,3 +263,29 @@ export function getCurrentAsyncSequence(): null | AsyncSequence {
   }
   return currentNode;
 }
+
+export function getAsyncSequenceFromPromise(
+  promise: any,
+): null | AsyncSequence {
+  if (!__DEV__ || !enableAsyncDebugInfo) {
+    return null;
+  }
+  // A Promise is conceptually an AsyncResource but doesn't have its own methods.
+  // We use this hack to extract the internal asyncId off the Promise.
+  let asyncId: void | number;
+  try {
+    asyncId = getAsyncId.call(promise);
+  } catch (x) {
+    // Ignore errors extracting the ID. We treat it as missing.
+    // This could happen if our hack stops working or in the case where this is
+    // a Proxy that throws such as our own ClientReference proxies.
+  }
+  if (asyncId === undefined) {
+    return null;
+  }
+  const node = pendingOperations.get(asyncId);
+  if (node === undefined) {
+    return null;
+  }
+  return node;
+}

packages/react-server/src/ReactFlightServerConfigDebugNoop.js

@@ -15,3 +15,8 @@ export function markAsyncSequenceRootTask(): void {}
 export function getCurrentAsyncSequence(): null | AsyncSequence {
   return null;
 }
+export function getAsyncSequenceFromPromise(
+  promise: any,
+): null | AsyncSequence {
+  return null;
+}

packages/react-server/src/ReactFlightServerTemporaryReferences.js

@@ -52,6 +52,9 @@ const proxyHandlers = {
       // reference.
       case 'defaultProps':
         return undefined;
+      // React looks for debugInfo on thenables.
+      case '_debugInfo':
+        return undefined;
       // Avoid this attempting to be serialized.
       case 'toJSON':
         return undefined;

packages/react-server/src/ReactFlightThenable.js

@@ -20,9 +20,11 @@ import type {
   RejectedThenable,
 } from 'shared/ReactTypes';
 
+import {enableAsyncDebugInfo} from 'shared/ReactFeatureFlags';
+
 import noop from 'shared/noop';
 
-export opaque type ThenableState = Array<Thenable<any>>;
+export type ThenableState = Array<Thenable<any>>;
 
 // An error that is thrown (e.g. by `use`) to trigger Suspense. If we
 // detect this is caught by userspace, we'll log a warning in development.
@@ -50,6 +52,11 @@ export function trackUsedThenable<T>(
   const previous = thenableState[index];
   if (previous === undefined) {
     thenableState.push(thenable);
+    if (__DEV__ && enableAsyncDebugInfo) {
+      const stacks: Array<Error> =
+        (thenableState: any)._stacks || ((thenableState: any)._stacks = []);
+      stacks.push(new Error());
+    }
   } else {
     if (previous !== thenable) {
       // Reuse the previous thenable, and drop the new one. We can assume

scripts/flow/environment.js

@@ -356,7 +356,9 @@ declare module 'async_hooks' {
     run<R>(store: T, callback: (...args: any[]) => R, ...args: any[]): R;
     enterWith(store: T): void;
   }
-  declare interface AsyncResource {}
+  declare class AsyncResource {
+    asyncId(): number;
+  }
   declare function executionAsyncId(): number;
   declare function executionAsyncResource(): AsyncResource;
   declare function triggerAsyncId(): number;