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

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,8 +1569,8 @@ 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(
@@ -1397,7 +1578,7 @@ describe('ReactFlightDOMNode', () => {
           ).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);
       }
-      */
     }
   }
 

packages/react-server/src/ReactFizzThenable.js

@@ -7,20 +7,19 @@
  * @flow
  */
 
-// Corresponds to ReactFiberWakeable and ReactFlightWakeable modules. Generally,
+// Corresponds to ReactFiberThenable and ReactFlightThenable modules. Generally,
 // changes to one module should be reflected in the others.
 
-// TODO: Rename this module and the corresponding Fiber one to "Thenable"
-// instead of "Wakeable". Or some other more appropriate name.
-
 import type {
   Thenable,
   PendingThenable,
   FulfilledThenable,
   RejectedThenable,
 } from 'shared/ReactTypes';
+import type {ComponentStackNode} from './ReactFizzComponentStack';
 
 import noop from 'shared/noop';
+import {currentTaskInDEV} from './ReactFizzCurrentTask';
 
 export opaque type ThenableState = Array<Thenable<any>>;
 
@@ -126,6 +125,9 @@ export function trackUsedThenable<T>(
       // get captured by the work loop, log a warning, because that means
       // something in userspace must have caught it.
       suspendedThenable = thenable;
+      if (__DEV__ && shouldCaptureSuspendedCallSite) {
+        captureSuspendedCallSite();
+      }
       throw SuspenseException;
     }
   }
@@ -163,3 +165,130 @@ export function getSuspendedThenable(): Thenable<mixed> {
   suspendedThenable = null;
   return thenable;
 }
+
+let shouldCaptureSuspendedCallSite: boolean = false;
+export function setCaptureSuspendedCallSiteDEV(capture: boolean): void {
+  if (!__DEV__) {
+    // eslint-disable-next-line react-internal/prod-error-codes
+    throw new Error(
+      'setCaptureSuspendedCallSiteDEV was called in a production environment. ' +
+        'This is a bug in React.',
+    );
+  }
+  shouldCaptureSuspendedCallSite = capture;
+}
+
+// DEV-only
+let suspendedCallSiteStack: ComponentStackNode | null = null;
+let suspendedCallSiteDebugTask: ConsoleTask | null = null;
+function captureSuspendedCallSite(): void {
+  // This is currently only used when aborting in Fizz.
+  // You can only abort the render in Fizz and Flight.
+  // In Fiber we only track suspended use via DevTools.
+  // In Flight, we track suspended use via async debug info.
+  const currentTask = currentTaskInDEV;
+  if (currentTask === null) {
+    // eslint-disable-next-line react-internal/prod-error-codes -- not a prod error
+    throw new Error(
+      'Expected to have a current task when tracking a suspend call site. ' +
+        'This is a bug in React.',
+    );
+  }
+  const currentComponentStack = currentTask.componentStack;
+  if (currentComponentStack === null) {
+    // eslint-disable-next-line react-internal/prod-error-codes -- not a prod error
+    throw new Error(
+      'Expected to have a component stack on the current task when ' +
+        'tracking a suspended call site. This is a bug in React.',
+    );
+  }
+  suspendedCallSiteStack = {
+    parent: currentComponentStack.parent,
+    type: currentComponentStack.type,
+    owner: currentComponentStack.owner,
+    stack: Error('react-stack-top-frame'),
+  };
+  // TODO: If this is used in error handlers, the ConsoleTask stack
+  // will just be this debugTask + the stack of the abort() call which usually means
+  // it's just this debugTask.
+  // Ideally we'd be able to reconstruct the owner ConsoleTask as well.
+  // The stack of the debugTask would not point to the suspend location anyway.
+  // The focus is really on callsite which should be used in captureOwnerStack().
+  suspendedCallSiteDebugTask = currentTask.debugTask;
+}
+export function getSuspendedCallSiteStackDEV(): ComponentStackNode | null {
+  if (__DEV__) {
+    if (suspendedCallSiteStack === null) {
+      return null;
+    }
+    const callSite = suspendedCallSiteStack;
+    suspendedCallSiteStack = null;
+    return callSite;
+  } else {
+    // eslint-disable-next-line react-internal/prod-error-codes
+    throw new Error(
+      'getSuspendedCallSiteDEV was called in a production environment. ' +
+        'This is a bug in React.',
+    );
+  }
+}
+
+export function getSuspendedCallSiteDebugTaskDEV(): ConsoleTask | null {
+  if (__DEV__) {
+    if (suspendedCallSiteDebugTask === null) {
+      return null;
+    }
+    const debugTask = suspendedCallSiteDebugTask;
+    suspendedCallSiteDebugTask = null;
+    return debugTask;
+  } else {
+    // eslint-disable-next-line react-internal/prod-error-codes
+    throw new Error(
+      'getSuspendedCallSiteDebugTaskDEV was called in a production environment. ' +
+        'This is a bug in React.',
+    );
+  }
+}
+
+export function ensureSuspendableThenableStateDEV(
+  thenableState: ThenableState,
+): () => void {
+  if (__DEV__) {
+    const lastThenable = thenableState[thenableState.length - 1];
+    // Reset the last thenable back to pending.
+    switch (lastThenable.status) {
+      case 'fulfilled':
+        const previousThenableValue = lastThenable.value;
+        // $FlowIgnore[method-unbinding] We rebind .then immediately.
+        const previousThenableThen = lastThenable.then.bind(lastThenable);
+        delete lastThenable.value;
+        delete (lastThenable: any).status;
+        // We'll call .then again if we resuspend. Since we potentially corrupted
+        // the internal state of unknown classes, we need to diffuse the potential
+        // crash by replacing the .then method with a noop.
+        // $FlowFixMe[cannot-write] Custom userspace Thenables may not be but native Promises are.
+        lastThenable.then = noop;
+        return () => {
+          // $FlowFixMe[cannot-write] Custom userspace Thenables may not be but native Promises are.
+          lastThenable.then = previousThenableThen;
+          lastThenable.value = previousThenableValue;
+          lastThenable.status = 'fulfilled';
+        };
+      case 'rejected':
+        const previousThenableReason = lastThenable.reason;
+        delete lastThenable.reason;
+        delete (lastThenable: any).status;
+        return () => {
+          lastThenable.reason = previousThenableReason;
+          lastThenable.status = 'rejected';
+        };
+    }
+    return noop;
+  } else {
+    // eslint-disable-next-line react-internal/prod-error-codes
+    throw new Error(
+      'ensureSuspendableThenableStateDEV was called in a production environment. ' +
+        'This is a bug in React.',
+    );
+  }
+}

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

@@ -9,6 +9,7 @@
  */
 
 'use strict';
+import {AsyncLocalStorage} from 'node:async_hooks';
 
 let act;
 let React;
@@ -27,10 +28,43 @@ function normalizeCodeLocInfo(str) {
   );
 }
 
+/**
+ * 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;
+}
+
+const currentTask = new AsyncLocalStorage({defaultValue: null});
+
 describe('ReactServer', () => {
   beforeEach(() => {
     jest.resetModules();
 
+    console.createTask = jest.fn(taskName => {
+      return {
+        run: taskFn => {
+          const parentTask = currentTask.getStore() || '';
+          return currentTask.run(parentTask + '\n' + taskName, taskFn);
+        },
+      };
+    });
+
     act = require('internal-test-utils').act;
     React = require('react');
     ReactNoopServer = require('react-noop-renderer/server');
@@ -49,29 +83,67 @@ describe('ReactServer', () => {
   });
 
   it('has Owner Stacks in DEV when aborted', async () => {
-    function Component({promise}) {
-      React.use(promise);
+    const Context = React.createContext(null);
+
+    function Component({p1, p2, p3}) {
+      const context = React.use(Context);
+      if (context === null) {
+        throw new Error('Missing context');
+      }
+      React.use(p1);
+      React.use(p2);
+      React.use(p3);
       return <div>Hello, Dave!</div>;
     }
-    function App({promise}) {
-      return <Component promise={promise} />;
+    function Indirection({p1, p2, p3}) {
+      return (
+        <div>
+          <Component p1={p1} p2={p2} p3={p3} />
+        </div>
+      );
+    }
+    function App({p1, p2, p3}) {
+      return (
+        <section>
+          <div>
+            <Indirection p1={p1} p2={p2} p3={p3} />
+          </div>
+        </section>
+      );
     }
 
     let caughtError;
     let componentStack;
     let ownerStack;
+    let task;
+    const resolvedPromise = Promise.resolve('one');
+    resolvedPromise.status = 'fulfilled';
+    resolvedPromise.value = 'one';
+    let resolvePendingPromise;
+    const pendingPromise = new Promise(resolve => {
+      resolvePendingPromise = value => {
+        pendingPromise.status = 'fulfilled';
+        pendingPromise.value = value;
+        resolve(value);
+      };
+    });
+    const hangingPromise = new Promise(() => {});
     const result = ReactNoopServer.render(
-      <App promise={new Promise(() => {})} />,
+      <Context value="provided">
+        <App p1={resolvedPromise} p2={pendingPromise} p3={hangingPromise} />
+      </Context>,
       {
         onError: (error, errorInfo) => {
           caughtError = error;
           componentStack = errorInfo.componentStack;
           ownerStack = __DEV__ ? React.captureOwnerStack() : null;
+          task = currentTask.getStore();
         },
       },
     );
 
     await act(async () => {
+      resolvePendingPromise('two');
       result.abort();
     });
     expect(caughtError).toEqual(
@@ -80,10 +152,35 @@ describe('ReactServer', () => {
       }),
     );
     expect(normalizeCodeLocInfo(componentStack)).toEqual(
-      '\n    in Component (at **)' + '\n    in App (at **)',
-    );
-    expect(normalizeCodeLocInfo(ownerStack)).toEqual(
-      __DEV__ ? '\n    in App (at **)' : null,
+      '\n    in Component (at **)' +
+        '\n    in div' +
+        '\n    in Indirection (at **)' +
+        '\n    in div' +
+        '\n    in section' +
+        '\n    in App (at **)',
     );
+    if (__DEV__) {
+      // The concrete location may change as this test is updated.
+      // Just make sure they still point at the same code
+      if (gate(flags => flags.enableAsyncDebugInfo)) {
+        expect(ignoreListStack(ownerStack)).toEqual(
+          '' +
+            // Pointing at React.use(p2)
+            '\n    at Component (./ReactServer-test.js:94:13)' +
+            '\n    at Indirection (./ReactServer-test.js:101:44)' +
+            '\n    at App (./ReactServer-test.js:109:46)',
+        );
+      } else {
+        expect(ignoreListStack(ownerStack)).toEqual(
+          '' +
+            '\n    at Indirection (./ReactServer-test.js:101:44)' +
+            '\n    at App (./ReactServer-test.js:109:46)',
+        );
+      }
+      expect(task).toEqual('\n<Component>');
+    } else {
+      expect(ownerStack).toBeNull();
+      expect(task).toEqual(undefined);
+    }
   });
 });