Don't recreate the same fallback on the client if hydrating suspends (#24236)

* Delay showing fallback if hydrating suspends

* Fix up

* Include all non-urgent lanes

* Moar tests

* Add test for transitions

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

@@ -869,16 +869,16 @@ describe('ReactDOMFizzServer', () => {
     });
 
     // We still can't render it on the client.
-    expect(Scheduler).toFlushAndYield([
-      'The server could not finish this Suspense boundary, likely due to an ' +
-        'error during server rendering. Switched to client rendering.',
-    ]);
+    expect(Scheduler).toFlushAndYield([]);
     expect(getVisibleChildren(container)).toEqual(<div>Loading...</div>);
 
     // We now resolve it on the client.
     resolveText('Hello');
 
-    Scheduler.unstable_flushAll();
+    expect(Scheduler).toFlushAndYield([
+      'The server could not finish this Suspense boundary, likely due to an ' +
+        'error during server rendering. Switched to client rendering.',
+    ]);
 
     // The client rendered HTML is now in place.
     expect(getVisibleChildren(container)).toEqual(
@@ -2220,6 +2220,286 @@ describe('ReactDOMFizzServer', () => {
     },
   );
 
+  // @gate experimental
+  it('does not recreate the fallback if server errors and hydration suspends', async () => {
+    let isClient = false;
+
+    function Child() {
+      if (isClient) {
+        readText('Yay!');
+      } else {
+        throw Error('Oops.');
+      }
+      Scheduler.unstable_yieldValue('Yay!');
+      return 'Yay!';
+    }
+
+    const fallbackRef = React.createRef();
+    function App() {
+      return (
+        <div>
+          <Suspense fallback={<p ref={fallbackRef}>Loading...</p>}>
+            <span>
+              <Child />
+            </span>
+          </Suspense>
+        </div>
+      );
+    }
+    await act(async () => {
+      const {pipe} = ReactDOMFizzServer.renderToPipeableStream(<App />, {
+        onError(error) {
+          Scheduler.unstable_yieldValue('[s!] ' + error.message);
+        },
+      });
+      pipe(writable);
+    });
+    expect(Scheduler).toHaveYielded(['[s!] Oops.']);
+
+    // The server could not complete this boundary, so we'll retry on the client.
+    const serverFallback = container.getElementsByTagName('p')[0];
+    expect(serverFallback.innerHTML).toBe('Loading...');
+
+    // Hydrate the tree. This will suspend.
+    isClient = true;
+    ReactDOMClient.hydrateRoot(container, <App />, {
+      onRecoverableError(error) {
+        Scheduler.unstable_yieldValue('[c!] ' + error.message);
+      },
+    });
+    // This should not report any errors yet.
+    expect(Scheduler).toFlushAndYield([]);
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <p>Loading...</p>
+      </div>,
+    );
+
+    // Normally, hydrating after server error would force a clean client render.
+    // However, it suspended so at best we'd only get the same fallback anyway.
+    // We don't want to recreate the same fallback in the DOM again because
+    // that's extra work and would restart animations etc. Check we don't do that.
+    const clientFallback = container.getElementsByTagName('p')[0];
+    expect(serverFallback).toBe(clientFallback);
+
+    // When we're able to fully hydrate, we expect a clean client render.
+    await act(async () => {
+      resolveText('Yay!');
+    });
+    expect(Scheduler).toFlushAndYield([
+      'Yay!',
+      '[c!] The server could not finish this Suspense boundary, ' +
+        'likely due to an error during server rendering. ' +
+        'Switched to client rendering.',
+    ]);
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        <span>Yay!</span>
+      </div>,
+    );
+  });
+
+  // @gate experimental
+  it(
+    'does not recreate the fallback if server errors and hydration suspends ' +
+      'and root receives a transition',
+    async () => {
+      let isClient = false;
+
+      function Child({color}) {
+        if (isClient) {
+          readText('Yay!');
+        } else {
+          throw Error('Oops.');
+        }
+        Scheduler.unstable_yieldValue('Yay! (' + color + ')');
+        return 'Yay! (' + color + ')';
+      }
+
+      const fallbackRef = React.createRef();
+      function App({color}) {
+        return (
+          <div>
+            <Suspense fallback={<p ref={fallbackRef}>Loading...</p>}>
+              <span>
+                <Child color={color} />
+              </span>
+            </Suspense>
+          </div>
+        );
+      }
+      await act(async () => {
+        const {pipe} = ReactDOMFizzServer.renderToPipeableStream(
+          <App color="red" />,
+          {
+            onError(error) {
+              Scheduler.unstable_yieldValue('[s!] ' + error.message);
+            },
+          },
+        );
+        pipe(writable);
+      });
+      expect(Scheduler).toHaveYielded(['[s!] Oops.']);
+
+      // The server could not complete this boundary, so we'll retry on the client.
+      const serverFallback = container.getElementsByTagName('p')[0];
+      expect(serverFallback.innerHTML).toBe('Loading...');
+
+      // Hydrate the tree. This will suspend.
+      isClient = true;
+      const root = ReactDOMClient.hydrateRoot(container, <App color="red" />, {
+        onRecoverableError(error) {
+          Scheduler.unstable_yieldValue('[c!] ' + error.message);
+        },
+      });
+      // This should not report any errors yet.
+      expect(Scheduler).toFlushAndYield([]);
+      expect(getVisibleChildren(container)).toEqual(
+        <div>
+          <p>Loading...</p>
+        </div>,
+      );
+
+      // Normally, hydrating after server error would force a clean client render.
+      // However, it suspended so at best we'd only get the same fallback anyway.
+      // We don't want to recreate the same fallback in the DOM again because
+      // that's extra work and would restart animations etc. Check we don't do that.
+      const clientFallback = container.getElementsByTagName('p')[0];
+      expect(serverFallback).toBe(clientFallback);
+
+      // Transition updates shouldn't recreate the fallback either.
+      React.startTransition(() => {
+        root.render(<App color="blue" />);
+      });
+      Scheduler.unstable_flushAll();
+      jest.runAllTimers();
+      const clientFallback2 = container.getElementsByTagName('p')[0];
+      expect(clientFallback2).toBe(serverFallback);
+
+      // When we're able to fully hydrate, we expect a clean client render.
+      await act(async () => {
+        resolveText('Yay!');
+      });
+      expect(Scheduler).toFlushAndYield([
+        'Yay! (red)',
+        '[c!] The server could not finish this Suspense boundary, ' +
+          'likely due to an error during server rendering. ' +
+          'Switched to client rendering.',
+        'Yay! (blue)',
+      ]);
+      expect(getVisibleChildren(container)).toEqual(
+        <div>
+          <span>Yay! (blue)</span>
+        </div>,
+      );
+    },
+  );
+
+  // @gate experimental
+  it(
+    'recreates the fallback if server errors and hydration suspends but ' +
+      'client receives new props',
+    async () => {
+      let isClient = false;
+
+      function Child() {
+        const value = 'Yay!';
+        if (isClient) {
+          readText(value);
+        } else {
+          throw Error('Oops.');
+        }
+        Scheduler.unstable_yieldValue(value);
+        return value;
+      }
+
+      const fallbackRef = React.createRef();
+      function App({fallbackText}) {
+        return (
+          <div>
+            <Suspense fallback={<p ref={fallbackRef}>{fallbackText}</p>}>
+              <span>
+                <Child />
+              </span>
+            </Suspense>
+          </div>
+        );
+      }
+
+      await act(async () => {
+        const {pipe} = ReactDOMFizzServer.renderToPipeableStream(
+          <App fallbackText="Loading..." />,
+          {
+            onError(error) {
+              Scheduler.unstable_yieldValue('[s!] ' + error.message);
+            },
+          },
+        );
+        pipe(writable);
+      });
+      expect(Scheduler).toHaveYielded(['[s!] Oops.']);
+
+      const serverFallback = container.getElementsByTagName('p')[0];
+      expect(serverFallback.innerHTML).toBe('Loading...');
+
+      // Hydrate the tree. This will suspend.
+      isClient = true;
+      const root = ReactDOMClient.hydrateRoot(
+        container,
+        <App fallbackText="Loading..." />,
+        {
+          onRecoverableError(error) {
+            Scheduler.unstable_yieldValue('[c!] ' + error.message);
+          },
+        },
+      );
+      // This should not report any errors yet.
+      expect(Scheduler).toFlushAndYield([]);
+      expect(getVisibleChildren(container)).toEqual(
+        <div>
+          <p>Loading...</p>
+        </div>,
+      );
+
+      // Normally, hydration after server error would force a clean client render.
+      // However, that suspended so at best we'd only get a fallback anyway.
+      // We don't want to replace a fallback with the same fallback because
+      // that's extra work and would restart animations etc. Verify we don't do that.
+      const clientFallback1 = container.getElementsByTagName('p')[0];
+      expect(serverFallback).toBe(clientFallback1);
+
+      // However, an update may have changed the fallback props. In that case we have to
+      // actually force it to re-render on the client and throw away the server one.
+      root.render(<App fallbackText="More loading..." />);
+      Scheduler.unstable_flushAll();
+      jest.runAllTimers();
+      expect(Scheduler).toHaveYielded([
+        '[c!] The server could not finish this Suspense boundary, ' +
+          'likely due to an error during server rendering. ' +
+          'Switched to client rendering.',
+      ]);
+      expect(getVisibleChildren(container)).toEqual(
+        <div>
+          <p>More loading...</p>
+        </div>,
+      );
+      // This should be a clean render without reusing DOM.
+      const clientFallback2 = container.getElementsByTagName('p')[0];
+      expect(clientFallback2).not.toBe(clientFallback1);
+
+      // Verify we can still do a clean content render after.
+      await act(async () => {
+        resolveText('Yay!');
+      });
+      expect(Scheduler).toFlushAndYield(['Yay!']);
+      expect(getVisibleChildren(container)).toEqual(
+        <div>
+          <span>Yay!</span>
+        </div>,
+      );
+    },
+  );
+
   // @gate experimental
   it(
     'errors during hydration force a client render at the nearest Suspense ' +
@@ -2293,17 +2573,12 @@ describe('ReactDOMFizzServer', () => {
         },
       });
 
-      // An error logged but instead of surfacing it to the UI, we switched
-      // to client rendering.
-      expect(Scheduler).toFlushAndYield([
-        'Hydration error',
-        'There was an error while hydrating this Suspense boundary. Switched ' +
-          'to client rendering.',
-      ]);
+      // An error happened but instead of surfacing it to the UI, we suspended.
+      expect(Scheduler).toFlushAndYield([]);
       expect(getVisibleChildren(container)).toEqual(
         <div>
           <span />
-          Loading...
+          <span>Yay!</span>
           <span />
         </div>,
       );
@@ -2311,7 +2586,12 @@ describe('ReactDOMFizzServer', () => {
       await act(async () => {
         resolveText('Yay!');
       });
-      expect(Scheduler).toFlushAndYield(['Yay!']);
+      expect(Scheduler).toFlushAndYield([
+        'Yay!',
+        'Hydration error',
+        'There was an error while hydrating this Suspense boundary. Switched ' +
+          'to client rendering.',
+      ]);
       expect(getVisibleChildren(container)).toEqual(
         <div>
           <span />

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

@@ -1046,7 +1046,7 @@ describe('ReactDOMServerHydration', () => {
       });
 
       // @gate __DEV__
-      it('warns when client renders an extra node inside Suspense fallback', () => {
+      it('does not warn when client renders an extra node inside Suspense fallback', () => {
         function Mismatch({isClient}) {
           return (
             <div className="parent">
@@ -1063,27 +1063,18 @@ describe('ReactDOMServerHydration', () => {
             </div>
           );
         }
-        // TODO: Why does this not show a fallback mismatch?
-        // And why is this message different from the other ones?
         if (
           gate(flags => flags.enableClientRenderFallbackOnHydrationMismatch)
         ) {
-          expect(testMismatch(Mismatch)).toMatchInlineSnapshot(`
-            Array [
-              "Caught [The server could not finish this Suspense boundary, likely due to an error during server rendering. Switched to client rendering.]",
-            ]
-          `);
+          // There is no error because we don't actually hydrate fallbacks.
+          expect(testMismatch(Mismatch)).toMatchInlineSnapshot(`Array []`);
         } else {
-          expect(testMismatch(Mismatch)).toMatchInlineSnapshot(`
-            Array [
-              "Caught [The server could not finish this Suspense boundary, likely due to an error during server rendering. Switched to client rendering.]",
-            ]
-          `);
+          expect(testMismatch(Mismatch)).toMatchInlineSnapshot(`Array []`);
         }
       });
 
       // @gate __DEV__
-      it('warns when server renders an extra node inside Suspense fallback', () => {
+      it('does not warn when server renders an extra node inside Suspense fallback', () => {
         function Mismatch({isClient}) {
           return (
             <div className="parent">
@@ -1100,22 +1091,13 @@ describe('ReactDOMServerHydration', () => {
             </div>
           );
         }
-        // TODO: Why does this not show a fallback mismatch?
-        // And why is this message different from the other ones?
         if (
           gate(flags => flags.enableClientRenderFallbackOnHydrationMismatch)
         ) {
-          expect(testMismatch(Mismatch)).toMatchInlineSnapshot(`
-            Array [
-              "Caught [The server could not finish this Suspense boundary, likely due to an error during server rendering. Switched to client rendering.]",
-            ]
-          `);
+          // There is no error because we don't actually hydrate fallbacks.
+          expect(testMismatch(Mismatch)).toMatchInlineSnapshot(`Array []`);
         } else {
-          expect(testMismatch(Mismatch)).toMatchInlineSnapshot(`
-            Array [
-              "Caught [The server could not finish this Suspense boundary, likely due to an error during server rendering. Switched to client rendering.]",
-            ]
-          `);
+          expect(testMismatch(Mismatch)).toMatchInlineSnapshot(`Array []`);
         }
       });
     });

packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js

@@ -2137,10 +2137,7 @@ describe('ReactDOMServerPartialHydration', () => {
     });
 
     suspend = true;
-    expect(Scheduler).toFlushAndYield([
-      'The server could not finish this Suspense boundary, likely due to ' +
-        'an error during server rendering. Switched to client rendering.',
-    ]);
+    expect(Scheduler).toFlushAndYield([]);
 
     // We haven't hydrated the second child but the placeholder is still in the list.
     expect(container.textContent).toBe('ALoading B');

packages/react-reconciler/src/ReactFiberBeginWork.new.js

@@ -2241,6 +2241,7 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
           } else {
             // Suspended but we should no longer be in dehydrated mode.
             // Therefore we now have to render the fallback.
+            renderDidSuspendDelayIfPossible();
             const nextPrimaryChildren = nextProps.children;
             const nextFallbackChildren = nextProps.fallback;
             const fallbackChildFragment = mountSuspenseFallbackAfterRetryWithoutHydrating(

packages/react-reconciler/src/ReactFiberBeginWork.old.js

@@ -2241,6 +2241,7 @@ function updateSuspenseComponent(current, workInProgress, renderLanes) {
           } else {
             // Suspended but we should no longer be in dehydrated mode.
             // Therefore we now have to render the fallback.
+            renderDidSuspendDelayIfPossible();
             const nextPrimaryChildren = nextProps.children;
             const nextFallbackChildren = nextProps.fallback;
             const fallbackChildFragment = mountSuspenseFallbackAfterRetryWithoutHydrating(

packages/react-reconciler/src/ReactFiberLane.new.js

@@ -458,8 +458,9 @@ export function includesNonIdleWork(lanes: Lanes) {
 export function includesOnlyRetries(lanes: Lanes) {
   return (lanes & RetryLanes) === lanes;
 }
-export function includesOnlyTransitions(lanes: Lanes) {
-  return (lanes & TransitionLanes) === lanes;
+export function includesOnlyNonUrgentLanes(lanes: Lanes) {
+  const UrgentLanes = SyncLane | InputContinuousLane | DefaultLane;
+  return (lanes & UrgentLanes) === NoLanes;
 }
 
 export function includesBlockingLane(root: FiberRoot, lanes: Lanes) {

packages/react-reconciler/src/ReactFiberLane.old.js

@@ -458,8 +458,9 @@ export function includesNonIdleWork(lanes: Lanes) {
 export function includesOnlyRetries(lanes: Lanes) {
   return (lanes & RetryLanes) === lanes;
 }
-export function includesOnlyTransitions(lanes: Lanes) {
-  return (lanes & TransitionLanes) === lanes;
+export function includesOnlyNonUrgentLanes(lanes: Lanes) {
+  const UrgentLanes = SyncLane | InputContinuousLane | DefaultLane;
+  return (lanes & UrgentLanes) === NoLanes;
 }
 
 export function includesBlockingLane(root: FiberRoot, lanes: Lanes) {

packages/react-reconciler/src/ReactFiberWorkLoop.new.js

@@ -132,7 +132,7 @@ import {
   pickArbitraryLane,
   includesNonIdleWork,
   includesOnlyRetries,
-  includesOnlyTransitions,
+  includesOnlyNonUrgentLanes,
   includesBlockingLane,
   includesExpiredLane,
   getNextLanes,
@@ -1110,7 +1110,7 @@ function finishConcurrentRender(root, exitStatus, lanes) {
     case RootSuspendedWithDelay: {
       markRootSuspended(root, lanes);
 
-      if (includesOnlyTransitions(lanes)) {
+      if (includesOnlyNonUrgentLanes(lanes)) {
         // This is a transition, so we should exit without committing a
         // placeholder and without scheduling a timeout. Delay indefinitely
         // until we receive more data.

packages/react-reconciler/src/ReactFiberWorkLoop.old.js

@@ -132,7 +132,7 @@ import {
   pickArbitraryLane,
   includesNonIdleWork,
   includesOnlyRetries,
-  includesOnlyTransitions,
+  includesOnlyNonUrgentLanes,
   includesBlockingLane,
   includesExpiredLane,
   getNextLanes,
@@ -1110,7 +1110,7 @@ function finishConcurrentRender(root, exitStatus, lanes) {
     case RootSuspendedWithDelay: {
       markRootSuspended(root, lanes);
 
-      if (includesOnlyTransitions(lanes)) {
+      if (includesOnlyNonUrgentLanes(lanes)) {
         // This is a transition, so we should exit without committing a
         // placeholder and without scheduling a timeout. Delay indefinitely
         // until we receive more data.