Fix bug that wasn't picking up the right stack at suspended boundaries

This makes it more explicit which stack we pass in to be retained by the
task.

Author: sebmarkbage

packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js

@@ -1058,6 +1058,110 @@ describe('ReactDOMFizzServer', () => {
     );
   });
 
+  function normalizeCodeLocInfo(str) {
+    return (
+      str &&
+      str.replace(/\n +(?:at|in) ([\S]+)[^\n]*/g, function(m, name) {
+        return '\n    in ' + name + ' (at **)';
+      })
+    );
+  }
+
+  // @gate experimental
+  it('should include a component stack across suspended boundaries', async () => {
+    function B() {
+      const children = [readText('Hello'), readText('World')];
+      // Intentionally trigger a key warning here.
+      return (
+        <div>
+          {children.map(t => (
+            <span>{t}</span>
+          ))}
+        </div>
+      );
+    }
+    function C() {
+      return (
+        <inCorrectTag>
+          <Text text="Loading" />
+        </inCorrectTag>
+      );
+    }
+    function A() {
+      return (
+        <div>
+          <Suspense fallback={<C />}>
+            <B />
+          </Suspense>
+        </div>
+      );
+    }
+
+    // We can't use the toErrorDev helper here because this is an async act.
+    const originalConsoleError = console.error;
+    const mockError = jest.fn();
+    console.error = (...args) => {
+      mockError(...args.map(normalizeCodeLocInfo));
+    };
+
+    try {
+      await act(async () => {
+        const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
+          <A />,
+          writable,
+        );
+        startWriting();
+      });
+
+      expect(getVisibleChildren(container)).toEqual(
+        <div>
+          <incorrecttag>Loading</incorrecttag>
+        </div>,
+      );
+
+      expect(mockError).toHaveBeenCalledWith(
+        'Warning: <%s /> is using incorrect casing. Use PascalCase for React components, or lowercase for HTML elements.%s',
+        'inCorrectTag',
+        '\n' +
+          '    in inCorrectTag (at **)\n' +
+          '    in C (at **)\n' +
+          '    in Suspense (at **)\n' +
+          '    in div (at **)\n' +
+          '    in A (at **)',
+      );
+      mockError.mockClear();
+
+      await act(async () => {
+        resolveText('Hello');
+        resolveText('World');
+      });
+
+      expect(mockError).toHaveBeenCalledWith(
+        'Warning: Each child in a list should have a unique "key" prop.%s%s' +
+          ' See https://reactjs.org/link/warning-keys for more information.%s',
+        '\n\nCheck the top-level render call using <div>.',
+        '',
+        '\n' +
+          '    in span (at **)\n' +
+          '    in B (at **)\n' +
+          '    in Suspense (at **)\n' +
+          '    in div (at **)\n' +
+          '    in A (at **)',
+      );
+
+      expect(getVisibleChildren(container)).toEqual(
+        <div>
+          <div>
+            <span>Hello</span>
+            <span>World</span>
+          </div>
+        </div>,
+      );
+    } finally {
+      console.error = originalConsoleError;
+    }
+  });
+
   // @gate experimental
   it('should can suspend in a class component with legacy context', async () => {
     class TestProvider extends React.Component {

packages/react-server/src/ReactFizzServer.js

@@ -314,7 +314,7 @@ function createTask(
     assignID,
   }: any);
   if (__DEV__) {
-    task.componentStack = currentTaskInDEV ? task.componentStack : null;
+    task.componentStack = null;
   }
   abortSet.add(task);
   return task;
@@ -465,6 +465,7 @@ function renderSuspenseBoundary(
       // This must have been the last segment we were waiting on. This boundary is now complete.
       // Therefore we won't need the fallback. We early return so that we don't have to create
       // the fallback.
+      popComponentStackInDEV(task);
       return;
     }
   } catch (error) {
@@ -477,7 +478,6 @@ function renderSuspenseBoundary(
   } finally {
     task.blockedBoundary = parentBoundary;
     task.blockedSegment = parentSegment;
-    popComponentStackInDEV(task);
   }
 
   // This injects an extra segment just to contain an empty tag with an ID.
@@ -505,9 +505,14 @@ function renderSuspenseBoundary(
     task.context,
     null,
   );
+  if (__DEV__) {
+    suspendedFallbackTask.componentStack = task.componentStack;
+  }
   // TODO: This should be queued at a separate lower priority queue so that we only work
   // on preparing fallbacks if we don't have any more main content to task on.
   request.pingedTasks.push(suspendedFallbackTask);
+
+  popComponentStackInDEV(task);
 }
 
 function renderHostElement(
@@ -1233,6 +1238,13 @@ function spawnNewSuspendedTask(
     task.context,
     task.assignID,
   );
+  if (__DEV__) {
+    if (task.componentStack !== null) {
+      // We pop one task off the stack because the node that suspended will be tried again,
+      // which will add it back onto the stack.
+      newTask.componentStack = task.componentStack.parent;
+    }
+  }
   // We've delegated the assignment.
   task.assignID = null;
   const ping = newTask.ping;