Include env in ReactAsyncInfo and ReactIOInfo

This lets us keep track of which environment this was fetched and awaited.

Currently the IO and await is in the same environment.

Once we support forwarding information from a Promise fetched from another
environment, the await can end up being in a different environment.

Author: sebmarkbage

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);

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