[Flight] Forward debugInfo from awaited instrumented Promises (#33415)

Stacked on #33403.

When a Promise is coming from React such as when it's passed from
another environment, we should forward the debug information from that
environment. We already do that when rendered as a child.

This makes it possible to also `await promise` and have the information
from that instrumented promise carry through to the next render.

This is a bit tricky because the current protocol is that we have to
read it from the Promise after it resolves so it has time to be assigned
to the promise. `async_hooks` doesn't pass us the instance (even though
it has it) when it gets resolved so we need to keep it around. However,
we have to be very careful because if we get this wrong it'll cause a
memory leak since we retain things by `asyncId` and then manually listen
for `destroy()` which can only be called once a Promise is GC:ed, which
it can't be if we retain it. We have to therefore use a `WeakRef` in
case it never resolves, and then read the `_debugInfo` when it resolves.
We could maybe install a setter or something instead but that's also
heavy.

The other issues is that we don't use native Promises in
ReactFlightClient so our instrumented promises aren't picked up by the
`async_hooks` implementation and so we never get a handle to our
thenable instance. To solve this we can create a native wrapper only in
DEV.

Author: sebmarkbage

.eslintrc.js

@@ -561,6 +561,7 @@ module.exports = {
     ConsoleTask: 'readonly', // TOOD: Figure out what the official name of this will be.
     ReturnType: 'readonly',
     AnimationFrameID: 'readonly',
+    WeakRef: 'readonly',
     // For Flow type annotation. Only `BigInt` is valid at runtime.
     bigint: 'readonly',
     BigInt: 'readonly',

packages/react-client/src/ReactFlightClient.js

@@ -266,6 +266,27 @@ ReactPromise.prototype.then = function <T>(
       initializeModuleChunk(chunk);
       break;
   }
+  if (__DEV__ && enableAsyncDebugInfo) {
+    // Because only native Promises get picked up when we're awaiting we need to wrap
+    // this in a native Promise in DEV. This means that these callbacks are no longer sync
+    // but the lazy initialization is still sync and the .value can be inspected after,
+    // allowing it to be read synchronously anyway.
+    const resolveCallback = resolve;
+    const rejectCallback = reject;
+    const wrapperPromise: Promise<T> = new Promise((res, rej) => {
+      resolve = value => {
+        // $FlowFixMe
+        wrapperPromise._debugInfo = this._debugInfo;
+        res(value);
+      };
+      reject = reason => {
+        // $FlowFixMe
+        wrapperPromise._debugInfo = this._debugInfo;
+        rej(reason);
+      };
+    });
+    wrapperPromise.then(resolveCallback, rejectCallback);
+  }
   // The status might have changed after initialization.
   switch (chunk.status) {
     case INITIALIZED:

packages/react-server/src/ReactFlightAsyncSequence.js

@@ -7,25 +7,33 @@
  * @flow
  */
 
-import type {ReactComponentInfo} from 'shared/ReactTypes';
+import type {ReactDebugInfo, ReactComponentInfo} from 'shared/ReactTypes';
 
 export const IO_NODE = 0;
 export const PROMISE_NODE = 1;
 export const AWAIT_NODE = 2;
+export const UNRESOLVED_PROMISE_NODE = 3;
+export const UNRESOLVED_AWAIT_NODE = 4;
+
+type PromiseWithDebugInfo = interface extends Promise<any> {
+  _debugInfo?: ReactDebugInfo,
+};
 
 export type IONode = {
   tag: 0,
   owner: null | ReactComponentInfo,
   stack: Error, // callsite that spawned the I/O
+  debugInfo: null, // not used on 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.
   awaited: null, // I/O is only blocked on external.
-  previous: null | AwaitNode, // the preceeding await that spawned this new work
+  previous: null | AwaitNode | UnresolvedAwaitNode, // the preceeding await that spawned this new work
 };
 
 export type PromiseNode = {
   tag: 1,
   owner: null | ReactComponentInfo,
+  debugInfo: null | ReactDebugInfo, // forwarded debugInfo from the Promise
   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.
@@ -36,11 +44,39 @@ export type PromiseNode = {
 export type AwaitNode = {
   tag: 2,
   owner: null | ReactComponentInfo,
+  debugInfo: null | ReactDebugInfo, // forwarded debugInfo from the Promise
   stack: Error, // callsite that awaited (using await, .then(), Promise.all(), ...)
   start: number, // when we started blocking. This might be later than the I/O started.
   end: number, // when we unblocked. This might be later than the I/O resolved if there's CPU time.
   awaited: null | AsyncSequence, // the promise we were waiting on
   previous: null | AsyncSequence, // the sequence that was blocking us from awaiting in the first place
 };
 
-export type AsyncSequence = IONode | PromiseNode | AwaitNode;
+export type UnresolvedPromiseNode = {
+  tag: 3,
+  owner: null | ReactComponentInfo,
+  debugInfo: WeakRef<PromiseWithDebugInfo>, // holds onto the Promise until we can extract debugInfo when it resolves
+  stack: Error, // callsite that created the Promise
+  start: number, // start time when the Promise was created
+  end: -1.1, // set when we resolve.
+  awaited: null | AsyncSequence, // the thing that ended up resolving this promise
+  previous: null, // where we created the promise is not interesting since creating it doesn't mean waiting.
+};
+
+export type UnresolvedAwaitNode = {
+  tag: 4,
+  owner: null | ReactComponentInfo,
+  debugInfo: WeakRef<PromiseWithDebugInfo>, // holds onto the Promise until we can extract debugInfo when it resolves
+  stack: Error, // callsite that awaited (using await, .then(), Promise.all(), ...)
+  start: number, // when we started blocking. This might be later than the I/O started.
+  end: -1.1, // set when we resolve.
+  awaited: null | AsyncSequence, // the promise we were waiting on
+  previous: null | AsyncSequence, // the sequence that was blocking us from awaiting in the first place
+};
+
+export type AsyncSequence =
+  | IONode
+  | PromiseNode
+  | AwaitNode
+  | UnresolvedPromiseNode
+  | UnresolvedAwaitNode;

packages/react-server/src/ReactFlightServer.js

@@ -89,6 +89,7 @@ import {
   requestStorage,
   createHints,
   initAsyncDebugInfo,
+  markAsyncSequenceRootTask,
   getCurrentAsyncSequence,
   parseStackTrace,
   supportsComponentStorage,
@@ -149,7 +150,13 @@ import binaryToComparableString from 'shared/binaryToComparableString';
 
 import {SuspenseException, getSuspendedThenable} from './ReactFlightThenable';
 
-import {IO_NODE, PROMISE_NODE, AWAIT_NODE} from './ReactFlightAsyncSequence';
+import {
+  IO_NODE,
+  PROMISE_NODE,
+  AWAIT_NODE,
+  UNRESOLVED_AWAIT_NODE,
+  UNRESOLVED_PROMISE_NODE,
+} from './ReactFlightAsyncSequence';
 
 // DEV-only set containing internal objects that should not be limited and turned into getters.
 const doNotLimit: WeakSet<Reference> = __DEV__ ? new WeakSet() : (null: any);
@@ -1879,6 +1886,9 @@ function visitAsyncNode(
     case IO_NODE: {
       return node;
     }
+    case UNRESOLVED_PROMISE_NODE: {
+      return null;
+    }
     case PROMISE_NODE: {
       if (node.end < cutOff) {
         // This was already resolved when we started this sequence. It must have been
@@ -1888,6 +1898,7 @@ function visitAsyncNode(
         return null;
       }
       const awaited = node.awaited;
+      let match = null;
       if (awaited !== null) {
         const ioNode = visitAsyncNode(request, task, awaited, cutOff, visited);
         if (ioNode !== null) {
@@ -1907,72 +1918,104 @@ function visitAsyncNode(
               // If we haven't defined an end time, use the resolve of the outer Promise.
               ioNode.end = node.end;
             }
-            return ioNode;
+            match = ioNode;
+          } else {
+            match = node;
           }
-          return node;
         }
       }
-      return null;
+      // 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) {
+        forwardDebugInfo(request, task.id, debugInfo);
+      }
+      return match;
     }
+    case UNRESOLVED_AWAIT_NODE:
+    // We could be inside the .then() which is about to resolve this node.
+    // TODO: We could call emitAsyncSequence in a microtask to avoid this issue.
+    // Fallthrough to the resolved path.
     case AWAIT_NODE: {
       const awaited = node.awaited;
+      let match = null;
       if (awaited !== null) {
         const ioNode = visitAsyncNode(request, task, awaited, cutOff, visited);
         if (ioNode !== null) {
-          if (node.end < 0) {
+          let endTime: number;
+          if (node.tag === UNRESOLVED_AWAIT_NODE) {
             // If we haven't defined an end time, use the resolve of the inner Promise.
             // This can happen because the ping gets invoked before the await gets resolved.
             if (ioNode.end < node.start) {
               // If we're awaiting a resolved Promise it could have finished before we started.
-              node.end = node.start;
+              endTime = node.start;
             } else {
-              node.end = ioNode.end;
+              endTime = ioNode.end;
             }
+          } else {
+            endTime = node.end;
           }
-          if (node.end < cutOff) {
+          if (endTime < cutOff) {
             // This was already resolved when we started this sequence. It must have been
             // part of a different component.
             // TODO: Think of some other way to exclude irrelevant data since if we awaited
             // a cached promise, we should still log this component as being dependent on that data.
-            return null;
-          }
-
-          const stack = filterStackTrace(
-            request,
-            parseStackTrace(node.stack, 1),
-          );
-          if (stack.length === 0) {
-            // If this await was fully filtered out, then it was inside third party code
-            // such as in an external library. We return the I/O node and try another await.
-            return ioNode;
-          }
-          // Outline the IO node.
-          serializeIONode(request, ioNode);
-          // We log the environment at the time when the last promise pigned ping which may
-          // be later than what the environment was when we actually started awaiting.
-          const env = (0, request.environmentName)();
-          if (node.start <= cutOff) {
-            // If this was an await that started before this sequence but finished after,
-            // then we clamp it to the start of this sequence. We don't need to emit a time
-            // TODO: Typically we'll already have a previous time stamp with the cutOff time
-            // so we shouldn't need to emit another one. But not always.
-            emitTimingChunk(request, task.id, cutOff);
           } else {
-            emitTimingChunk(request, task.id, node.start);
+            const stack = filterStackTrace(
+              request,
+              parseStackTrace(node.stack, 1),
+            );
+            if (stack.length === 0) {
+              // If this await was fully filtered out, then it was inside third party code
+              // such as in an external library. We return the I/O node and try another await.
+              match = ioNode;
+            } else {
+              // Outline the IO node.
+              if (ioNode.end < 0) {
+                ioNode.end = endTime;
+              }
+              serializeIONode(request, ioNode);
+              // We log the environment at the time when the last promise pigned ping which may
+              // be later than what the environment was when we actually started awaiting.
+              const env = (0, request.environmentName)();
+              if (node.start <= cutOff) {
+                // If this was an await that started before this sequence but finished after,
+                // then we clamp it to the start of this sequence. We don't need to emit a time
+                // TODO: Typically we'll already have a previous time stamp with the cutOff time
+                // so we shouldn't need to emit another one. But not always.
+                emitTimingChunk(request, task.id, cutOff);
+              } else {
+                emitTimingChunk(request, task.id, node.start);
+              }
+              // Then emit a reference to us awaiting it in the current task.
+              request.pendingChunks++;
+              emitDebugChunk(request, task.id, {
+                awaited: ((ioNode: any): ReactIOInfo), // This is deduped by this reference.
+                env: env,
+                owner: node.owner,
+                stack: stack,
+              });
+              emitTimingChunk(request, task.id, node.end);
+            }
           }
-          // Then emit a reference to us awaiting it in the current task.
-          request.pendingChunks++;
-          emitDebugChunk(request, task.id, {
-            awaited: ((ioNode: any): ReactIOInfo), // This is deduped by this reference.
-            env: env,
-            owner: node.owner,
-            stack: stack,
-          });
-          emitTimingChunk(request, task.id, node.end);
         }
       }
-      // If we had awaited anything we would have written it now.
-      return null;
+      // 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.
+      let debugInfo: null | ReactDebugInfo;
+      if (node.tag === UNRESOLVED_AWAIT_NODE) {
+        const promise = node.debugInfo.deref();
+        debugInfo =
+          promise === undefined || promise._debugInfo === undefined
+            ? null
+            : promise._debugInfo;
+      } else {
+        debugInfo = node.debugInfo;
+      }
+      if (debugInfo !== null) {
+        forwardDebugInfo(request, task.id, debugInfo);
+      }
+      return match;
     }
     default: {
       // eslint-disable-next-line react-internal/prod-error-codes
@@ -4513,6 +4556,8 @@ function tryStreamTask(request: Request, task: Task): void {
 }
 
 function performWork(request: Request): void {
+  markAsyncSequenceRootTask();
+
   const prevDispatcher = ReactSharedInternals.H;
   ReactSharedInternals.H = HooksDispatcher;
   const prevRequest = currentRequest;