Increase retryTime for increased priority dehydrated boundaries

Author: sebmarkbage

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

@@ -114,4 +114,88 @@ describe('ReactDOMServerSelectiveHydration', () => {
 
     document.body.removeChild(container);
   });
+
+  it('hydrates at higher pri if sync did not work first time', async () => {
+    let suspend = false;
+    let resolve;
+    let promise = new Promise(resolvePromise => (resolve = resolvePromise));
+
+    function Child({text}) {
+      if ((text === 'A' || text === 'D') && suspend) {
+        throw promise;
+      }
+      Scheduler.unstable_yieldValue(text);
+      return (
+        <span
+          onClick={e => {
+            e.preventDefault();
+            Scheduler.unstable_yieldValue('Clicked ' + text);
+          }}>
+          {text}
+        </span>
+      );
+    }
+
+    function App() {
+      Scheduler.unstable_yieldValue('App');
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            <Child text="A" />
+          </Suspense>
+          <Suspense fallback="Loading...">
+            <Child text="B" />
+          </Suspense>
+          <Suspense fallback="Loading...">
+            <Child text="C" />
+          </Suspense>
+          <Suspense fallback="Loading...">
+            <Child text="D" />
+          </Suspense>
+        </div>
+      );
+    }
+
+    let finalHTML = ReactDOMServer.renderToString(<App />);
+
+    expect(Scheduler).toHaveYielded(['App', 'A', 'B', 'C', 'D']);
+
+    let container = document.createElement('div');
+    // We need this to be in the document since we'll dispatch events on it.
+    document.body.appendChild(container);
+
+    container.innerHTML = finalHTML;
+
+    let span = container.getElementsByTagName('span')[3];
+
+    suspend = true;
+
+    // A and D will be suspended. We'll click on D which should take
+    // priority, after we unsuspend.
+    let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+    root.render(<App />);
+
+    // Nothing has been hydrated so far.
+    expect(Scheduler).toHaveYielded([]);
+
+    // This click target cannot be hydrated yet because it's suspended.
+    let result = dispatchClickEvent(span);
+
+    expect(Scheduler).toHaveYielded(['App']);
+
+    expect(result).toBe(true);
+
+    // Continuing rendering will render B next.
+    expect(Scheduler).toFlushAndYield(['B', 'C']);
+
+    suspend = false;
+    resolve();
+    await promise;
+
+    // After the click, we should prioritize D and the Click first,
+    // and only after that render A and C.
+    expect(Scheduler).toFlushAndYield(['D', 'Clicked D', 'A']);
+
+    document.body.removeChild(container);
+  });
 });

packages/react-dom/src/client/ReactDOM.js

@@ -40,6 +40,7 @@ import {
   flushPassiveEffects,
   IsThisRendererActing,
   attemptSynchronousHydration,
+  attemptHydration,
 } from 'react-reconciler/inline.dom';
 import {createPortal as createPortalImpl} from 'shared/ReactPortal';
 import {canUseDOM} from 'shared/ExecutionEnvironment';
@@ -75,7 +76,10 @@ import {
 } from './ReactDOMComponentTree';
 import {restoreControlledState} from './ReactDOMComponent';
 import {dispatchEvent} from '../events/ReactDOMEventListener';
-import {setAttemptSynchronousHydration} from '../events/ReactDOMEventReplaying';
+import {
+  setAttemptSynchronousHydration,
+  setAttemptHydration,
+} from '../events/ReactDOMEventReplaying';
 import {eagerlyTrapReplayableEvents} from '../events/ReactDOMEventReplaying';
 import {
   ELEMENT_NODE,
@@ -86,6 +90,7 @@ import {
 import {ROOT_ATTRIBUTE_NAME} from '../shared/DOMProperty';
 
 setAttemptSynchronousHydration(attemptSynchronousHydration);
+setAttemptHydration(attemptHydration);
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
 

packages/react-dom/src/events/ReactDOMEventReplaying.js

@@ -37,6 +37,12 @@ export function setAttemptSynchronousHydration(fn: (fiber: Object) => void) {
   attemptSynchronousHydration = fn;
 }
 
+let attemptHydration: (fiber: Object) => void;
+
+export function setAttemptHydration(fn: (fiber: Object) => void) {
+  attemptHydration = fn;
+}
+
 // TODO: Upgrade this definition once we're on a newer version of Flow that
 // has this definition built-in.
 type PointerEvent = Event & {

packages/react-reconciler/src/ReactFiberReconciler.js

@@ -20,7 +20,10 @@ import {FundamentalComponent} from 'shared/ReactWorkTags';
 import type {ReactNodeList} from 'shared/ReactTypes';
 import type {ExpirationTime} from './ReactFiberExpirationTime';
 import type {SuspenseConfig} from './ReactFiberSuspenseConfig';
-import type {SuspenseHydrationCallbacks} from './ReactFiberSuspenseComponent';
+import type {
+  SuspenseHydrationCallbacks,
+  SuspenseState,
+} from './ReactFiberSuspenseComponent';
 
 import {
   findCurrentHostFiber,
@@ -378,10 +381,43 @@ export function attemptSynchronousHydration(fiber: Fiber): void {
       break;
     case SuspenseComponent:
       flushSync(() => scheduleWork(fiber, Sync));
+      // If we're still blocked after this, we need to increase
+      // the priority of any promises resolving within this
+      // boundary so that they next attempt also has higher pri.
+      markRetryTimeIfNotHydrated(fiber, Sync);
       break;
   }
 }
 
+function markRetryTimeImpl(fiber: Fiber, retryTime: ExpirationTime) {
+  let suspenseState: null | SuspenseState = fiber.memoizedState;
+  if (suspenseState !== null && suspenseState.dehydrated !== null) {
+    if (suspenseState.retryTime < retryTime) {
+      suspenseState.retryTime = retryTime;
+    }
+  }
+}
+
+// Increases the priority of thennables when they resolve within this boundary.
+function markRetryTimeIfNotHydrated(fiber: Fiber, retryTime: ExpirationTime) {
+  markRetryTimeImpl(fiber, retryTime);
+  let alternate = fiber.alternate;
+  if (alternate) {
+    markRetryTimeImpl(alternate, retryTime);
+  }
+}
+
+export function attemptHydration(fiber: Fiber): void {
+  if (fiber.tag !== SuspenseComponent) {
+    // We ignore HostRoots here because we can't increase
+    // their priority and they should not suspend on I/O,
+    // since you have to wrap anything that might suspend in
+    // Suspense.
+    return;
+  }
+  // TODO
+}
+
 export {findHostInstance};
 
 export {findHostInstanceWithWarning};