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

Author: eps1lon

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,37 @@ function pushHaltedAwaitOnComponentStack(
   }
 }
 
+function pushSuspendedCallSiteOnComponentStack(
+  request: Request,
+  task: Task,
+): void {
+  setCaptureSuspendedCallSiteDEV(true);
+  let suspendCallSiteStack: ComponentStackNode | null = null;
+  let suspendCallSiteDebugTask: ConsoleTask | null = null;
+  const previousPingedTasks = request.pingedTasks;
+  try {
+    // TODO: Use a dedicated method to re-render instead of abusing ping.
+    request.pingedTasks = [task];
+    performWork(request);
+    suspendCallSiteStack = getSuspendedCallSiteStackDEV();
+    suspendCallSiteDebugTask = getSuspendedCallSiteDebugTaskDEV();
+  } finally {
+    request.pingedTasks = previousPingedTasks;
+    setCaptureSuspendedCallSiteDEV(false);
+  }
+
+  if (suspendCallSiteStack !== null) {
+    const ownerStack = task.componentStack;
+    task.componentStack = {
+      owner: ownerStack,
+      parent: suspendCallSiteStack.parent,
+      stack: suspendCallSiteStack.stack,
+      type: suspendCallSiteStack.type,
+    };
+  }
+  task.debugTask = suspendCallSiteDebugTask;
+}
+
 function pushServerComponentStack(
   task: Task,
   debugInfo: void | null | ReactDebugInfo,
@@ -4535,12 +4572,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);
       }
-      */
     }
   }
 
@@ -4962,7 +4996,9 @@ function retryRenderTask(
   task: RenderTask,
   segment: Segment,
 ): void {
-  if (segment.status !== PENDING) {
+  // TODO: We only retry when aborted to get the suspended callsite.
+  // Use a dedicated mechanism to re-render.
+  if (segment.status !== PENDING && segment.status !== ABORTED) {
     // We completed this by other means before we had a chance to retry it.
     return;
   }

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,77 @@ 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'),
+  };
+  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

@@ -83,7 +83,15 @@ describe('ReactServer', () => {
       '\n    in Component (at **)' + '\n    in App (at **)',
     );
     expect(normalizeCodeLocInfo(ownerStack)).toEqual(
-      __DEV__ ? '\n    in App (at **)' : null,
+      __DEV__
+        ? '' +
+            '\n    in trackUsedThenable (at **)' +
+            '\n    in unwrapThenable (at **)' +
+            '\n    in use (at **)' +
+            '\n    in use (at **)' +
+            '\n    in Component (at **)' +
+            '\n    in App (at **)'
+        : null,
     );
   });
 });