[Partial Hydration] Attempt hydration at a higher pri first if props/context changes (#16352)

* Test that we can suspend updates while waiting to hydrate

* Attempt hydration at a higher pri first if props/context changes

* Retrying a dehydrated boundary pings at the earliest forced time

This might quickly become an already expired time.

* Mark the render as delayed if we have to retry

This allows the suspense config to kick in and we can wait for much longer
before we're forced to give up on hydrating.

fixtures/ssr/src/components/Chrome.js

@@ -26,7 +26,16 @@ export default class Chrome extends Component {
           <Theme.Provider value={this.state.theme}>
             {this.props.children}
             <div>
-              <ThemeToggleButton onChange={theme => this.setState({theme})} />
+              <ThemeToggleButton
+                onChange={theme => {
+                  React.unstable_withSuspenseConfig(
+                    () => {
+                      this.setState({theme});
+                    },
+                    {timeoutMs: 6000}
+                  );
+                }}
+              />
             </div>
           </Theme.Provider>
           <script

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

@@ -277,7 +277,7 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(container.firstChild.children[1].textContent).toBe('After');
   });
 
-  it('regenerates the content if props have changed before hydration completes', async () => {
+  it('blocks updates to hydrate the content first if props have changed', async () => {
     let suspend = false;
     let resolve;
     let promise = new Promise(resolvePromise => (resolve = resolvePromise));
@@ -331,14 +331,14 @@ describe('ReactDOMServerPartialHydration', () => {
     resolve();
     await promise;
 
-    // Flushing both of these in the same batch won't be able to hydrate so we'll
-    // probably throw away the existing subtree.
+    // This should first complete the hydration and then flush the update onto the hydrated state.
     Scheduler.unstable_flushAll();
     jest.runAllTimers();
 
-    // Pick up the new span. In an ideal implementation this might be the same span
-    // but patched up. At the time of writing, this will be a new span though.
-    span = container.getElementsByTagName('span')[0];
+    // The new span should be the same since we should have successfully hydrated
+    // before changing it.
+    let newSpan = container.getElementsByTagName('span')[0];
+    expect(span).toBe(newSpan);
 
     // We should now have fully rendered with a ref on the new span.
     expect(ref.current).toBe(span);
@@ -562,7 +562,87 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(container.textContent).toBe('Hi Hi');
   });
 
-  it('regenerates the content if context has changed before hydration completes', async () => {
+  it('hydrates first if props changed but we are able to resolve within a timeout', async () => {
+    let suspend = false;
+    let resolve;
+    let promise = new Promise(resolvePromise => (resolve = resolvePromise));
+    let ref = React.createRef();
+
+    function Child({text}) {
+      if (suspend) {
+        throw promise;
+      } else {
+        return text;
+      }
+    }
+
+    function App({text, className}) {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <span ref={ref} className={className}>
+              <Child text={text} />
+            </span>
+          </Suspense>
+        </div>
+      );
+    }
+
+    suspend = false;
+    let finalHTML = ReactDOMServer.renderToString(
+      <App text="Hello" className="hello" />,
+    );
+    let container = document.createElement('div');
+    container.innerHTML = finalHTML;
+
+    let span = container.getElementsByTagName('span')[0];
+
+    // On the client we don't have all data yet but we want to start
+    // hydrating anyway.
+    suspend = true;
+    let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+    root.render(<App text="Hello" className="hello" />);
+    Scheduler.unstable_flushAll();
+    jest.runAllTimers();
+
+    expect(ref.current).toBe(null);
+    expect(container.textContent).toBe('Hello');
+
+    // Render an update with a long timeout.
+    React.unstable_withSuspenseConfig(
+      () => root.render(<App text="Hi" className="hi" />),
+      {timeoutMs: 5000},
+    );
+
+    // This shouldn't force the fallback yet.
+    Scheduler.unstable_flushAll();
+
+    expect(ref.current).toBe(null);
+    expect(container.textContent).toBe('Hello');
+
+    // Resolving the promise so that rendering can complete.
+    suspend = false;
+    resolve();
+    await promise;
+
+    // This should first complete the hydration and then flush the update onto the hydrated state.
+    Scheduler.unstable_flushAll();
+    jest.runAllTimers();
+
+    // The new span should be the same since we should have successfully hydrated
+    // before changing it.
+    let newSpan = container.getElementsByTagName('span')[0];
+    expect(span).toBe(newSpan);
+
+    // We should now have fully rendered with a ref on the new span.
+    expect(ref.current).toBe(span);
+    expect(container.textContent).toBe('Hi');
+    // If we ended up hydrating the existing content, we won't have properly
+    // patched up the tree, which might mean we haven't patched the className.
+    expect(span.className).toBe('hi');
+  });
+
+  it('blocks the update to hydrate first if context has changed', async () => {
     let suspend = false;
     let resolve;
     let promise = new Promise(resolvePromise => (resolve = resolvePromise));
@@ -630,14 +710,13 @@ describe('ReactDOMServerPartialHydration', () => {
     resolve();
     await promise;
 
-    // Flushing both of these in the same batch won't be able to hydrate so we'll
-    // probably throw away the existing subtree.
+    // This should first complete the hydration and then flush the update onto the hydrated state.
     Scheduler.unstable_flushAll();
     jest.runAllTimers();
 
-    // Pick up the new span. In an ideal implementation this might be the same span
-    // but patched up. At the time of writing, this will be a new span though.
-    span = container.getElementsByTagName('span')[0];
+    // Since this should have been hydrated, this should still be the same span.
+    let newSpan = container.getElementsByTagName('span')[0];
+    expect(newSpan).toBe(span);
 
     // We should now have fully rendered with a ref on the new span.
     expect(ref.current).toBe(span);
@@ -1421,4 +1500,85 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(ref1.current).toBe(span1);
     expect(ref2.current).toBe(span2);
   });
+
+  it('regenerates if it cannot hydrate before changes to props/context expire', async () => {
+    let suspend = false;
+    let promise = new Promise(resolvePromise => {});
+    let ref = React.createRef();
+    let ClassName = React.createContext(null);
+
+    function Child({text}) {
+      let className = React.useContext(ClassName);
+      if (suspend && className !== 'hi' && text !== 'Hi') {
+        // Never suspends on the newer data.
+        throw promise;
+      } else {
+        return (
+          <span ref={ref} className={className}>
+            {text}
+          </span>
+        );
+      }
+    }
+
+    function App({text, className}) {
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <Child text={text} />
+          </Suspense>
+        </div>
+      );
+    }
+
+    suspend = false;
+    let finalHTML = ReactDOMServer.renderToString(
+      <ClassName.Provider value={'hello'}>
+        <App text="Hello" />
+      </ClassName.Provider>,
+    );
+    let container = document.createElement('div');
+    container.innerHTML = finalHTML;
+
+    let span = container.getElementsByTagName('span')[0];
+
+    // On the client we don't have all data yet but we want to start
+    // hydrating anyway.
+    suspend = true;
+    let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+    root.render(
+      <ClassName.Provider value={'hello'}>
+        <App text="Hello" />
+      </ClassName.Provider>,
+    );
+    Scheduler.unstable_flushAll();
+    jest.runAllTimers();
+
+    expect(ref.current).toBe(null);
+    expect(span.textContent).toBe('Hello');
+
+    // Render an update, which will be higher or the same priority as pinging the hydration.
+    // The new update doesn't suspend.
+    root.render(
+      <ClassName.Provider value={'hi'}>
+        <App text="Hi" />
+      </ClassName.Provider>,
+    );
+
+    // Since we're still suspended on the original data, we can't hydrate.
+    // This will force all expiration times to flush.
+    Scheduler.unstable_flushAll();
+    jest.runAllTimers();
+
+    // This will now be a new span because we weren't able to hydrate before
+    let newSpan = container.getElementsByTagName('span')[0];
+    expect(newSpan).not.toBe(span);
+
+    // We should now have fully rendered with a ref on the new span.
+    expect(ref.current).toBe(newSpan);
+    expect(newSpan.textContent).toBe('Hi');
+    // If we ended up hydrating the existing content, we won't have properly
+    // patched up the tree, which might mean we haven't patched the className.
+    expect(newSpan.className).toBe('hi');
+  });
 });

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -171,7 +171,9 @@ import {
 import {
   markSpawnedWork,
   requestCurrentTime,
-  retryTimedOutBoundary,
+  retryDehydratedSuspenseBoundary,
+  scheduleWork,
+  renderDidSuspendDelayIfPossible,
 } from './ReactFiberWorkLoop';
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
@@ -1476,6 +1478,7 @@ function validateFunctionComponentInDev(workInProgress: Fiber, Component: any) {
 
 const SUSPENDED_MARKER: SuspenseState = {
   dehydrated: null,
+  retryTime: Never,
 };
 
 function shouldRemainOnFallback(
@@ -1672,6 +1675,7 @@ function updateSuspenseComponent(
               current,
               workInProgress,
               dehydrated,
+              prevState,
               renderExpirationTime,
             );
           } else if (
@@ -2004,6 +2008,7 @@ function updateDehydratedSuspenseComponent(
   current: Fiber,
   workInProgress: Fiber,
   suspenseInstance: SuspenseInstance,
+  suspenseState: SuspenseState,
   renderExpirationTime: ExpirationTime,
 ): null | Fiber {
   // We should never be hydrating at this point because it is the first pass,
@@ -2033,11 +2038,31 @@ function updateDehydratedSuspenseComponent(
   const hasContextChanged = current.childExpirationTime >= renderExpirationTime;
   if (didReceiveUpdate || hasContextChanged) {
     // This boundary has changed since the first render. This means that we are now unable to
-    // hydrate it. We might still be able to hydrate it using an earlier expiration time but
-    // during this render we can't. Instead, we're going to delete the whole subtree and
-    // instead inject a new real Suspense boundary to take its place, which may render content
-    // or fallback. The real Suspense boundary will suspend for a while so we have some time
-    // to ensure it can produce real content, but all state and pending events will be lost.
+    // hydrate it. We might still be able to hydrate it using an earlier expiration time, if
+    // we are rendering at lower expiration than sync.
+    if (renderExpirationTime < Sync) {
+      if (suspenseState.retryTime <= renderExpirationTime) {
+        // This render is even higher pri than we've seen before, let's try again
+        // at even higher pri.
+        let attemptHydrationAtExpirationTime = renderExpirationTime + 1;
+        suspenseState.retryTime = attemptHydrationAtExpirationTime;
+        scheduleWork(current, attemptHydrationAtExpirationTime);
+        // TODO: Early abort this render.
+      } else {
+        // We have already tried to ping at a higher priority than we're rendering with
+        // so if we got here, we must have failed to hydrate at those levels. We must
+        // now give up. Instead, we're going to delete the whole subtree and instead inject
+        // a new real Suspense boundary to take its place, which may render content
+        // or fallback. This might suspend for a while and if it does we might still have
+        // an opportunity to hydrate before this pass commits.
+      }
+    }
+    // If we have scheduled higher pri work above, this will probably just abort the render
+    // since we now have higher priority work, but in case it doesn't, we need to prepare to
+    // render something, if we time out. Even if that requires us to delete everything and
+    // skip hydration.
+    // Delay having to do this as long as the suspense timeout allows us.
+    renderDidSuspendDelayIfPossible();
     return retrySuspenseComponentWithoutHydrating(
       current,
       workInProgress,
@@ -2059,7 +2084,7 @@ function updateDehydratedSuspenseComponent(
     // Register a callback to retry this boundary once the server has sent the result.
     registerSuspenseInstanceRetry(
       suspenseInstance,
-      retryTimedOutBoundary.bind(null, current),
+      retryDehydratedSuspenseBoundary.bind(null, current),
     );
     return null;
   } else {

packages/react-reconciler/src/ReactFiberHydrationContext.js

@@ -55,6 +55,7 @@ import {
 } from './ReactFiberHostConfig';
 import {enableSuspenseServerRenderer} from 'shared/ReactFeatureFlags';
 import warning from 'shared/warning';
+import {Never} from './ReactFiberExpirationTime';
 
 // The deepest Fiber on the stack involved in a hydration context.
 // This may have been an insertion or a hydration.
@@ -229,6 +230,7 @@ function tryHydrate(fiber, nextInstance) {
         if (suspenseInstance !== null) {
           const suspenseState: SuspenseState = {
             dehydrated: suspenseInstance,
+            retryTime: Never,
           };
           fiber.memoizedState = suspenseState;
           // Store the dehydrated fragment as a child fiber.

packages/react-reconciler/src/ReactFiberSuspenseComponent.js

@@ -9,6 +9,7 @@
 
 import type {Fiber} from './ReactFiber';
 import type {SuspenseInstance} from './ReactFiberHostConfig';
+import type {ExpirationTime} from './ReactFiberExpirationTime';
 import {SuspenseComponent, SuspenseListComponent} from 'shared/ReactWorkTags';
 import {NoEffect, DidCapture} from 'shared/ReactSideEffectTags';
 import {
@@ -28,6 +29,9 @@ export type SuspenseState = {|
   // here to indicate that it is dehydrated (flag) and for quick access
   // to check things like isSuspenseInstancePending.
   dehydrated: null | SuspenseInstance,
+  // Represents the earliest expiration time we should attempt to hydrate
+  // a dehydrated boundary at. Never is the default.
+  retryTime: ExpirationTime,
 |};
 
 export type SuspenseListTailMode = 'collapsed' | 'hidden' | void;

packages/react-reconciler/src/ReactFiberWorkLoop.js

@@ -16,6 +16,7 @@ import type {
 } from './SchedulerWithReactIntegration';
 import type {Interaction} from 'scheduler/src/Tracing';
 import type {SuspenseConfig} from './ReactFiberSuspenseConfig';
+import type {SuspenseState} from './ReactFiberSuspenseComponent';
 
 import {
   warnAboutDeprecatedLifecycles,
@@ -2160,18 +2161,23 @@ export function pingSuspendedRoot(
   scheduleCallbackForRoot(root, priorityLevel, suspendedTime);
 }
 
-export function retryTimedOutBoundary(boundaryFiber: Fiber) {
+function retryTimedOutBoundary(
+  boundaryFiber: Fiber,
+  retryTime: ExpirationTime,
+) {
   // The boundary fiber (a Suspense component or SuspenseList component)
   // previously was rendered in its fallback state. One of the promises that
   // suspended it has resolved, which means at least part of the tree was
   // likely unblocked. Try rendering again, at a new expiration time.
   const currentTime = requestCurrentTime();
-  const suspenseConfig = null; // Retries don't carry over the already committed update.
-  const retryTime = computeExpirationForFiber(
-    currentTime,
-    boundaryFiber,
-    suspenseConfig,
-  );
+  if (retryTime === Never) {
+    const suspenseConfig = null; // Retries don't carry over the already committed update.
+    retryTime = computeExpirationForFiber(
+      currentTime,
+      boundaryFiber,
+      suspenseConfig,
+    );
+  }
   // TODO: Special case idle priority?
   const priorityLevel = inferPriorityFromExpirationTime(currentTime, retryTime);
   const root = markUpdateTimeFromFiberToRoot(boundaryFiber, retryTime);
@@ -2180,12 +2186,26 @@ export function retryTimedOutBoundary(boundaryFiber: Fiber) {
   }
 }
 
+export function retryDehydratedSuspenseBoundary(boundaryFiber: Fiber) {
+  const suspenseState: null | SuspenseState = boundaryFiber.memoizedState;
+  let retryTime = Never;
+  if (suspenseState !== null) {
+    retryTime = suspenseState.retryTime;
+  }
+  retryTimedOutBoundary(boundaryFiber, retryTime);
+}
+
 export function resolveRetryThenable(boundaryFiber: Fiber, thenable: Thenable) {
+  let retryTime = Never; // Default
   let retryCache: WeakSet<Thenable> | Set<Thenable> | null;
   if (enableSuspenseServerRenderer) {
     switch (boundaryFiber.tag) {
       case SuspenseComponent:
         retryCache = boundaryFiber.stateNode;
+        const suspenseState: null | SuspenseState = boundaryFiber.memoizedState;
+        if (suspenseState !== null) {
+          retryTime = suspenseState.retryTime;
+        }
         break;
       default:
         invariant(
@@ -2204,7 +2224,7 @@ export function resolveRetryThenable(boundaryFiber: Fiber, thenable: Thenable) {
     retryCache.delete(thenable);
   }
 
-  retryTimedOutBoundary(boundaryFiber);
+  retryTimedOutBoundary(boundaryFiber, retryTime);
 }
 
 // Computes the next Just Noticeable Difference (JND) boundary.