Repository: facebook/react
Pull request: [Flight] Include env in ReactAsyncInfo and ReactIOInfo

Stacked on #33395.

This lets us keep track of which environment this was fetched and awaited.

Currently the IO and await is in the same environment. It's just kept when forwarded. Once we support forwarding information from a Promise fetched from another environment and awaited in this environment then the await can end up being in a different environment.

There's a question of when the await is inside Flight itself such as when you return a promise fetched from another environment whether that should mean that the await is in the current environment. I don't think so since the original stack trace is the best stack trace. It's only if you `await` it in user space in this environment first that this might happen and even then it should only be considered if there wasn't a better await earlier or if reading from the other environment was itself I/O.

The timing of *when* we read `environmentName()` is a little interesting here too.

Changed files:
filename: packages/react-client/src/ReactFlightClient.js
patch: @@ -2758,9 +2758,7 @@ function resolveConsoleEntry(
 
 function initializeIOInfo(response: Response, ioInfo: ReactIOInfo): void {
   const env =
-    // TODO: Pass env through I/O info.
-    // ioInfo.env !== undefined ? ioInfo.env :
-    response._rootEnvironmentName;
+    ioInfo.env === undefined ? response._rootEnvironmentName : ioInfo.env;
   if (ioInfo.stack !== undefined) {
     initializeFakeTask(response, ioInfo, env);
     initializeFakeStack(response, ioInfo);
@@ -2771,7 +2769,7 @@ function initializeIOInfo(response: Response, ioInfo: ReactIOInfo): void {
   // $FlowFixMe[cannot-write]
   ioInfo.end += response._timeOrigin;
 
-  logIOInfo(ioInfo);
+  logIOInfo(ioInfo, response._rootEnvironmentName);
 }
 
 function resolveIOInfo(

filename: packages/react-client/src/ReactFlightPerformanceTrack.js
patch: @@ -224,19 +224,23 @@ function getIOColor(
   }
 }
 
-export function logIOInfo(ioInfo: ReactIOInfo): void {
+export function logIOInfo(ioInfo: ReactIOInfo, rootEnv: string): void {
   const startTime = ioInfo.start;
   const endTime = ioInfo.end;
   if (supportsUserTiming && endTime >= 0) {
     const name = ioInfo.name;
+    const env = ioInfo.env;
+    const isPrimaryEnv = env === rootEnv;
+    const entryName =
+      isPrimaryEnv || env === undefined ? name : name + ' [' + env + ']';
     const debugTask = ioInfo.debugTask;
     const color = getIOColor(name);
     if (__DEV__ && debugTask) {
       debugTask.run(
         // $FlowFixMe[method-unbinding]
         console.timeStamp.bind(
           console,
-          name,
+          entryName,
           startTime < 0 ? 0 : startTime,
           endTime,
           IO_TRACK,
@@ -246,7 +250,7 @@ export function logIOInfo(ioInfo: ReactIOInfo): void {
       );
     } else {
       console.timeStamp(
-        name,
+        entryName,
         startTime < 0 ? 0 : startTime,
         endTime,
         IO_TRACK,

filename: packages/react-server/src/ReactFlightServer.js
patch: @@ -1930,10 +1930,14 @@ function visitAsyncNode(
           }
           // 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)();
           // 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,
           });
@@ -1969,8 +1973,12 @@ function emitAsyncSequence(
     }
     serializeIONode(request, awaitedNode);
     request.pendingChunks++;
+    // 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)();
     emitDebugChunk(request, task.id, {
       awaited: ((awaitedNode: any): ReactIOInfo), // This is deduped by this reference.
+      env: env,
     });
   }
 }
@@ -3524,6 +3532,7 @@ function emitIOInfoChunk(
   name: string,
   start: number,
   end: number,
+  env: ?string,
   owner: ?ReactComponentInfo,
   stack: ?ReactStackTrace,
 ): void {
@@ -3563,6 +3572,10 @@ function emitIOInfoChunk(
     start: relativeStartTimestamp,
     end: relativeEndTimestamp,
   };
+  if (env != null) {
+    // $FlowFixMe[cannot-write]
+    debugIOInfo.env = env;
+  }
   if (stack != null) {
     // $FlowFixMe[cannot-write]
     debugIOInfo.stack = stack;
@@ -3597,6 +3610,7 @@ function outlineIOInfo(request: Request, ioInfo: ReactIOInfo): void {
     ioInfo.name,
     ioInfo.start,
     ioInfo.end,
+    ioInfo.env,
     owner,
     ioInfo.stack,
   );
@@ -3633,9 +3647,22 @@ function serializeIONode(
     outlineComponentInfo(request, owner);
   }
 
+  // We log the environment at the time when we serialize the I/O node.
+  // The environment name may have changed from when the I/O was actually started.
+  const env = (0, request.environmentName)();
+
   request.pendingChunks++;
   const id = request.nextChunkId++;
-  emitIOInfoChunk(request, id, name, ioNode.start, ioNode.end, owner, stack);
+  emitIOInfoChunk(
+    request,
+    id,
+    name,
+    ioNode.start,
+    ioNode.end,
+    env,
+    owner,
+    stack,
+  );
   const ref = serializeByValueID(id);
   request.writtenObjects.set(ioNode, ref);
   return ref;
@@ -4161,6 +4188,7 @@ function forwardDebugInfo(
         const debugAsyncInfo: Omit<ReactAsyncInfo, 'debugTask' | 'debugStack'> =
           {
             awaited: ioInfo,
+            env: debugInfo[i].env,
             owner: debugInfo[i].owner,
             stack: debugInfo[i].stack,
           };

filename: packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js
patch: @@ -173,6 +173,7 @@ describe('ReactFlightAsyncDebugInfo', () => {
           {
             "awaited": {
               "end": 0,
+              "env": "Server",
               "name": "delay",
               "owner": {
                 "env": "Server",
@@ -219,6 +220,7 @@ describe('ReactFlightAsyncDebugInfo', () => {
               ],
               "start": 0,
             },
+            "env": "Server",
             "owner": {
               "env": "Server",
               "key": null,
@@ -258,6 +260,7 @@ describe('ReactFlightAsyncDebugInfo', () => {
           {
             "awaited": {
               "end": 0,
+              "env": "Server",
               "name": "delay",
               "owner": {
                 "env": "Server",
@@ -304,6 +307,7 @@ describe('ReactFlightAsyncDebugInfo', () => {
               ],
               "start": 0,
             },
+            "env": "Server",
             "owner": {
               "env": "Server",
               "key": null,
@@ -394,16 +398,17 @@ describe('ReactFlightAsyncDebugInfo', () => {
               [
                 "Object.<anonymous>",
                 "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
-                364,
+                368,
                 109,
-                351,
+                355,
                 67,
               ],
             ],
           },
           {
             "awaited": {
               "end": 0,
+              "env": "Server",
               "name": "setTimeout",
               "owner": {
                 "env": "Server",
@@ -415,9 +420,9 @@ describe('ReactFlightAsyncDebugInfo', () => {
                   [
                     "Object.<anonymous>",
                     "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
-                    364,
+                    368,
                     109,
-                    351,
+                    355,
                     67,
                   ],
                 ],
@@ -426,14 +431,15 @@ describe('ReactFlightAsyncDebugInfo', () => {
                 [
                   "Component",
                   "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
-                  354,
+                  358,
                   7,
-                  352,
+                  356,
                   5,
                 ],
               ],
               "start": 0,
             },
+            "env": "Server",
           },
           {
             "time": 0,

filename: packages/shared/ReactTypes.js
patch: @@ -234,6 +234,7 @@ export type ReactIOInfo = {
   +name: string, // the name of the async function being called (e.g. "fetch")
   +start: number, // the start time
   +end: number, // the end time (this might be different from the time the await was unblocked)
+  +env?: string, // the environment where this I/O was spawned.
   +owner?: null | ReactComponentInfo,
   +stack?: null | ReactStackTrace,
   // Stashed Data for the Specific Execution Environment. Not part of the transport protocol
@@ -243,6 +244,7 @@ export type ReactIOInfo = {
 
 export type ReactAsyncInfo = {
   +awaited: ReactIOInfo,
+  +env?: string, // the environment where this was awaited. This might not be the same as where it was spawned.
   +owner?: null | ReactComponentInfo,
   +stack?: null | ReactStackTrace,
   // Stashed Data for the Specific Execution Environment. Not part of the transport protocol