[Selective Hydration] Increase priority for non-synchronous discrete events and retries (#16935)

* Increase retryTime for increased priority dehydrated boundaries

* Increaese the priority to user blocking for every next discrete boundary

Author: sebmarkbage

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

@@ -114,4 +114,180 @@ 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 spanD = 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(spanD);
+
+    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);
+  });
+
+  it('hydrates at higher pri for secondary discrete events', 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 spanA = container.getElementsByTagName('span')[0];
+    let spanC = container.getElementsByTagName('span')[2];
+    let spanD = 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 the first is Suspended.
+    dispatchClickEvent(spanA);
+    dispatchClickEvent(spanC);
+    dispatchClickEvent(spanD);
+
+    expect(Scheduler).toHaveYielded(['App']);
+
+    suspend = false;
+    resolve();
+    await promise;
+
+    // We should prioritize hydrating A, C and D first since we clicked in
+    // them. Only after they're done will we hydrate B.
+    expect(Scheduler).toFlushAndYield([
+      'A',
+      'Clicked A',
+      'C',
+      'Clicked C',
+      'D',
+      'Clicked D',
+      // B should render last since it wasn't clicked.
+      'B',
+    ]);
+
+    document.body.removeChild(container);
+  });
 });

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

@@ -40,6 +40,7 @@ import {
   flushPassiveEffects,
   IsThisRendererActing,
   attemptSynchronousHydration,
+  attemptUserBlockingHydration,
 } 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,
+  setAttemptUserBlockingHydration,
+} 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);
+setAttemptUserBlockingHydration(attemptUserBlockingHydration);
 
 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 attemptUserBlockingHydration: (fiber: Object) => void;
+
+export function setAttemptUserBlockingHydration(fn: (fiber: Object) => void) {
+  attemptUserBlockingHydration = fn;
+}
+
 // TODO: Upgrade this definition once we're on a newer version of Flow that
 // has this definition built-in.
 type PointerEvent = Event & {
@@ -436,6 +442,12 @@ function replayUnblockedEvents() {
     let nextDiscreteEvent = queuedDiscreteEvents[0];
     if (nextDiscreteEvent.blockedOn !== null) {
       // We're still blocked.
+      // Increase the priority of this boundary to unblock
+      // the next discrete event.
+      let fiber = getInstanceFromNode(nextDiscreteEvent.blockedOn);
+      if (fiber !== null) {
+        attemptUserBlockingHydration(fiber);
+      }
       break;
     }
     let nextBlockedOn = attemptToDispatchEvent(

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,
@@ -75,7 +78,7 @@ import {
   current as ReactCurrentFiberCurrent,
 } from './ReactCurrentFiber';
 import {StrictMode} from './ReactTypeOfMode';
-import {Sync} from './ReactFiberExpirationTime';
+import {Sync, computeInteractiveExpiration} from './ReactFiberExpirationTime';
 import {requestCurrentSuspenseConfig} from './ReactFiberSuspenseConfig';
 import {
   scheduleRefresh,
@@ -378,10 +381,46 @@ 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.
+      let retryExpTime = computeInteractiveExpiration(requestCurrentTime());
+      markRetryTimeIfNotHydrated(fiber, retryExpTime);
       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 attemptUserBlockingHydration(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;
+  }
+  let expTime = computeInteractiveExpiration(requestCurrentTime());
+  scheduleWork(fiber, expTime);
+  markRetryTimeIfNotHydrated(fiber, expTime);
+}
+
 export {findHostInstance};
 
 export {findHostInstanceWithWarning};