Repository: facebook/react
Pull request: [Flight] Insert an extra await node for awaiting on the promise returned by then callback

When a `.then()` callback returns another Promise, there's effectively another "await" on that Promise that happens in the internals but that was not modeled. In effect the Promise returned by `.then()` is blocked on both the original Promise AND the promise returned by the callback.

This models that by cloning the original node and treat that as the await on the original Promise. Then we use the existing Node to await the new Promise but its "previous" points to the clone. That way we have a forked node that awaits both.



Changed files:
filename: packages/react-server/src/ReactFlightServer.js
patch: @@ -2163,13 +2163,21 @@ function visitAsyncNode(
           } else {
             let isAwaitInUserspace = false;
             const fullStack = node.stack;
-            if (fullStack.length > 0) {
+            let firstFrame = 0;
+            while (
+              fullStack.length > firstFrame &&
+              fullStack[firstFrame][0] === 'Promise.then'
+            ) {
+              // Skip Promise.then frame itself.
+              firstFrame++;
+            }
+            if (fullStack.length > firstFrame) {
               // Check if the very first stack frame that awaited this Promise was in user space.
               // TODO: This doesn't take into account wrapper functions such as our fake .then()
               // in FlightClient which will always be considered third party awaits if you call
               // .then directly.
               const filterStackFrame = request.filterStackFrame;
-              const callsite = fullStack[0];
+              const callsite = fullStack[firstFrame];
               const functionName = callsite[0];
               const url = devirtualizeURL(callsite[1]);
               isAwaitInUserspace = filterStackFrame(url, functionName);

filename: packages/react-server/src/ReactFlightServerConfigDebugNode.js
patch: @@ -257,7 +257,36 @@ export function initAsyncDebugInfo(): void {
             // the trigger that we originally stored wasn't actually the dependency.
             // Instead, the current execution context is what ultimately unblocked it.
             const awaited = pendingOperations.get(currentAsyncId);
-            resolvedNode.awaited = awaited === undefined ? null : awaited;
+            if (resolvedNode.tag === PROMISE_NODE) {
+              // For a Promise we just override the await. We're not interested in
+              // what created the Promise itself.
+              resolvedNode.awaited = awaited === undefined ? null : awaited;
+            } else {
+              // For an await, there's really two things awaited here. It's the trigger
+              // that .then() was called on but there seems to also be something else
+              // in the .then() callback that blocked the returned Promise from resolving
+              // immediately. We create a fork node which essentially represents an await
+              // of the Promise returned from the .then() callback. That Promise was blocked
+              // on the original awaited thing which we stored as "previous".
+              if (awaited !== undefined) {
+                const clonedNode: AwaitNode = {
+                  tag: AWAIT_NODE,
+                  owner: resolvedNode.owner,
+                  stack: resolvedNode.stack,
+                  start: resolvedNode.start,
+                  end: resolvedNode.end,
+                  promise: resolvedNode.promise,
+                  awaited: resolvedNode.awaited,
+                  previous: resolvedNode.previous,
+                };
+                // We started awaiting on the callback when the original .then() resolved.
+                resolvedNode.start = resolvedNode.end;
+                // It resolved now. We could use the end time of "awaited" maybe.
+                resolvedNode.end = performance.now();
+                resolvedNode.previous = clonedNode;
+                resolvedNode.awaited = awaited;
+              }
+            }
           }
         }
       },

filename: packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js
patch: @@ -2226,4 +2226,158 @@ describe('ReactFlightAsyncDebugInfo', () => {
       `);
     }
   });
+
+  it('can track IO that is chained via then(async ...)', async () => {
+    function getData(text) {
+      return delay(1).then(async () => {
+        return text.toUpperCase();
+      });
+    }
+
+    async function Component({text, promise}) {
+      return await getData('hi, sebbie');
+    }
+
+    const stream = ReactServerDOMServer.renderToPipeableStream(<Component />);
+
+    const readable = new Stream.PassThrough(streamOptions);
+
+    const result = ReactServerDOMClient.createFromNodeStream(readable, {
+      moduleMap: {},
+      moduleLoading: {},
+    });
+    stream.pipe(readable);
+
+    expect(await result).toBe('HI, SEBBIE');
+
+    await finishLoadingStream(readable);
+    if (
+      __DEV__ &&
+      gate(
+        flags =>
+          flags.enableComponentPerformanceTrack && flags.enableAsyncDebugInfo,
+      )
+    ) {
+      expect(getDebugInfo(result)).toMatchInlineSnapshot(`
+        [
+          {
+            "time": 0,
+          },
+          {
+            "env": "Server",
+            "key": null,
+            "name": "Component",
+            "props": {},
+            "stack": [
+              [
+                "Object.<anonymous>",
+                "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
+                2241,
+                109,
+                2230,
+                58,
+              ],
+            ],
+          },
+          {
+            "time": 0,
+          },
+          {
+            "awaited": {
+              "end": 0,
+              "env": "Server",
+              "name": "delay",
+              "owner": {
+                "env": "Server",
+                "key": null,
+                "name": "Component",
+                "props": {},
+                "stack": [
+                  [
+                    "Object.<anonymous>",
+                    "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
+                    2241,
+                    109,
+                    2230,
+                    58,
+                  ],
+                ],
+              },
+              "stack": [
+                [
+                  "delay",
+                  "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
+                  160,
+                  12,
+                  159,
+                  3,
+                ],
+                [
+                  "getData",
+                  "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
+                  2232,
+                  14,
+                  2231,
+                  5,
+                ],
+                [
+                  "Component",
+                  "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
+                  2238,
+                  20,
+                  2237,
+                  5,
+                ],
+              ],
+              "start": 0,
+              "value": {
+                "value": undefined,
+              },
+            },
+            "env": "Server",
+            "owner": {
+              "env": "Server",
+              "key": null,
+              "name": "Component",
+              "props": {},
+              "stack": [
+                [
+                  "Object.<anonymous>",
+                  "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
+                  2241,
+                  109,
+                  2230,
+                  58,
+                ],
+              ],
+            },
+            "stack": [
+              [
+                "getData",
+                "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
+                2232,
+                23,
+                2231,
+                5,
+              ],
+              [
+                "Component",
+                "/packages/react-server/src/__tests__/ReactFlightAsyncDebugInfo-test.js",
+                2238,
+                20,
+                2237,
+                5,
+              ],
+            ],
+          },
+          {
+            "time": 0,
+          },
+          {
+            "time": 0,
+          },
+        ]
+      `);
+    }
+  });
 });