[Flight] Emit the time we awaited something inside a Server Component

packages/react-client/src/ReactFlightClient.js

@@ -79,6 +79,7 @@ import {
   logDedupedComponentRender,
   logComponentErrored,
   logIOInfo,
+  logComponentAwait,
 } from './ReactFlightPerformanceTrack';
 
 import {
@@ -680,7 +681,7 @@ function getIOInfoTaskName(ioInfo: ReactIOInfo): string {
 }
 
 function getAsyncInfoTaskName(asyncInfo: ReactAsyncInfo): string {
-  return 'await'; // We could be smarter about this and give it a name like `then` or `Promise.all`.
+  return 'await ' + getIOInfoTaskName(asyncInfo.awaited);
 }
 
 function getServerComponentTaskName(componentInfo: ReactComponentInfo): string {
@@ -2971,9 +2972,12 @@ function flushComponentPerformance(
     for (let i = debugInfo.length - 1; i >= 0; i--) {
       const info = debugInfo[i];
       if (typeof info.time === 'number') {
-        endTime = info.time;
-        if (endTime > childrenEndTime) {
-          childrenEndTime = endTime;
+        if (info.time > childrenEndTime) {
+          childrenEndTime = info.time;
+        }
+        if (endTime === 0) {
+          // Last timestamp is the end of the last component.
+          endTime = info.time;
         }
       }
       if (typeof info.name === 'string' && i > 0) {
@@ -3011,8 +3015,29 @@ function flushComponentPerformance(
           }
           // Track the root most component of the result for deduping logging.
           result.component = componentInfo;
+          // Set the end time of the previous component to the start of the previous.
+          endTime = startTime;
         }
         isLastComponent = false;
+      } else if (info.awaited && i > 0 && i < debugInfo.length - 2) {
+        // $FlowFixMe: Refined.
+        const asyncInfo: ReactAsyncInfo = info;
+        const startTimeInfo = debugInfo[i - 1];
+        const endTimeInfo = debugInfo[i + 1];
+        if (
+          typeof startTimeInfo.time === 'number' &&
+          typeof endTimeInfo.time === 'number'
+        ) {
+          const awaitStartTime = startTimeInfo.time;
+          const awaitEndTime = endTimeInfo.time;
+          logComponentAwait(
+            asyncInfo,
+            trackIdx,
+            awaitStartTime,
+            awaitEndTime,
+            response._rootEnvironmentName,
+          );
+        }
       }
     }
   }

packages/react-client/src/ReactFlightPerformanceTrack.js

@@ -9,7 +9,11 @@
 
 /* eslint-disable react-internal/no-production-logging */
 
-import type {ReactComponentInfo, ReactIOInfo} from 'shared/ReactTypes';
+import type {
+  ReactComponentInfo,
+  ReactIOInfo,
+  ReactAsyncInfo,
+} from 'shared/ReactTypes';
 
 import {enableProfilerTimer} from 'shared/ReactFeatureFlags';
 
@@ -224,6 +228,48 @@ function getIOColor(
   }
 }
 
+export function logComponentAwait(
+  asyncInfo: ReactAsyncInfo,
+  trackIdx: number,
+  startTime: number,
+  endTime: number,
+  rootEnv: string,
+): void {
+  if (supportsUserTiming && endTime > 0) {
+    const env = asyncInfo.env;
+    const name = asyncInfo.awaited.name;
+    const isPrimaryEnv = env === rootEnv;
+    const color = getIOColor(name);
+    const entryName =
+      'await ' +
+      (isPrimaryEnv || env === undefined ? name : name + ' [' + env + ']');
+    const debugTask = asyncInfo.debugTask;
+    if (__DEV__ && debugTask) {
+      debugTask.run(
+        // $FlowFixMe[method-unbinding]
+        console.timeStamp.bind(
+          console,
+          entryName,
+          startTime < 0 ? 0 : startTime,
+          endTime,
+          trackNames[trackIdx],
+          COMPONENTS_TRACK,
+          color,
+        ),
+      );
+    } else {
+      console.timeStamp(
+        entryName,
+        startTime < 0 ? 0 : startTime,
+        endTime,
+        trackNames[trackIdx],
+        COMPONENTS_TRACK,
+        color,
+      );
+    }
+  }
+}
+
 export function logIOInfo(ioInfo: ReactIOInfo, rootEnv: string): void {
   const startTime = ioInfo.start;
   const endTime = ioInfo.end;

packages/react-server/src/ReactFlightAsyncSequence.js

@@ -37,8 +37,8 @@ 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.
+  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
 };

packages/react-server/src/ReactFlightServer.js

@@ -1919,6 +1919,24 @@ function visitAsyncNode(
       if (awaited !== null) {
         const ioNode = visitAsyncNode(request, task, awaited, cutOff, visited);
         if (ioNode !== null) {
+          if (node.end < 0) {
+            // 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;
+            } else {
+              node.end = ioNode.end;
+            }
+          }
+          if (node.end < 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),
@@ -1933,6 +1951,15 @@ function visitAsyncNode(
           // 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, {
@@ -1941,6 +1968,7 @@ function visitAsyncNode(
             owner: node.owner,
             stack: stack,
           });
+          emitTimingChunk(request, task.id, node.end);
         }
       }
       // If we had awaited anything we would have written it now.
@@ -1976,10 +2004,14 @@ function emitAsyncSequence(
     // We log the environment at the time when we ping which may be later than what the
     // environment was when we actually started awaiting.
     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 cutOff time is basically that.
+    emitTimingChunk(request, task.id, cutOff);
     emitDebugChunk(request, task.id, {
       awaited: ((awaitedNode: any): ReactIOInfo), // This is deduped by this reference.
       env: env,
     });
+    emitTimingChunk(request, task.id, awaitedNode.end);
   }
 }
 

packages/react-server/src/ReactFlightServerConfigDebugNode.js

@@ -49,8 +49,8 @@ export function initAsyncDebugInfo(): void {
               tag: AWAIT_NODE,
               owner: resolveOwner(),
               stack: new Error(),
-              start: -1.1,
-              end: -1.1,
+              start: performance.now(),
+              end: -1.1, // set when resolved.
               awaited: trigger, // The thing we're awaiting on. Might get overrriden when we resolve.
               previous: current === undefined ? null : current, // The path that led us here.
             }: AwaitNode);
@@ -118,10 +118,8 @@ export function initAsyncDebugInfo(): void {
               'A Promise should never be an IO_NODE. This is a bug in React.',
             );
           }
-          if (resolvedNode.tag === PROMISE_NODE) {
-            // Log the end time when we resolved the promise.
-            resolvedNode.end = performance.now();
-          }
+          // Log the end time when we resolved the promise.
+          resolvedNode.end = performance.now();
           const currentAsyncId = executionAsyncId();
           if (asyncId !== currentAsyncId) {
             // If the promise was not resolved by itself, then that means that

packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js

@@ -170,6 +170,9 @@ describe('ReactFlightAsyncDebugInfo', () => {
               ],
             ],
           },
+          {
+            "time": 0,
+          },
           {
             "awaited": {
               "end": 0,
@@ -257,6 +260,12 @@ describe('ReactFlightAsyncDebugInfo', () => {
               ],
             ],
           },
+          {
+            "time": 0,
+          },
+          {
+            "time": 0,
+          },
           {
             "awaited": {
               "end": 0,
@@ -347,6 +356,9 @@ describe('ReactFlightAsyncDebugInfo', () => {
           {
             "time": 0,
           },
+          {
+            "time": 0,
+          },
         ]
       `);
     }
@@ -398,13 +410,16 @@ describe('ReactFlightAsyncDebugInfo', () => {
               [
                 "Object.<anonymous>",
                 "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
-                368,
+                380,
                 109,
-                355,
+                367,
                 67,
               ],
             ],
           },
+          {
+            "time": 0,
+          },
           {
             "awaited": {
               "end": 0,
@@ -420,9 +435,9 @@ describe('ReactFlightAsyncDebugInfo', () => {
                   [
                     "Object.<anonymous>",
                     "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
-                    368,
+                    380,
                     109,
-                    355,
+                    367,
                     67,
                   ],
                 ],
@@ -431,9 +446,9 @@ describe('ReactFlightAsyncDebugInfo', () => {
                 [
                   "Component",
                   "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
-                  358,
+                  370,
                   7,
-                  356,
+                  368,
                   5,
                 ],
               ],
@@ -444,6 +459,9 @@ describe('ReactFlightAsyncDebugInfo', () => {
           {
             "time": 0,
           },
+          {
+            "time": 0,
+          },
         ]
       `);
     }