Decouple update priority tracking from Scheduler package (#19121)

* Initial currentLanePriority implementation

* Minor updates from review

* Fix typos and enable flag

* Fix feature flags and lint

* Fix simple event tests by switching to withSuspenseConfig

* Don't lower the priority of setPending in startTransition below InputContinuous

* Move currentUpdateLanePriority in commit root into the first effect block

* Refactor requestUpdateLane to log for priority mismatches

Also verifies that the update lane priority matches the scheduler lane priority before using it

* Fix four tests by adding ReactDOM.unstable_runWithPriority

* Fix partial hydration when using update lane priority

* Fix partial hydration when using update lane priority

* Rename feature flag and only log for now

* Move unstable_runWithPriority to ReactFiberReconciler

* Add unstable_runWithPriority to ReactNoopPersistent too

* Bug fixes and performance improvements

* Initial currentLanePriority implementation

* Minor updates from review

* Fix typos and enable flag

* Remove higherLanePriority from ReactDOMEventReplaying.js

* Change warning implementation and startTransition update lane priority

* Inject reconciler functions to avoid importing src/

* Fix feature flags and lint

* Fix simple event tests by switching to withSuspenseConfig

* Don't lower the priority of setPending in startTransition below InputContinuous

* Move currentUpdateLanePriority in commit root into the first effect block

* Refactor requestUpdateLane to log for priority mismatches

Also verifies that the update lane priority matches the scheduler lane priority before using it

* Fix four tests by adding ReactDOM.unstable_runWithPriority

* Fix partial hydration when using update lane priority

* Fix partial hydration when using update lane priority

* Rename feature flag and only log for now

* Move unstable_runWithPriority to ReactFiberReconciler

* Bug fixes and performance improvements

* Remove higherLanePriority from ReactDOMEventReplaying.js

* Change warning implementation and startTransition update lane priority

* Inject reconciler functions to avoid importing src/

* Fixes from bad rebase

Author: rickhanlonii

packages/react-dom/index.classic.fb.js

@@ -34,6 +34,7 @@ export {
   createBlockingRoot as unstable_createBlockingRoot,
   unstable_flushControlled,
   unstable_scheduleHydration,
+  unstable_runWithPriority,
   unstable_renderSubtreeIntoContainer,
   unstable_createPortal,
   unstable_createEventHandle,

packages/react-dom/index.experimental.js

@@ -23,6 +23,8 @@ export {
   createBlockingRoot as unstable_createBlockingRoot,
   unstable_flushControlled,
   unstable_scheduleHydration,
+  // DO NOT USE: Temporarily exposing this to migrate off of Scheduler.runWithPriority.
+  unstable_runWithPriority,
   // Disabled behind disableUnstableRenderSubtreeIntoContainer
   unstable_renderSubtreeIntoContainer,
   // Disabled behind disableUnstableCreatePortal

packages/react-dom/index.js

@@ -25,6 +25,7 @@ export {
   createBlockingRoot as unstable_createBlockingRoot,
   unstable_flushControlled,
   unstable_scheduleHydration,
+  unstable_runWithPriority,
   unstable_renderSubtreeIntoContainer,
   unstable_createPortal,
   unstable_createEventHandle,

packages/react-dom/index.modern.fb.js

@@ -19,6 +19,7 @@ export {
   createBlockingRoot as unstable_createBlockingRoot,
   unstable_flushControlled,
   unstable_scheduleHydration,
+  unstable_runWithPriority,
   unstable_createEventHandle,
   unstable_isNewReconciler,
 } from './src/client/ReactDOM';

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

@@ -35,6 +35,8 @@ import {
   attemptUserBlockingHydration,
   attemptContinuousHydration,
   attemptHydrationAtCurrentPriority,
+  runWithPriority,
+  getCurrentUpdatePriority,
 } from 'react-reconciler/src/ReactFiberReconciler';
 import {createPortal as createPortalImpl} from 'react-reconciler/src/ReactPortal';
 import {canUseDOM} from 'shared/ExecutionEnvironment';
@@ -58,6 +60,8 @@ import {
   setAttemptContinuousHydration,
   setAttemptHydrationAtCurrentPriority,
   queueExplicitHydrationTarget,
+  setGetCurrentUpdatePriority,
+  setAttemptHydrationAtPriority,
 } from '../events/ReactDOMEventReplaying';
 import {setBatchingImplementation} from '../events/ReactDOMUpdateBatching';
 import {
@@ -70,6 +74,8 @@ setAttemptSynchronousHydration(attemptSynchronousHydration);
 setAttemptUserBlockingHydration(attemptUserBlockingHydration);
 setAttemptContinuousHydration(attemptContinuousHydration);
 setAttemptHydrationAtCurrentPriority(attemptHydrationAtCurrentPriority);
+setGetCurrentUpdatePriority(getCurrentUpdatePriority);
+setAttemptHydrationAtPriority(runWithPriority);
 
 let didWarnAboutUnstableCreatePortal = false;
 let didWarnAboutUnstableRenderSubtreeIntoContainer = false;
@@ -205,6 +211,9 @@ export {
   unstable_createPortal,
   // enableCreateEventHandleAPI
   createEventHandle as unstable_createEventHandle,
+  // TODO: Remove this once callers migrate to alternatives.
+  // This should only be used by React internals.
+  runWithPriority as unstable_runWithPriority,
 };
 
 const foundDevTools = injectIntoDevTools({

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

@@ -12,7 +12,10 @@ import type {Container, SuspenseInstance} from '../client/ReactDOMHostConfig';
 import type {DOMTopLevelEventType} from '../events/TopLevelEventTypes';
 import type {ElementListenerMap} from '../client/ReactDOMComponentTree';
 import type {EventSystemFlags} from './EventSystemFlags';
-import type {FiberRoot} from 'react-reconciler/src/ReactInternalTypes';
+import type {
+  FiberRoot,
+  ReactPriorityLevel,
+} from 'react-reconciler/src/ReactInternalTypes';
 
 import {
   enableDeprecatedFlareAPI,
@@ -64,6 +67,23 @@ export function setAttemptHydrationAtCurrentPriority(
   attemptHydrationAtCurrentPriority = fn;
 }
 
+let getCurrentUpdatePriority: () => ReactPriorityLevel;
+
+export function setGetCurrentUpdatePriority(fn: () => ReactPriorityLevel) {
+  getCurrentUpdatePriority = fn;
+}
+
+let attemptHydrationAtPriority: <T>(
+  priority: ReactPriorityLevel,
+  fn: () => T,
+) => T;
+
+export function setAttemptHydrationAtPriority(
+  fn: <T>(priority: ReactPriorityLevel, fn: () => T) => T,
+) {
+  attemptHydrationAtPriority = fn;
+}
+
 // TODO: Upgrade this definition once we're on a newer version of Flow that
 // has this definition built-in.
 type PointerEvent = Event & {
@@ -147,6 +167,7 @@ type QueuedHydrationTarget = {|
   blockedOn: null | Container | SuspenseInstance,
   target: Node,
   priority: number,
+  lanePriority: ReactPriorityLevel,
 |};
 const queuedExplicitHydrationTargets: Array<QueuedHydrationTarget> = [];
 
@@ -508,9 +529,12 @@ function attemptExplicitHydrationTarget(
           // We're blocked on hydrating this boundary.
           // Increase its priority.
           queuedTarget.blockedOn = instance;
-          runWithPriority(queuedTarget.priority, () => {
-            attemptHydrationAtCurrentPriority(nearestMounted);
+          attemptHydrationAtPriority(queuedTarget.lanePriority, () => {
+            runWithPriority(queuedTarget.priority, () => {
+              attemptHydrationAtCurrentPriority(nearestMounted);
+            });
           });
+
           return;
         }
       } else if (tag === HostRoot) {
@@ -529,15 +553,17 @@ function attemptExplicitHydrationTarget(
 
 export function queueExplicitHydrationTarget(target: Node): void {
   if (enableSelectiveHydration) {
-    const priority = getCurrentPriorityLevel();
+    const schedulerPriority = getCurrentPriorityLevel();
+    const updateLanePriority = getCurrentUpdatePriority();
     const queuedTarget: QueuedHydrationTarget = {
       blockedOn: null,
       target: target,
-      priority: priority,
+      priority: schedulerPriority,
+      lanePriority: updateLanePriority,
     };
     let i = 0;
     for (; i < queuedExplicitHydrationTargets.length; i++) {
-      if (priority <= queuedExplicitHydrationTargets[i].priority) {
+      if (schedulerPriority <= queuedExplicitHydrationTargets[i].priority) {
         break;
       }
     }

packages/react-dom/src/events/plugins/__tests__/ModernSimpleEventPlugin-test.js

@@ -231,6 +231,8 @@ describe('SimpleEventPlugin', function() {
   describe('interactive events, in concurrent mode', () => {
     beforeEach(() => {
       jest.resetModules();
+
+      React = require('react');
       ReactDOM = require('react-dom');
       Scheduler = require('scheduler');
     });
@@ -377,11 +379,14 @@ describe('SimpleEventPlugin', function() {
             <button
               ref={el => (button = el)}
               onClick={() => {
-                Scheduler.unstable_next(() => {
-                  this.setState(state => ({
-                    lowPriCount: state.lowPriCount + 1,
-                  }));
-                });
+                React.unstable_withSuspenseConfig(
+                  () => {
+                    this.setState(state => ({
+                      lowPriCount: state.lowPriCount + 1,
+                    }));
+                  },
+                  {timeoutMs: 5000},
+                );
               }}>
               {text}
             </button>

packages/react-noop-renderer/src/ReactNoop.js

@@ -47,6 +47,9 @@ export const {
   act,
   dumpTree,
   getRoot,
+  // TODO: Remove this once callers migrate to alternatives.
+  // This should only be used by React internals.
+  unstable_runWithPriority,
 } = createReactNoop(
   ReactFiberReconciler, // reconciler
   true, // useMutation

packages/react-noop-renderer/src/ReactNoopPersistent.js

@@ -47,6 +47,9 @@ export const {
   act,
   dumpTree,
   getRoot,
+  // TODO: Remove this once callers migrate to alternatives.
+  // This should only be used by React internals.
+  unstable_runWithPriority,
 } = createReactNoop(
   ReactFiberReconciler, // reconciler
   false, // useMutation

packages/react-noop-renderer/src/createReactNoop.js

@@ -954,6 +954,8 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
       return Scheduler.unstable_flushExpired();
     },
 
+    unstable_runWithPriority: NoopRenderer.runWithPriority,
+
     batchedUpdates: NoopRenderer.batchedUpdates,
 
     deferredUpdates: NoopRenderer.deferredUpdates,