Decouple update priority tracking from Scheduler package

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,

packages/react-reconciler/src/ReactFiberHooks.new.js

@@ -30,11 +30,16 @@ import {NoMode, BlockingMode} from './ReactTypeOfMode';
 import {
   NoLane,
   NoLanes,
+  InputContinuousLanePriority,
   isSubsetOfLanes,
   mergeLanes,
   removeLanes,
   markRootEntangled,
   markRootMutableRead,
+  getCurrentUpdateLanePriority,
+  setCurrentUpdateLanePriority,
+  higherLanePriority,
+  DefaultLanePriority,
 } from './ReactFiberLane';
 import {readContext} from './ReactFiberNewContext.new';
 import {createDeprecatedResponderListener} from './ReactFiberDeprecatedEvents.new';
@@ -1498,12 +1503,20 @@ function rerenderDeferredValue<T>(
 
 function startTransition(setPending, config, callback) {
   const priorityLevel = getCurrentPriorityLevel();
+  const previousLanePriority = getCurrentUpdateLanePriority();
+  setCurrentUpdateLanePriority(
+    higherLanePriority(previousLanePriority, InputContinuousLanePriority),
+  );
   runWithPriority(
     priorityLevel < UserBlockingPriority ? UserBlockingPriority : priorityLevel,
     () => {
       setPending(true);
     },
   );
+
+  // If there's no SuspenseConfig set, we'll use the DefaultLanePriority for this transition.
+  setCurrentUpdateLanePriority(DefaultLanePriority);
+
   runWithPriority(
     priorityLevel > NormalPriority ? NormalPriority : priorityLevel,
     () => {
@@ -1513,6 +1526,7 @@ function startTransition(setPending, config, callback) {
         setPending(false);
         callback();
       } finally {
+        setCurrentUpdateLanePriority(previousLanePriority);
         ReactCurrentBatchConfig.suspense = previousConfig;
       }
     },

packages/react-reconciler/src/ReactFiberHooks.old.js

@@ -33,11 +33,16 @@ import {NoMode, BlockingMode, DebugTracingMode} from './ReactTypeOfMode';
 import {
   NoLane,
   NoLanes,
+  InputContinuousLanePriority,
   isSubsetOfLanes,
   mergeLanes,
   removeLanes,
   markRootEntangled,
   markRootMutableRead,
+  getCurrentUpdateLanePriority,
+  setCurrentUpdateLanePriority,
+  higherLanePriority,
+  DefaultLanePriority,
 } from './ReactFiberLane';
 import {readContext} from './ReactFiberNewContext.old';
 import {createDeprecatedResponderListener} from './ReactFiberDeprecatedEvents.old';
@@ -1502,12 +1507,20 @@ function rerenderDeferredValue<T>(
 
 function startTransition(setPending, config, callback) {
   const priorityLevel = getCurrentPriorityLevel();
+  const previousLanePriority = getCurrentUpdateLanePriority();
+  setCurrentUpdateLanePriority(
+    higherLanePriority(previousLanePriority, InputContinuousLanePriority),
+  );
   runWithPriority(
     priorityLevel < UserBlockingPriority ? UserBlockingPriority : priorityLevel,
     () => {
       setPending(true);
     },
   );
+
+  // If there's no SuspenseConfig set, we'll use the DefaultLanePriority for this transition.
+  setCurrentUpdateLanePriority(DefaultLanePriority);
+
   runWithPriority(
     priorityLevel > NormalPriority ? NormalPriority : priorityLevel,
     () => {
@@ -1517,6 +1530,7 @@ function startTransition(setPending, config, callback) {
         setPending(false);
         callback();
       } finally {
+        setCurrentUpdateLanePriority(previousLanePriority);
         ReactCurrentBatchConfig.suspense = previousConfig;
       }
     },

packages/react-reconciler/src/ReactFiberLane.js

@@ -49,10 +49,10 @@ const InputDiscreteHydrationLanePriority: LanePriority = 14;
 export const InputDiscreteLanePriority: LanePriority = 13;
 
 const InputContinuousHydrationLanePriority: LanePriority = 12;
-const InputContinuousLanePriority: LanePriority = 11;
+export const InputContinuousLanePriority: LanePriority = 11;
 
 const DefaultHydrationLanePriority: LanePriority = 10;
-const DefaultLanePriority: LanePriority = 9;
+export const DefaultLanePriority: LanePriority = 9;
 
 const TransitionShortHydrationLanePriority: LanePriority = 8;
 export const TransitionShortLanePriority: LanePriority = 7;
@@ -120,6 +120,16 @@ export const OffscreenLane: Lane = /*                   */ 0b1000000000000000000
 
 export const NoTimestamp = -1;
 
+let currentUpdateLanePriority: LanePriority = NoLanePriority;
+
+export function getCurrentUpdateLanePriority(): LanePriority {
+  return currentUpdateLanePriority;
+}
+
+export function setCurrentUpdateLanePriority(newLanePriority: LanePriority) {
+  currentUpdateLanePriority = newLanePriority;
+}
+
 // "Registers" used to "return" multiple values
 // Used by getHighestPriorityLanes and getNextLanes:
 let return_highestLanePriority: LanePriority = DefaultLanePriority;
@@ -651,6 +661,13 @@ export function higherPriorityLane(a: Lane, b: Lane) {
   return a !== NoLane && a < b ? a : b;
 }
 
+export function higherLanePriority(
+  a: LanePriority,
+  b: LanePriority,
+): LanePriority {
+  return a !== NoLanePriority && a > b ? a : b;
+}
+
 export function createLaneMap<T>(initial: T): LaneMap<T> {
   return new Array(TotalLanes).fill(initial);
 }

packages/react-reconciler/src/ReactFiberReconciler.js

@@ -50,6 +50,8 @@ import {
   focusWithin as focusWithin_old,
   observeVisibleRects as observeVisibleRects_old,
   registerMutableSourceForHydration as registerMutableSourceForHydration_old,
+  runWithPriority as runWithPriority_old,
+  getCurrentUpdatePriority as getCurrentUpdatePriority_old,
 } from './ReactFiberReconciler.old';
 
 import {
@@ -88,6 +90,8 @@ import {
   focusWithin as focusWithin_new,
   observeVisibleRects as observeVisibleRects_new,
   registerMutableSourceForHydration as registerMutableSourceForHydration_new,
+  runWithPriority as runWithPriority_new,
+  getCurrentUpdatePriority as getCurrentUpdatePriority_new,
 } from './ReactFiberReconciler.new';
 
 export const createContainer = enableNewReconciler
@@ -139,6 +143,9 @@ export const attemptContinuousHydration = enableNewReconciler
 export const attemptHydrationAtCurrentPriority = enableNewReconciler
   ? attemptHydrationAtCurrentPriority_new
   : attemptHydrationAtCurrentPriority_old;
+export const getCurrentUpdatePriority = enableNewReconciler
+  ? getCurrentUpdatePriority_new
+  : getCurrentUpdatePriority_old;
 export const findHostInstance = enableNewReconciler
   ? findHostInstance_new
   : findHostInstance_old;
@@ -194,3 +201,6 @@ export const observeVisibleRects = enableNewReconciler
 export const registerMutableSourceForHydration = enableNewReconciler
   ? registerMutableSourceForHydration_new
   : registerMutableSourceForHydration_old;
+export const runWithPriority = enableNewReconciler
+  ? runWithPriority_new
+  : runWithPriority_old;