Mark await end time when they begin

promiseResolve, for already resolved promises, is invoked at the end of
its execution phase. That means that any awaits that happen inside would
have their start time below the end time which cuts them off.

Author: sebmarkbage

packages/react-server/src/ReactFlightServer.js

@@ -1933,10 +1933,6 @@ function visitAsyncNode(
             // If this Promise was created inside only third party code, then try to use
             // the inner I/O node instead. This could happen if third party calls into first
             // party to perform some I/O.
-            if (ioNode.end < 0) {
-              // If we haven't defined an end time, use the resolve of the outer Promise.
-              ioNode.end = node.end;
-            }
             match = ioNode;
           } else {
             match = node;
@@ -1951,30 +1947,17 @@ function visitAsyncNode(
       }
       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 UNRESOLVED_AWAIT_NODE: {
+      return null;
+    }
     case AWAIT_NODE: {
       const awaited = node.awaited;
       let match = null;
       if (awaited !== null) {
         const ioNode = visitAsyncNode(request, task, awaited, visited);
         if (ioNode !== null) {
           const startTime: number = node.start;
-          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.
-              endTime = node.start;
-            } else {
-              endTime = ioNode.end;
-            }
-          } else {
-            endTime = node.end;
-          }
+          const endTime: number = node.end;
           if (endTime <= request.timeOrigin) {
             // This was already resolved when we started this render. It must have been either something
             // that's part of a start up sequence or externally cached data. We exclude that information.
@@ -1998,10 +1981,8 @@ function visitAsyncNode(
               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)();
@@ -2021,16 +2002,7 @@ 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.
-      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;
-      }
+      const debugInfo: null | ReactDebugInfo = node.debugInfo;
       if (debugInfo !== null) {
         forwardDebugInfo(request, task, debugInfo);
       }
@@ -2052,14 +2024,6 @@ function emitAsyncSequence(
   const awaitedNode = visitAsyncNode(request, task, node, visited);
   if (awaitedNode !== null) {
     // Nothing in user space (unfiltered stack) awaited this.
-    if (awaitedNode.end < 0) {
-      // If this was I/O directly without a Promise, then it means that some custom Thenable
-      // called our ping directly and not from a native .then(). We use the current ping time
-      // as the end time and treat it as an await with no stack.
-      // TODO: If this I/O is recurring then we really should have different entries for
-      // each occurrence. Right now we'll only track the first time it is invoked.
-      awaitedNode.end = performance.now();
-    }
     serializeIONode(request, awaitedNode);
     request.pendingChunks++;
     // We log the environment at the time when we ping which may be later than what the

packages/react-server/src/ReactFlightServerConfigDebugNode.js

@@ -30,6 +30,24 @@ import {enableAsyncDebugInfo} from 'shared/ReactFeatureFlags';
 const pendingOperations: Map<number, AsyncSequence> =
   __DEV__ && enableAsyncDebugInfo ? new Map() : (null: any);
 
+function resolvePromiseOrAwaitNode(
+  unresolvedNode: UnresolvedAwaitNode | UnresolvedPromiseNode,
+  endTime: number,
+): AwaitNode | PromiseNode {
+  const resolvedNode: AwaitNode | PromiseNode = (unresolvedNode: any);
+  resolvedNode.tag = ((unresolvedNode.tag === UNRESOLVED_PROMISE_NODE
+    ? PROMISE_NODE
+    : AWAIT_NODE): any);
+  // The Promise can be garbage collected after this so we should extract debugInfo first.
+  const promise = unresolvedNode.debugInfo.deref();
+  resolvedNode.debugInfo =
+    promise === undefined || promise._debugInfo === undefined
+      ? null
+      : promise._debugInfo;
+  resolvedNode.end = endTime;
+  return resolvedNode;
+}
+
 // Initialize the tracing of async operations.
 // We do this globally since the async work can potentially eagerly
 // start before the first request and once requests start they can interleave.
@@ -129,42 +147,50 @@ export function initAsyncDebugInfo(): void {
         }
         pendingOperations.set(asyncId, node);
       },
+      before(asyncId: number): void {
+        const node = pendingOperations.get(asyncId);
+        if (node !== undefined) {
+          switch (node.tag) {
+            case IO_NODE: {
+              // Log the end time when we resolved the I/O. This can happen
+              // more than once if it's a recurring resource like a connection.
+              const ioNode: IONode = (node: any);
+              ioNode.end = performance.now();
+              break;
+            }
+            case UNRESOLVED_AWAIT_NODE:
+            case UNRESOLVED_PROMISE_NODE: {
+              // If we begin before we resolve, that means that this is actually already resolved but
+              // the promiseResolve hook is called at the end of the execution. This means that it was
+              // actually already resolved when we started so we just use the start time as the end time.
+              resolvePromiseOrAwaitNode(node, node.start);
+              break;
+            }
+          }
+        }
+      },
       promiseResolve(asyncId: number): void {
         const node = pendingOperations.get(asyncId);
         if (node !== undefined) {
           let resolvedNode: AwaitNode | PromiseNode;
           switch (node.tag) {
-            case UNRESOLVED_AWAIT_NODE: {
-              const awaitNode: AwaitNode = (node: any);
-              awaitNode.tag = AWAIT_NODE;
-              resolvedNode = awaitNode;
+            case UNRESOLVED_AWAIT_NODE:
+            case UNRESOLVED_PROMISE_NODE: {
+              resolvedNode = resolvePromiseOrAwaitNode(node, performance.now());
               break;
             }
-            case UNRESOLVED_PROMISE_NODE: {
-              const promiseNode: PromiseNode = (node: any);
-              promiseNode.tag = PROMISE_NODE;
-              resolvedNode = promiseNode;
+            case AWAIT_NODE:
+            case PROMISE_NODE: {
+              resolvedNode = node;
+              // We already resolved this in the begin phase.
               break;
             }
-            case IO_NODE:
-              // eslint-disable-next-line react-internal/prod-error-codes
-              throw new Error(
-                'A Promise should never be an IO_NODE. This is a bug in React.',
-              );
             default:
               // eslint-disable-next-line react-internal/prod-error-codes
               throw new Error(
-                'A Promise should never be resolved twice. This is a bug in React or Node.js.',
+                'A Promise should never be an IO_NODE. This is a bug in React.',
               );
           }
-          // Log the end time when we resolved the promise.
-          resolvedNode.end = performance.now();
-          // The Promise can be garbage collected after this so we should extract debugInfo first.
-          const promise = node.debugInfo.deref();
-          resolvedNode.debugInfo =
-            promise === undefined || promise._debugInfo === undefined
-              ? null
-              : promise._debugInfo;
           const currentAsyncId = executionAsyncId();
           if (asyncId !== currentAsyncId) {
             // If the promise was not resolved by itself, then that means that