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

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;