[Fizz] Push halted await to the owner stack for late-arriving I/O info

To quote from facebook#33634:

> If an aborted task is not rendering, then this is an async abort.
> Conceptually it's as if the abort happened inside the async gap. The
> abort reason's stack frame won't have that on the stack so instead we
> use the owner stack and debug task of any halted async debug info.

This PR extends that logic to also try to resolve lazy components to
find debug info that has been transferred to the inner value. In
addition, we ignore any time and component info that might precede the
I/O info, effectively allowing resolved I/O to also be considered for
the owner stack. This is useful in a scenario where the Flight rendering
might have been completed (and not prematurely aborted), but then the
Fizz rendering is intentionally aborted before all chunks were received,
while still allowing the remaining chunks (including I/O info for halted
components) to be processed while the prerender is in the aborting
state.

Author: unstubbable

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMNode-test.js

@@ -86,12 +86,25 @@ describe('ReactFlightDOMNode', () => {
     );
   }
 
-  function normalizeCodeLocInfo(str) {
+  const relativeFilename = path.relative(__dirname, __filename);
+
+  function normalizeCodeLocInfo(str, {preserveLocation = false} = {}) {
     return (
       str &&
-      str.replace(/^ +(?:at|in) ([\S]+)[^\n]*/gm, function (m, name) {
-        return '    in ' + name + (/\d/.test(m) ? ' (at **)' : '');
-      })
+      str.replace(
+        /^ +(?:at|in) ([\S]+) ([^\n]*)/gm,
+        function (m, name, location) {
+          return (
+            '    in ' +
+            name +
+            (/\d/.test(m)
+              ? preserveLocation
+                ? ' ' + location.replace(__filename, relativeFilename)
+                : ' (at **)'
+              : '')
+          );
+        },
+      )
     );
   }
 
@@ -1169,4 +1182,198 @@ describe('ReactFlightDOMNode', () => {
     // Must not throw an error.
     await readable.pipeTo(writable);
   });
+
+  describe('with real timers', () => {
+    // These tests schedule their rendering in a way that requires real timers
+    // to be used to accurately represent how this interacts with React's
+    // internal scheduling.
+
+    beforeEach(() => {
+      jest.useRealTimers();
+    });
+
+    afterEach(() => {
+      jest.useFakeTimers();
+    });
+
+    it('should use late-arriving I/O debug info to enhance component and owner stacks when aborting a prerender', async () => {
+      let resolveDynamicData;
+
+      async function getCachedData() {
+        // Cached data resolves in microtasks.
+        return Promise.resolve('Hi');
+      }
+
+      async function getDynamicData() {
+        return new Promise(resolve => {
+          resolveDynamicData = resolve;
+        });
+      }
+
+      async function Dynamic() {
+        const cachedData = await getCachedData();
+        const dynamicData = await getDynamicData();
+
+        return (
+          <p>
+            {cachedData} {dynamicData}
+          </p>
+        );
+      }
+
+      function App() {
+        return ReactServer.createElement(
+          'html',
+          null,
+          ReactServer.createElement(
+            'body',
+            null,
+            ReactServer.createElement(Dynamic),
+          ),
+        );
+      }
+
+      const stream = await ReactServerDOMServer.renderToPipeableStream(
+        ReactServer.createElement(App),
+        webpackMap,
+        {filterStackFrame},
+      );
+
+      const staticChunks = [];
+      const dynamicChunks = [];
+      let isStatic = true;
+
+      const passThrough = new Stream.PassThrough(streamOptions);
+      stream.pipe(passThrough);
+
+      // Split chunks into static and dynamic chunks.
+      passThrough.on('data', chunk => {
+        if (isStatic) {
+          staticChunks.push(chunk);
+        } else {
+          dynamicChunks.push(chunk);
+        }
+      });
+
+      await new Promise(resolve => {
+        setTimeout(() => {
+          isStatic = false;
+          resolveDynamicData('Josh');
+          resolve();
+        });
+      });
+
+      await new Promise(resolve => {
+        passThrough.on('end', resolve);
+      });
+
+      // Create a new Readable and push all static chunks immediately.
+      const readable = new Stream.Readable({...streamOptions, read() {}});
+      for (let i = 0; i < staticChunks.length; i++) {
+        readable.push(staticChunks[i]);
+      }
+
+      const abortController = new AbortController();
+
+      // When prerendering is aborted, push all dynamic chunks.
+      abortController.signal.addEventListener(
+        'abort',
+        () => {
+          for (let i = 0; i < dynamicChunks.length; i++) {
+            readable.push(dynamicChunks[i]);
+          }
+        },
+        {once: true},
+      );
+
+      const response = ReactServerDOMClient.createFromNodeStream(readable, {
+        serverConsumerManifest: {
+          moduleMap: null,
+          moduleLoading: null,
+        },
+      });
+
+      function ClientRoot() {
+        return use(response);
+      }
+
+      let componentStack;
+      let ownerStack;
+
+      const {prelude} = await new Promise(resolve => {
+        let result;
+
+        setTimeout(() => {
+          result = ReactDOMFizzStatic.prerenderToNodeStream(
+            React.createElement(ClientRoot),
+            {
+              signal: abortController.signal,
+              onError(error, errorInfo) {
+                componentStack = errorInfo.componentStack;
+                ownerStack = React.captureOwnerStack
+                  ? React.captureOwnerStack()
+                  : null;
+              },
+            },
+          );
+        });
+
+        setTimeout(() => {
+          abortController.abort();
+          resolve(result);
+        });
+      });
+
+      const prerenderHTML = await readResult(prelude);
+
+      expect(prerenderHTML).toContain('');
+
+      if (__DEV__) {
+        expect(
+          normalizeCodeLocInfo(componentStack, {preserveLocation: true}),
+        ).toBe(
+          '\n' +
+            '    in Dynamic' +
+            (gate(flags => flags.enableAsyncDebugInfo)
+              ? ' (file://ReactFlightDOMNode-test.js:1215:33)\n'
+              : '\n') +
+            '    in body\n' +
+            '    in html\n' +
+            '    in App (file://ReactFlightDOMNode-test.js:1231:25)\n' +
+            '    in ClientRoot (ReactFlightDOMNode-test.js:1297:16)',
+        );
+      } else {
+        expect(
+          normalizeCodeLocInfo(componentStack, {preserveLocation: true}),
+        ).toBe(
+          '\n' +
+            '    in body\n' +
+            '    in html\n' +
+            '    in ClientRoot (ReactFlightDOMNode-test.js:1297:16)',
+        );
+      }
+
+      if (__DEV__) {
+        if (gate(flags => flags.enableAsyncDebugInfo)) {
+          expect(
+            normalizeCodeLocInfo(ownerStack, {preserveLocation: true}),
+          ).toBe(
+            '\n' +
+              '    in Dynamic (file://ReactFlightDOMNode-test.js:1215:33)\n' +
+              '    in App (file://ReactFlightDOMNode-test.js:1231:25)',
+          );
+        } else {
+          expect(
+            normalizeCodeLocInfo(ownerStack, {preserveLocation: true}),
+          ).toBe(
+            '' +
+              '\n' +
+              '    in App (file://ReactFlightDOMNode-test.js:1231:25)',
+          );
+        }
+      } else {
+        expect(ownerStack).toBeNull();
+      }
+    });
+  });
 });

packages/react-server/src/ReactFizzServer.js

@@ -1004,14 +1004,6 @@ function pushHaltedAwaitOnComponentStack(
   if (debugInfo != null) {
     for (let i = debugInfo.length - 1; i >= 0; i--) {
       const info = debugInfo[i];
-      if (typeof info.name === 'string') {
-        // This is a Server Component. Any awaits in previous Server Components already resolved.
-        break;
-      }
-      if (typeof info.time === 'number') {
-        // This had an end time. Any awaits before this must have already resolved.
-        break;
-      }
       if (info.awaited != null) {
         const asyncInfo: ReactAsyncInfo = (info: any);
         const bestStack =
@@ -4653,10 +4645,38 @@ function abortTask(task: Task, request: Request, error: mixed): void {
     // If the task is not rendering, then this is an async abort. Conceptually it's as if
     // the abort happened inside the async gap. The abort reason's stack frame won't have that
     // on the stack so instead we use the owner stack and debug task of any halted async debug info.
-    const node: any = task.node;
+    let node: any = task.node;
     if (node !== null && typeof node === 'object') {
       // Push a fake component stack frame that represents the await.
-      pushHaltedAwaitOnComponentStack(task, node._debugInfo);
+      let debugInfo = node._debugInfo;
+      if (debugInfo == null || debugInfo.length === 0) {
+        // If there's no debug info, try resolving lazy components to find debug
+        // info that has been transferred to the inner value.
+        while (
+          typeof node === 'object' &&
+          node !== null &&
+          node.$$typeof === REACT_LAZY_TYPE
+        ) {
+          const payload = node._payload;
+          if (payload.status === 'fulfilled') {
+            node = payload.value;
+            continue;
+          }
+          break;
+        }
+        if (
+          typeof node === 'object' &&
+          node !== null &&
+          (isArray(node) ||
+            typeof node[ASYNC_ITERATOR] === 'function' ||
+            node.$$typeof === REACT_ELEMENT_TYPE ||
+            node.$$typeof === REACT_LAZY_TYPE) &&
+          isArray(node._debugInfo)
+        ) {
+          debugInfo = node._debugInfo;
+        }
+      }
+      pushHaltedAwaitOnComponentStack(task, debugInfo);
       /*
       if (task.thenableState !== null) {
         // TODO: If we were stalled inside use() of a Client Component then we should