[Fizz] Push a stalled use() to the ownerStack/debugTask (facebook#35226)

Author: eps1lon

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

@@ -108,6 +108,28 @@ describe('ReactFlightDOMNode', () => {
     );
   }
 
+  /**
+   * Removes all stackframes not pointing into this file
+   */
+  function ignoreListStack(str) {
+    if (!str) {
+      return str;
+    }
+
+    let ignoreListedStack = '';
+    const lines = str.split('\n');
+
+    // eslint-disable-next-line no-for-of-loops/no-for-of-loops
+    for (const line of lines) {
+      if (line.indexOf(__filename) === -1) {
+      } else {
+        ignoreListedStack += '\n' + line.replace(__dirname, '.');
+      }
+    }
+
+    return ignoreListedStack;
+  }
+
   function readResult(stream) {
     return new Promise((resolve, reject) => {
       let buffer = '';
@@ -784,6 +806,165 @@ describe('ReactFlightDOMNode', () => {
     }
   });
 
+  // @gate enableHalt
+  it('includes source locations in component and owner stacks for halted Client components', async () => {
+    function SharedComponent({p1, p2, p3}) {
+      use(p1);
+      use(p2);
+      use(p3);
+      return <div>Hello, Dave!</div>;
+    }
+    const ClientComponentOnTheServer = clientExports(SharedComponent);
+    const ClientComponentOnTheClient = clientExports(
+      SharedComponent,
+      123,
+      'path/to/chunk.js',
+    );
+
+    let resolvePendingPromise;
+    function ServerComponent() {
+      const p1 = Promise.resolve();
+      const p2 = new Promise(resolve => {
+        resolvePendingPromise = value => {
+          p2.status = 'fulfilled';
+          p2.value = value;
+          resolve(value);
+        };
+      });
+      const p3 = new Promise(() => {});
+      return ReactServer.createElement(ClientComponentOnTheClient, {
+        p1: p1,
+        p2: p2,
+        p3: p3,
+      });
+    }
+
+    function App() {
+      return ReactServer.createElement(
+        'html',
+        null,
+        ReactServer.createElement(
+          'body',
+          null,
+          ReactServer.createElement(
+            ReactServer.Suspense,
+            {fallback: 'Loading...'},
+            ReactServer.createElement(ServerComponent, null),
+          ),
+        ),
+      );
+    }
+
+    const errors = [];
+    const rscStream = await serverAct(() =>
+      ReactServerDOMServer.renderToPipeableStream(
+        ReactServer.createElement(App, null),
+        webpackMap,
+      ),
+    );
+
+    const readable = new Stream.PassThrough(streamOptions);
+    rscStream.pipe(readable);
+
+    function ClientRoot({response}) {
+      return use(response);
+    }
+
+    const serverConsumerManifest = {
+      moduleMap: {
+        [webpackMap[ClientComponentOnTheClient.$$id].id]: {
+          '*': webpackMap[ClientComponentOnTheServer.$$id],
+        },
+      },
+      moduleLoading: webpackModuleLoading,
+    };
+
+    expect(errors).toEqual([]);
+
+    function ClientRoot({response}) {
+      return use(response);
+    }
+
+    const response = ReactServerDOMClient.createFromNodeStream(
+      readable,
+      serverConsumerManifest,
+    );
+
+    let componentStack;
+    let ownerStack;
+
+    const clientAbortController = new AbortController();
+
+    const fizzPrerenderStreamResult = ReactDOMFizzStatic.prerender(
+      React.createElement(ClientRoot, {response}),
+      {
+        signal: clientAbortController.signal,
+        onError(error, errorInfo) {
+          componentStack = errorInfo.componentStack;
+          ownerStack = React.captureOwnerStack
+            ? React.captureOwnerStack()
+            : null;
+        },
+      },
+    );
+
+    resolvePendingPromise('custom-instrum-resolve');
+    await serverAct(
+      async () =>
+        new Promise(resolve => {
+          setImmediate(() => {
+            clientAbortController.abort();
+            resolve();
+          });
+        }),
+    );
+
+    const fizzPrerenderStream = await fizzPrerenderStreamResult;
+    const prerenderHTML = await readWebResult(fizzPrerenderStream.prelude);
+
+    expect(prerenderHTML).toContain('Loading...');
+
+    if (__DEV__) {
+      expect(normalizeCodeLocInfo(componentStack)).toBe(
+        '\n' +
+          '    in SharedComponent (at **)\n' +
+          '    in ServerComponent' +
+          (gate(flags => flags.enableAsyncDebugInfo) ? ' (at **)' : '') +
+          '\n' +
+          '    in Suspense\n' +
+          '    in body\n' +
+          '    in html\n' +
+          '    in App (at **)\n' +
+          '    in ClientRoot (at **)',
+      );
+    } else {
+      expect(normalizeCodeLocInfo(componentStack)).toBe(
+        '\n' +
+          '    in SharedComponent (at **)\n' +
+          '    in Suspense\n' +
+          '    in body\n' +
+          '    in html\n' +
+          '    in ClientRoot (at **)',
+      );
+    }
+
+    if (__DEV__) {
+      expect(ignoreListStack(ownerStack)).toBe(
+        // eslint-disable-next-line react-internal/safe-string-coercion
+        '' +
+          // The concrete location may change as this test is updated.
+          // Just make sure they still point at React.use(p2)
+          (gate(flags => flags.enableAsyncDebugInfo)
+            ? '\n    at SharedComponent (./ReactFlightDOMNode-test.js:813:7)'
+            : '') +
+          '\n    at ServerComponent (file://./ReactFlightDOMNode-test.js:835:26)' +
+          '\n    at App (file://./ReactFlightDOMNode-test.js:852:25)',
+      );
+    } else {
+      expect(ownerStack).toBeNull();
+    }
+  });
+
   // @gate enableHalt
   it('includes deeper location for aborted stacks', async () => {
     async function getData() {
@@ -1364,12 +1545,12 @@ describe('ReactFlightDOMNode', () => {
           '\n' +
             '    in Dynamic' +
             (gate(flags => flags.enableAsyncDebugInfo)
-              ? ' (file://ReactFlightDOMNode-test.js:1238:27)\n'
+              ? ' (file://ReactFlightDOMNode-test.js:1419:27)\n'
               : '\n') +
             '    in body\n' +
             '    in html\n' +
-            '    in App (file://ReactFlightDOMNode-test.js:1251:25)\n' +
-            '    in ClientRoot (ReactFlightDOMNode-test.js:1326:16)',
+            '    in App (file://ReactFlightDOMNode-test.js:1432:25)\n' +
+            '    in ClientRoot (ReactFlightDOMNode-test.js:1507:16)',
         );
       } else {
         expect(
@@ -1378,7 +1559,7 @@ describe('ReactFlightDOMNode', () => {
           '\n' +
             '    in body\n' +
             '    in html\n' +
-            '    in ClientRoot (ReactFlightDOMNode-test.js:1326:16)',
+            '    in ClientRoot (ReactFlightDOMNode-test.js:1507:16)',
         );
       }
 
@@ -1388,16 +1569,16 @@ describe('ReactFlightDOMNode', () => {
             normalizeCodeLocInfo(ownerStack, {preserveLocation: true}),
           ).toBe(
             '\n' +
-              '    in Dynamic (file://ReactFlightDOMNode-test.js:1238:27)\n' +
-              '    in App (file://ReactFlightDOMNode-test.js:1251:25)',
+              '    in Dynamic (file://ReactFlightDOMNode-test.js:1419:27)\n' +
+              '    in App (file://ReactFlightDOMNode-test.js:1432:25)',
           );
         } else {
           expect(
             normalizeCodeLocInfo(ownerStack, {preserveLocation: true}),
           ).toBe(
             '' +
               '\n' +
-              '    in App (file://ReactFlightDOMNode-test.js:1251:25)',
+              '    in App (file://ReactFlightDOMNode-test.js:1432:25)',
           );
         }
       } else {

packages/react-server/src/ReactFizzServer.js

@@ -190,7 +190,14 @@ import assign from 'shared/assign';
 import noop from 'shared/noop';
 import getComponentNameFromType from 'shared/getComponentNameFromType';
 import isArray from 'shared/isArray';
-import {SuspenseException, getSuspendedThenable} from './ReactFizzThenable';
+import {
+  SuspenseException,
+  getSuspendedThenable,
+  ensureSuspendableThenableStateDEV,
+  getSuspendedCallSiteStackDEV,
+  getSuspendedCallSiteDebugTaskDEV,
+  setCaptureSuspendedCallSiteDEV,
+} from './ReactFizzThenable';
 
 // Linked list representing the identity of a component given the component/tag name and key.
 // The name might be minified but we assume that it's going to be the same generated name. Typically
@@ -355,6 +362,7 @@ const OPEN = 11;
 const ABORTING = 12;
 const CLOSING = 13;
 const CLOSED = 14;
+const STALLED_DEV = 15;
 
 export opaque type Request = {
   destination: null | Destination,
@@ -363,7 +371,7 @@ export opaque type Request = {
   +renderState: RenderState,
   +rootFormatContext: FormatContext,
   +progressiveChunkSize: number,
-  status: 10 | 11 | 12 | 13 | 14,
+  status: 10 | 11 | 12 | 13 | 14 | 15,
   fatalError: mixed,
   nextSegmentId: number,
   allPendingTasks: number, // when it reaches zero, we can close the connection.
@@ -1023,6 +1031,89 @@ function pushHaltedAwaitOnComponentStack(
   }
 }
 
+// performWork + retryTask without mutation
+function rerenderStalledTask(request: Request, task: Task): void {
+  const prevStatus = request.status;
+  request.status = STALLED_DEV;
+
+  const prevContext = getActiveContext();
+  const prevDispatcher = ReactSharedInternals.H;
+  ReactSharedInternals.H = HooksDispatcher;
+  const prevAsyncDispatcher = ReactSharedInternals.A;
+  ReactSharedInternals.A = DefaultAsyncDispatcher;
+
+  const prevRequest = currentRequest;
+  currentRequest = request;
+
+  const prevGetCurrentStackImpl = ReactSharedInternals.getCurrentStack;
+  ReactSharedInternals.getCurrentStack = getCurrentStackInDEV;
+
+  const prevResumableState = currentResumableState;
+  setCurrentResumableState(request.resumableState);
+  switchContext(task.context);
+  const prevTaskInDEV = currentTaskInDEV;
+  setCurrentTaskInDEV(task);
+  try {
+    retryNode(request, task);
+  } catch (x) {
+    // Suspended again.
+    resetHooksState();
+  } finally {
+    setCurrentTaskInDEV(prevTaskInDEV);
+    setCurrentResumableState(prevResumableState);
+
+    ReactSharedInternals.H = prevDispatcher;
+    ReactSharedInternals.A = prevAsyncDispatcher;
+
+    ReactSharedInternals.getCurrentStack = prevGetCurrentStackImpl;
+    if (prevDispatcher === HooksDispatcher) {
+      // This means that we were in a reentrant work loop. This could happen
+      // in a renderer that supports synchronous work like renderToString,
+      // when it's called from within another renderer.
+      // Normally we don't bother switching the contexts to their root/default
+      // values when leaving because we'll likely need the same or similar
+      // context again. However, when we're inside a synchronous loop like this
+      // we'll to restore the context to what it was before returning.
+      switchContext(prevContext);
+    }
+    currentRequest = prevRequest;
+    request.status = prevStatus;
+  }
+}
+
+function pushSuspendedCallSiteOnComponentStack(
+  request: Request,
+  task: Task,
+): void {
+  setCaptureSuspendedCallSiteDEV(true);
+  const restoreThenableState = ensureSuspendableThenableStateDEV(
+    // refined at the callsite
+    ((task.thenableState: any): ThenableState),
+  );
+  try {
+    rerenderStalledTask(request, task);
+  } finally {
+    restoreThenableState();
+    setCaptureSuspendedCallSiteDEV(false);
+  }
+
+  const suspendCallSiteStack = getSuspendedCallSiteStackDEV();
+  const suspendCallSiteDebugTask = getSuspendedCallSiteDebugTaskDEV();
+
+  if (suspendCallSiteStack !== null) {
+    const ownerStack = task.componentStack;
+    task.componentStack = {
+      // The owner of the suspended call site would be the owner of this task.
+      // We need the task itself otherwise we'd miss a frame.
+      owner: ownerStack,
+      parent: suspendCallSiteStack.parent,
+      stack: suspendCallSiteStack.stack,
+      type: suspendCallSiteStack.type,
+    };
+  }
+  task.debugTask = suspendCallSiteDebugTask;
+}
+
 function pushServerComponentStack(
   task: Task,
   debugInfo: void | null | ReactDebugInfo,
@@ -2723,7 +2814,12 @@ function renderLazyComponent(
     const init = lazyComponent._init;
     Component = init(payload);
   }
-  if (request.status === ABORTING) {
+  if (
+    request.status === ABORTING &&
+    // We're going to discard this render anyway.
+    // We just need to reach the point where we suspended in dev.
+    (!__DEV__ || request.status !== STALLED_DEV)
+  ) {
     // eslint-disable-next-line no-throw-literal
     throw null;
   }
@@ -4535,12 +4631,9 @@ function abortTask(task: Task, request: Request, error: mixed): void {
         debugInfo = node._debugInfo;
       }
       pushHaltedAwaitOnComponentStack(task, debugInfo);
-      /*
       if (task.thenableState !== null) {
-        // TODO: If we were stalled inside use() of a Client Component then we should
-        // rerender to get the stack trace from the use() call.
+        pushSuspendedCallSiteOnComponentStack(request, task);
       }
-      */
     }
   }