[Selective Hydration] ReactDOM.unstable_scheduleHydration(domNode) (#17004)

Adds an API to explicitly prioritize hydrating the path to a particular DOM node without relying on events to do it.

The API uses the current scheduler priority to schedule it. For the same priority, the last one wins. This allows a similar effect as continuous events. This is useful for example to hydrate based on scroll position, or prioritize components that will upgrade to client-rendered-only content.

I considered having an API that explicitly overrides the current target(s). However that makes it difficult to coordinate across components in an app.

This just hydrates one target at a time but if it is blocked on I/O we could consider increasing priority of later targets too.

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

@@ -440,4 +440,52 @@ describe('ReactDOMServerSelectiveHydration', () => {
 
     document.body.removeChild(container);
   });
+
+  it('hydrates the last explicitly hydrated target at higher priority', async () => {
+    function Child({text}) {
+      Scheduler.unstable_yieldValue(text);
+      return <span>{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>
+        </div>
+      );
+    }
+
+    let finalHTML = ReactDOMServer.renderToString(<App />);
+
+    expect(Scheduler).toHaveYielded(['App', 'A', 'B', 'C']);
+
+    let container = document.createElement('div');
+    container.innerHTML = finalHTML;
+
+    let spanB = container.getElementsByTagName('span')[1];
+    let spanC = container.getElementsByTagName('span')[2];
+
+    let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+    root.render(<App />);
+
+    // Nothing has been hydrated so far.
+    expect(Scheduler).toHaveYielded([]);
+
+    // Increase priority of B and then C.
+    ReactDOM.unstable_scheduleHydration(spanB);
+    ReactDOM.unstable_scheduleHydration(spanC);
+
+    // We should prioritize hydrating C first because the last added
+    // gets highest priority followed by the next added.
+    expect(Scheduler).toFlushAndYield(['App', 'C', 'B', 'A']);
+  });
 });

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

@@ -42,6 +42,7 @@ import {
   attemptSynchronousHydration,
   attemptUserBlockingHydration,
   attemptContinuousHydration,
+  attemptHydrationAtCurrentPriority,
 } from 'react-reconciler/inline.dom';
 import {createPortal as createPortalImpl} from 'shared/ReactPortal';
 import {canUseDOM} from 'shared/ExecutionEnvironment';
@@ -81,8 +82,10 @@ import {
   setAttemptSynchronousHydration,
   setAttemptUserBlockingHydration,
   setAttemptContinuousHydration,
+  setAttemptHydrationAtCurrentPriority,
+  eagerlyTrapReplayableEvents,
+  queueExplicitHydrationTarget,
 } from '../events/ReactDOMEventReplaying';
-import {eagerlyTrapReplayableEvents} from '../events/ReactDOMEventReplaying';
 import {
   ELEMENT_NODE,
   COMMENT_NODE,
@@ -94,6 +97,7 @@ import {ROOT_ATTRIBUTE_NAME} from '../shared/DOMProperty';
 setAttemptSynchronousHydration(attemptSynchronousHydration);
 setAttemptUserBlockingHydration(attemptUserBlockingHydration);
 setAttemptContinuousHydration(attemptContinuousHydration);
+setAttemptHydrationAtCurrentPriority(attemptHydrationAtCurrentPriority);
 
 const ReactCurrentOwner = ReactSharedInternals.ReactCurrentOwner;
 
@@ -841,6 +845,12 @@ const ReactDOM: Object = {
   unstable_createSyncRoot: createSyncRoot,
   unstable_flushControlled: flushControlled,
 
+  unstable_scheduleHydration(target: Node) {
+    if (target) {
+      queueExplicitHydrationTarget(target);
+    }
+  },
+
   __SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED: {
     // Keep in sync with ReactDOMUnstableNativeDependencies.js
     // ReactTestUtils.js, and ReactTestUtilsAct.js. This is an array for better minification.

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

@@ -11,15 +11,23 @@ import type {AnyNativeEvent} from 'legacy-events/PluginModuleType';
 import type {Container, SuspenseInstance} from '../client/ReactDOMHostConfig';
 import type {DOMTopLevelEventType} from 'legacy-events/TopLevelEventTypes';
 import type {EventSystemFlags} from 'legacy-events/EventSystemFlags';
+import type {FiberRoot} from 'react-reconciler/src/ReactFiberRoot';
 
 import {
   enableFlareAPI,
   enableSelectiveHydration,
 } from 'shared/ReactFeatureFlags';
 import {
+  unstable_runWithPriority as runWithPriority,
   unstable_scheduleCallback as scheduleCallback,
   unstable_NormalPriority as NormalPriority,
+  unstable_getCurrentPriorityLevel as getCurrentPriorityLevel,
 } from 'scheduler';
+import {
+  getNearestMountedFiber,
+  getContainerFromFiber,
+  getSuspenseInstanceFromFiber,
+} from 'react-reconciler/reflection';
 import {
   attemptToDispatchEvent,
   trapEventForResponderEventSystem,
@@ -28,8 +36,12 @@ import {
   getListeningSetForElement,
   listenToTopLevel,
 } from './ReactBrowserEventEmitter';
-import {getInstanceFromNode} from '../client/ReactDOMComponentTree';
+import {
+  getInstanceFromNode,
+  getClosestInstanceFromNode,
+} from '../client/ReactDOMComponentTree';
 import {unsafeCastDOMTopLevelTypeToString} from 'legacy-events/TopLevelEventTypes';
+import {HostRoot, SuspenseComponent} from 'shared/ReactWorkTags';
 
 let attemptSynchronousHydration: (fiber: Object) => void;
 
@@ -49,6 +61,14 @@ export function setAttemptContinuousHydration(fn: (fiber: Object) => void) {
   attemptContinuousHydration = fn;
 }
 
+let attemptHydrationAtCurrentPriority: (fiber: Object) => void;
+
+export function setAttemptHydrationAtCurrentPriority(
+  fn: (fiber: Object) => void,
+) {
+  attemptHydrationAtCurrentPriority = fn;
+}
+
 // TODO: Upgrade this definition once we're on a newer version of Flow that
 // has this definition built-in.
 type PointerEvent = Event & {
@@ -124,6 +144,13 @@ let queuedPointers: Map<number, QueuedReplayableEvent> = new Map();
 let queuedPointerCaptures: Map<number, QueuedReplayableEvent> = new Map();
 // We could consider replaying selectionchange and touchmoves too.
 
+type QueuedHydrationTarget = {|
+  blockedOn: null | Container | SuspenseInstance,
+  target: Node,
+  priority: number,
+|};
+let queuedExplicitHydrationTargets: Array<QueuedHydrationTarget> = [];
+
 export function hasQueuedDiscreteEvents(): boolean {
   return queuedDiscreteEvents.length > 0;
 }
@@ -422,6 +449,64 @@ export function queueIfContinuousEvent(
   return false;
 }
 
+// Check if this target is unblocked. Returns true if it's unblocked.
+function attemptExplicitHydrationTarget(
+  queuedTarget: QueuedHydrationTarget,
+): void {
+  // TODO: This function shares a lot of logic with attemptToDispatchEvent.
+  // Try to unify them. It's a bit tricky since it would require two return
+  // values.
+  let targetInst = getClosestInstanceFromNode(queuedTarget.target);
+  if (targetInst !== null) {
+    let nearestMounted = getNearestMountedFiber(targetInst);
+    if (nearestMounted !== null) {
+      const tag = nearestMounted.tag;
+      if (tag === SuspenseComponent) {
+        let instance = getSuspenseInstanceFromFiber(nearestMounted);
+        if (instance !== null) {
+          // We're blocked on hydrating this boundary.
+          // Increase its priority.
+          queuedTarget.blockedOn = instance;
+          runWithPriority(queuedTarget.priority, () => {
+            attemptHydrationAtCurrentPriority(nearestMounted);
+          });
+          return;
+        }
+      } else if (tag === HostRoot) {
+        const root: FiberRoot = nearestMounted.stateNode;
+        if (root.hydrate) {
+          queuedTarget.blockedOn = getContainerFromFiber(nearestMounted);
+          // We don't currently have a way to increase the priority of
+          // a root other than sync.
+          return;
+        }
+      }
+    }
+  }
+  queuedTarget.blockedOn = null;
+}
+
+export function queueExplicitHydrationTarget(target: Node): void {
+  if (enableSelectiveHydration) {
+    let priority = getCurrentPriorityLevel();
+    const queuedTarget: QueuedHydrationTarget = {
+      blockedOn: null,
+      target: target,
+      priority: priority,
+    };
+    let i = 0;
+    for (; i < queuedExplicitHydrationTargets.length; i++) {
+      if (priority <= queuedExplicitHydrationTargets[i].priority) {
+        break;
+      }
+    }
+    queuedExplicitHydrationTargets.splice(i, 0, queuedTarget);
+    if (i === 0) {
+      attemptExplicitHydrationTarget(queuedTarget);
+    }
+  }
+}
+
 function attemptReplayContinuousQueuedEvent(
   queuedEvent: QueuedReplayableEvent,
 ): boolean {
@@ -544,4 +629,25 @@ export function retryIfBlockedOn(
     scheduleCallbackIfUnblocked(queuedEvent, unblocked);
   queuedPointers.forEach(unblock);
   queuedPointerCaptures.forEach(unblock);
+
+  for (let i = 0; i < queuedExplicitHydrationTargets.length; i++) {
+    let queuedTarget = queuedExplicitHydrationTargets[i];
+    if (queuedTarget.blockedOn === unblocked) {
+      queuedTarget.blockedOn = null;
+    }
+  }
+
+  while (queuedExplicitHydrationTargets.length > 0) {
+    let nextExplicitTarget = queuedExplicitHydrationTargets[0];
+    if (nextExplicitTarget.blockedOn !== null) {
+      // We're still blocked.
+      break;
+    } else {
+      attemptExplicitHydrationTarget(nextExplicitTarget);
+      if (nextExplicitTarget.blockedOn === null) {
+        // We're unblocked.
+        queuedExplicitHydrationTargets.shift();
+      }
+    }
+  }
 }

packages/react-reconciler/src/ReactFiberReconciler.js

@@ -438,6 +438,18 @@ export function attemptContinuousHydration(fiber: Fiber): void {
   markRetryTimeIfNotHydrated(fiber, expTime);
 }
 
+export function attemptHydrationAtCurrentPriority(fiber: Fiber): void {
+  if (fiber.tag !== SuspenseComponent) {
+    // We ignore HostRoots here because we can't increase
+    // their priority other than synchronously flush it.
+    return;
+  }
+  const currentTime = requestCurrentTime();
+  const expTime = computeExpirationForFiber(currentTime, fiber, null);
+  scheduleWork(fiber, expTime);
+  markRetryTimeIfNotHydrated(fiber, expTime);
+}
+
 export {findHostInstance};
 
 export {findHostInstanceWithWarning};