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

fixtures/fizz/server/render-to-stream.js

@@ -45,6 +45,11 @@ module.exports = function render(url, res) {
     onError(x) {
       didError = true;
       console.error(x);
+      // Redundant with `console.createTask`. Only added for debugging.
+      console.error(
+        'The above error occurred during server rendering: %s',
+        React.captureOwnerStack()
+      );
     },
   });
   // Abandon and switch to client rendering if enough time passes.

fixtures/fizz/src/App.js

@@ -6,10 +6,17 @@
  *
  */
 
+import {Suspense} from 'react';
 import Html from './Html';
 import BigComponent from './BigComponent';
+import MaybeHaltedComponent from './MaybeHaltedComponent';
 
-export default function App({assets, title}) {
+const serverHalt =
+  typeof window === 'undefined'
+    ? new Promise(() => {})
+    : Promise.resolve('client');
+
+export default function App({assets, promise, title}) {
   const components = [];
 
   for (let i = 0; i <= 250; i++) {
@@ -21,6 +28,10 @@ export default function App({assets, title}) {
       <h1>{title}</h1>
       {components}
       <h1>all done</h1>
+      <h2>or maybe not</h2>
+      <Suspense fallback="loading more...">
+        <MaybeHaltedComponent promise={serverHalt} />
+      </Suspense>
     </Html>
   );
 }

fixtures/fizz/src/MaybeHaltedComponent.js

@@ -0,0 +1,6 @@
+import {use} from 'react';
+
+export default function MaybeHaltedComponent({promise}) {
+  use(promise);
+  return <div>Did not halt</div>;
+}

packages/react-server/src/ReactFizzServer.js

@@ -190,7 +190,13 @@ 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,
+  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
@@ -1023,6 +1029,80 @@ function pushHaltedAwaitOnComponentStack(
   }
 }
 
+// performWork + retryTask without mutation
+function rerenderHaltedTask(request: Request, task: Task): void {
+  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;
+  }
+}
+
+function pushSuspendedCallSiteOnComponentStack(
+  request: Request,
+  task: Task,
+): void {
+  setCaptureSuspendedCallSiteDEV(true);
+  try {
+    rerenderHaltedTask(request, task);
+  } finally {
+    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,
@@ -4535,12 +4615,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

@@ -19,8 +19,10 @@ import type {
   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 +128,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 +168,83 @@ 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 {
+  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.',
+    );
+  }
+}

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

@@ -9,6 +9,7 @@
  */
 
 'use strict';
+import {AsyncLocalStorage} from 'node:async_hooks';
 
 let act;
 let React;
@@ -17,20 +18,43 @@ let ReactNoopServer;
 function normalizeCodeLocInfo(str) {
   return (
     str &&
-    str.replace(/^ +(?:at|in) ([\S]+)[^\n]*/gm, function (m, name) {
-      const dot = name.lastIndexOf('.');
-      if (dot !== -1) {
-        name = name.slice(dot + 1);
-      }
-      return '    in ' + name + (/\d/.test(m) ? ' (at **)' : '');
-    })
+    str
+      .split('\n')
+      .filter(frame => {
+        // These frames should be ignore-listed since they point into
+        // React internals i.e. node_modules.
+        return (
+          frame.indexOf('ReactFizzHooks') === -1 &&
+          frame.indexOf('ReactFizzThenable') === -1 &&
+          frame.indexOf('ReactHooks') === -1
+        );
+      })
+      .join('\n')
+      .replace(/^ +(?:at|in) ([\S]+)[^\n]*/gm, function (m, name) {
+        const dot = name.lastIndexOf('.');
+        if (dot !== -1) {
+          name = name.slice(dot + 1);
+        }
+        return '    in ' + name + (/\d/.test(m) ? ' (at **)' : '');
+      })
   );
 }
 
+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,24 +73,47 @@ describe('ReactServer', () => {
   });
 
   it('has Owner Stacks in DEV when aborted', async () => {
+    const Context = React.createContext(null);
+
     function Component({promise}) {
+      const context = React.use(Context);
+      if (context === null) {
+        throw new Error('Missing context');
+      }
       React.use(promise);
       return <div>Hello, Dave!</div>;
     }
+    function Indirection({promise}) {
+      return (
+        <div>
+          <Component promise={promise} />
+        </div>
+      );
+    }
     function App({promise}) {
-      return <Component promise={promise} />;
+      return (
+        <section>
+          <div>
+            <Indirection promise={promise} />
+          </div>
+        </section>
+      );
     }
 
     let caughtError;
     let componentStack;
     let ownerStack;
+    let task;
     const result = ReactNoopServer.render(
-      <App promise={new Promise(() => {})} />,
+      <Context value="provided">
+        <App promise={new Promise(() => {})} />
+      </Context>,
       {
         onError: (error, errorInfo) => {
           caughtError = error;
           componentStack = errorInfo.componentStack;
           ownerStack = __DEV__ ? React.captureOwnerStack() : null;
+          task = currentTask.getStore();
         },
       },
     );
@@ -80,10 +127,21 @@ describe('ReactServer', () => {
       }),
     );
     expect(normalizeCodeLocInfo(componentStack)).toEqual(
-      '\n    in Component (at **)' + '\n    in App (at **)',
+      '\n    in Component (at **)' +
+        '\n    in div' +
+        '\n    in Indirection (at **)' +
+        '\n    in div' +
+        '\n    in section' +
+        '\n    in App (at **)',
     );
     expect(normalizeCodeLocInfo(ownerStack)).toEqual(
-      __DEV__ ? '\n    in App (at **)' : null,
+      __DEV__
+        ? '' +
+            '\n    in Component (at **)' +
+            '\n    in Indirection (at **)' +
+            '\n    in App (at **)'
+        : null,
     );
+    expect(task).toEqual(__DEV__ ? '\n<Component>' : null);
   });
 });