Ensure we rerender with the original suspender

Author: eps1lon

packages/react-server/src/ReactFizzServer.js

@@ -193,6 +193,7 @@ import isArray from 'shared/isArray';
 import {
   SuspenseException,
   getSuspendedThenable,
+  ensureSuspendableThenableStateDEV,
   getSuspendedCallSiteStackDEV,
   getSuspendedCallSiteDebugTaskDEV,
   setCaptureSuspendedCallSiteDEV,
@@ -1030,7 +1031,7 @@ function pushHaltedAwaitOnComponentStack(
 }
 
 // performWork + retryTask without mutation
-function rerenderHaltedTask(request: Request, task: Task): void {
+function rerenderStalledTask(request: Request, task: Task): void {
   const prevContext = getActiveContext();
   const prevDispatcher = ReactSharedInternals.H;
   ReactSharedInternals.H = HooksDispatcher;
@@ -1080,9 +1081,14 @@ function pushSuspendedCallSiteOnComponentStack(
   task: Task,
 ): void {
   setCaptureSuspendedCallSiteDEV(true);
+  const restoreThenableState = ensureSuspendableThenableStateDEV(
+    // refined at the callsite
+    ((task.thenableState: any): ThenableState),
+  );
   try {
-    rerenderHaltedTask(request, task);
+    rerenderStalledTask(request, task);
   } finally {
+    restoreThenableState();
     setCaptureSuspendedCallSiteDEV(false);
   }
 
@@ -4616,13 +4622,6 @@ function abortTask(task: Task, request: Request, error: mixed): void {
       }
       pushHaltedAwaitOnComponentStack(task, debugInfo);
       if (task.thenableState !== null) {
-        // TODO: really?
-        // If the thenable was resolved in the meantime, we won't get a stack.
-        // We won't know which thenable in thenableState is newly settled though.
-        // We can't just clear status fields on each thenable because then the
-        // stack may point at a thenable that wasn't stalled. In those cases
-        // it's better to point at the callsite of the stalled Component as an
-        // entrypoint instead of the wrong thenable.
         pushSuspendedCallSiteOnComponentStack(request, task);
       }
     }

packages/react-server/src/ReactFizzThenable.js

@@ -7,12 +7,9 @@
  * @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,
@@ -185,6 +182,10 @@ export function setCaptureSuspendedCallSiteDEV(capture: boolean): void {
 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
@@ -248,3 +249,36 @@ export function getSuspendedCallSiteDebugTaskDEV(): ConsoleTask | null {
     );
   }
 }
+
+export function ensureSuspendableThenableStateDEV(
+  thenableState: ThenableState,
+): () => void {
+  if (__DEV__) {
+    const lastThenable = thenableState[thenableState.length - 1];
+    switch (lastThenable.status) {
+      case 'fulfilled':
+        const previousThenableValue = lastThenable.value;
+        delete lastThenable.value;
+        delete (lastThenable: any).status;
+        return () => {
+          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

@@ -163,8 +163,8 @@ describe('ReactServer', () => {
       __DEV__
         ? '' +
             // The concrete location may change as this test is updated.
-            // Just make sure they still point at the same code
-            '\n    at Component (./ReactServer-test.js:95:13)' +
+            // Just make sure they still point 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)'
         : null,