[Flight] Include env in ReactAsyncInfo and ReactIOInfo

packages/react-client/src/ReactFlightClient.js

@@ -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(

packages/react-client/src/ReactFlightPerformanceTrack.js

@@ -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,

packages/react-server/src/ReactFlightServer.js

@@ -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,
           };

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

@@ -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,

packages/shared/ReactTypes.js

@@ -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