Support nesting of startTransition and flushSync

acdlite · acdlite · commit f2a4e8c1da97 · 2021-03-29T22:37:55.000-05:00
Currently we track event priorities in two different places:

- For transitions, we use the ReactCurrentBatchConfig object, which
  lives in the secret internals object of the isomorphic package. This
  is how the hook-less `startTransition` works.
- For other types of events, we use a module-level variable in
  ReactEventPriorities, which is owned by the renderer.

The problem with this approach is if both priorities are set, we don't
know which one takes precedence. For example, what happens if wrap an
update in both `ReactDOM.flushSync` and `React.startTransition`?

The desired behavior is that we use the priority of whichever wrapper
is closer on the stack. So if `flushSync` is called inside of
`startTransition`, then we should use sync priority; if it's the other
way around, we should use transition priority.

To implement this, I updated the ReactEventPriorities module to track
priority on the ReactCurrentBatchConfig object instead of in a module-
level variable. Now when multiple event wrappers are nested, the inner-
most wrapper wins.
diff --git a/packages/react-reconciler/src/ReactEventPriorities.new.js b/packages/react-reconciler/src/ReactEventPriorities.new.js
@@ -7,10 +7,11 @@
  * @flow
  */
 
-import type {Lane, Lanes} from './ReactFiberLane.new';
+import type {Lanes} from './ReactFiberLane.new';
+
+import ReactSharedInternals from 'shared/ReactSharedInternals';
 
 import {
-  NoLane,
   SyncLane,
   InputContinuousLane,
   DefaultLane,
@@ -19,30 +20,32 @@ import {
   includesNonIdleWork,
 } from './ReactFiberLane.new';
 
-export opaque type EventPriority = Lane;
+const {ReactCurrentBatchConfig} = ReactSharedInternals;
+
+export opaque type EventPriority = number;
 
 export const DiscreteEventPriority: EventPriority = SyncLane;
 export const ContinuousEventPriority: EventPriority = InputContinuousLane;
 export const DefaultEventPriority: EventPriority = DefaultLane;
 export const IdleEventPriority: EventPriority = IdleLane;
 
-let currentUpdatePriority: EventPriority = NoLane;
+export const TransitionEventPriority: EventPriority = 3;
 
 export function getCurrentUpdatePriority(): EventPriority {
-  return currentUpdatePriority;
+  return ReactCurrentBatchConfig.transition;
 }
 
 export function setCurrentUpdatePriority(newPriority: EventPriority) {
-  currentUpdatePriority = newPriority;
+  ReactCurrentBatchConfig.transition = newPriority;
 }
 
 export function runWithPriority<T>(priority: EventPriority, fn: () => T): T {
-  const previousPriority = currentUpdatePriority;
+  const previousPriority = ReactCurrentBatchConfig.transition;
   try {
-    currentUpdatePriority = priority;
+    ReactCurrentBatchConfig.transition = priority;
     return fn();
   } finally {
-    currentUpdatePriority = previousPriority;
+    ReactCurrentBatchConfig.transition = previousPriority;
   }
 }
 
diff --git a/packages/react-reconciler/src/ReactEventPriorities.old.js b/packages/react-reconciler/src/ReactEventPriorities.old.js
@@ -7,10 +7,11 @@
  * @flow
  */
 
-import type {Lane, Lanes} from './ReactFiberLane.old';
+import type {Lanes} from './ReactFiberLane.old';
+
+import ReactSharedInternals from 'shared/ReactSharedInternals';
 
 import {
-  NoLane,
   SyncLane,
   InputContinuousLane,
   DefaultLane,
@@ -19,30 +20,32 @@ import {
   includesNonIdleWork,
 } from './ReactFiberLane.old';
 
-export opaque type EventPriority = Lane;
+const {ReactCurrentBatchConfig} = ReactSharedInternals;
+
+export opaque type EventPriority = number;
 
 export const DiscreteEventPriority: EventPriority = SyncLane;
 export const ContinuousEventPriority: EventPriority = InputContinuousLane;
 export const DefaultEventPriority: EventPriority = DefaultLane;
 export const IdleEventPriority: EventPriority = IdleLane;
 
-let currentUpdatePriority: EventPriority = NoLane;
+export const TransitionEventPriority: EventPriority = 3;
 
 export function getCurrentUpdatePriority(): EventPriority {
-  return currentUpdatePriority;
+  return ReactCurrentBatchConfig.transition;
 }
 
 export function setCurrentUpdatePriority(newPriority: EventPriority) {
-  currentUpdatePriority = newPriority;
+  ReactCurrentBatchConfig.transition = newPriority;
 }
 
 export function runWithPriority<T>(priority: EventPriority, fn: () => T): T {
-  const previousPriority = currentUpdatePriority;
+  const previousPriority = ReactCurrentBatchConfig.transition;
   try {
-    currentUpdatePriority = priority;
+    ReactCurrentBatchConfig.transition = priority;
     return fn();
   } finally {
-    currentUpdatePriority = previousPriority;
+    ReactCurrentBatchConfig.transition = previousPriority;
   }
 }
 
diff --git a/packages/react-reconciler/src/ReactFiberHooks.new.js b/packages/react-reconciler/src/ReactFiberHooks.new.js
@@ -50,6 +50,7 @@ import {
 } from './ReactFiberLane.new';
 import {
   ContinuousEventPriority,
+  TransitionEventPriority,
   getCurrentUpdatePriority,
   setCurrentUpdatePriority,
   higherEventPriority,
@@ -1664,8 +1665,9 @@ function updateMemo<T>(
 function mountDeferredValue<T>(value: T): T {
   const [prevValue, setValue] = mountState(value);
   mountEffect(() => {
+    // TODO: Replace with setCurrentUpdatePriority
     const prevTransition = ReactCurrentBatchConfig.transition;
-    ReactCurrentBatchConfig.transition = 1;
+    ReactCurrentBatchConfig.transition = TransitionEventPriority;
     try {
       setValue(value);
     } finally {
@@ -1678,8 +1680,9 @@ function mountDeferredValue<T>(value: T): T {
 function updateDeferredValue<T>(value: T): T {
   const [prevValue, setValue] = updateState(value);
   updateEffect(() => {
+    // TODO: Replace with setCurrentUpdatePriority
     const prevTransition = ReactCurrentBatchConfig.transition;
-    ReactCurrentBatchConfig.transition = 1;
+    ReactCurrentBatchConfig.transition = TransitionEventPriority;
     try {
       setValue(value);
     } finally {
@@ -1693,7 +1696,7 @@ function rerenderDeferredValue<T>(value: T): T {
   const [prevValue, setValue] = rerenderState(value);
   updateEffect(() => {
     const prevTransition = ReactCurrentBatchConfig.transition;
-    ReactCurrentBatchConfig.transition = 1;
+    ReactCurrentBatchConfig.transition = TransitionEventPriority;
     try {
       setValue(value);
     } finally {
@@ -1711,8 +1714,9 @@ function startTransition(setPending, callback) {
 
   setPending(true);
 
+  // TODO: Replace with setCurrentUpdatePriority
   const prevTransition = ReactCurrentBatchConfig.transition;
-  ReactCurrentBatchConfig.transition = 1;
+  ReactCurrentBatchConfig.transition = TransitionEventPriority;
   try {
     setPending(false);
     callback();
diff --git a/packages/react-reconciler/src/ReactFiberHooks.old.js b/packages/react-reconciler/src/ReactFiberHooks.old.js
@@ -50,6 +50,7 @@ import {
 } from './ReactFiberLane.old';
 import {
   ContinuousEventPriority,
+  TransitionEventPriority,
   getCurrentUpdatePriority,
   setCurrentUpdatePriority,
   higherEventPriority,
@@ -1664,8 +1665,9 @@ function updateMemo<T>(
 function mountDeferredValue<T>(value: T): T {
   const [prevValue, setValue] = mountState(value);
   mountEffect(() => {
+    // TODO: Replace with setCurrentUpdatePriority
     const prevTransition = ReactCurrentBatchConfig.transition;
-    ReactCurrentBatchConfig.transition = 1;
+    ReactCurrentBatchConfig.transition = TransitionEventPriority;
     try {
       setValue(value);
     } finally {
@@ -1678,8 +1680,9 @@ function mountDeferredValue<T>(value: T): T {
 function updateDeferredValue<T>(value: T): T {
   const [prevValue, setValue] = updateState(value);
   updateEffect(() => {
+    // TODO: Replace with setCurrentUpdatePriority
     const prevTransition = ReactCurrentBatchConfig.transition;
-    ReactCurrentBatchConfig.transition = 1;
+    ReactCurrentBatchConfig.transition = TransitionEventPriority;
     try {
       setValue(value);
     } finally {
@@ -1693,7 +1696,7 @@ function rerenderDeferredValue<T>(value: T): T {
   const [prevValue, setValue] = rerenderState(value);
   updateEffect(() => {
     const prevTransition = ReactCurrentBatchConfig.transition;
-    ReactCurrentBatchConfig.transition = 1;
+    ReactCurrentBatchConfig.transition = TransitionEventPriority;
     try {
       setValue(value);
     } finally {
@@ -1711,8 +1714,9 @@ function startTransition(setPending, callback) {
 
   setPending(true);
 
+  // TODO: Replace with setCurrentUpdatePriority
   const prevTransition = ReactCurrentBatchConfig.transition;
-  ReactCurrentBatchConfig.transition = 1;
+  ReactCurrentBatchConfig.transition = TransitionEventPriority;
   try {
     setPending(false);
     callback();
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.new.js b/packages/react-reconciler/src/ReactFiberWorkLoop.new.js
@@ -10,6 +10,7 @@
 import type {Thenable, Wakeable} from 'shared/ReactTypes';
 import type {Fiber, FiberRoot} from './ReactInternalTypes';
 import type {Lanes, Lane} from './ReactFiberLane.new';
+import type {EventPriority} from './ReactEventPriorities.new';
 import type {Interaction} from 'scheduler/src/Tracing';
 import type {SuspenseState} from './ReactFiberSuspenseComponent.new';
 import type {StackCursor} from './ReactFiberStack.new';
@@ -170,7 +171,6 @@ import {
   higherEventPriority,
   lanesToEventPriority,
 } from './ReactEventPriorities.new';
-import {requestCurrentTransition, NoTransition} from './ReactFiberTransition';
 import {beginWork as originalBeginWork} from './ReactFiberBeginWork.new';
 import {completeWork} from './ReactFiberCompleteWork.new';
 import {unwindWork, unwindInterruptedWork} from './ReactFiberUnwindWork.new';
@@ -401,7 +401,19 @@ export function requestUpdateLane(fiber: Fiber): Lane {
     return pickArbitraryLane(workInProgressRootRenderLanes);
   }
 
-  const isTransition = requestCurrentTransition() !== NoTransition;
+  const eventPriority: EventPriority = (getCurrentUpdatePriority(): any);
+
+  // Check if the event priority is a lane, by checking if it's a power of 2. If
+  // it's not, then we assume it's TransitionPriority. This is a trick to avoid
+  // leaking the Lane type into the shared reconciler package. Consumers should
+  // be versioning the isomorphic package in lockstep with the renderer — but
+  // for those rare exceptions where there's a mismatch (like, say, a
+  // multi-renderer app), this is less likely to break than if we used an
+  // arbitrary lane bit that later changes meaning.
+  // TODO: Move this type conversion to the event priority module.
+  const maybePowerOfTwo = (eventPriority: any);
+  const isTransition =
+    maybePowerOfTwo !== 0 && (maybePowerOfTwo & (maybePowerOfTwo - 1)) !== 0;
   if (isTransition) {
     // The algorithm for assigning an update to a lane should be stable for all
     // updates at the same priority within the same event. To do this, the
@@ -417,13 +429,13 @@ export function requestUpdateLane(fiber: Fiber): Lane {
     return currentEventTransitionLane;
   }
 
+  // If this isn't a transition, the opaque type returned by the host config is
+  // internally a lane, so we can use that directly.
+  // TODO: Move this type conversion to the event priority module.
+  const updateLane: Lane = (eventPriority: any);
+
   // Updates originating inside certain React methods, like flushSync, have
   // their priority set by tracking it with a context variable.
-  //
-  // The opaque type returned by the host config is internally a lane, so we can
-  // use that directly.
-  // TODO: Move this type conversion to the event priority module.
-  const updateLane: Lane = (getCurrentUpdatePriority(): any);
   if (updateLane !== NoLane) {
     return updateLane;
   }
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.old.js b/packages/react-reconciler/src/ReactFiberWorkLoop.old.js
@@ -10,6 +10,7 @@
 import type {Thenable, Wakeable} from 'shared/ReactTypes';
 import type {Fiber, FiberRoot} from './ReactInternalTypes';
 import type {Lanes, Lane} from './ReactFiberLane.old';
+import type {EventPriority} from './ReactEventPriorities.old';
 import type {Interaction} from 'scheduler/src/Tracing';
 import type {SuspenseState} from './ReactFiberSuspenseComponent.old';
 import type {StackCursor} from './ReactFiberStack.old';
@@ -170,7 +171,6 @@ import {
   higherEventPriority,
   lanesToEventPriority,
 } from './ReactEventPriorities.old';
-import {requestCurrentTransition, NoTransition} from './ReactFiberTransition';
 import {beginWork as originalBeginWork} from './ReactFiberBeginWork.old';
 import {completeWork} from './ReactFiberCompleteWork.old';
 import {unwindWork, unwindInterruptedWork} from './ReactFiberUnwindWork.old';
@@ -401,7 +401,19 @@ export function requestUpdateLane(fiber: Fiber): Lane {
     return pickArbitraryLane(workInProgressRootRenderLanes);
   }
 
-  const isTransition = requestCurrentTransition() !== NoTransition;
+  const eventPriority: EventPriority = (getCurrentUpdatePriority(): any);
+
+  // Check if the event priority is a lane, by checking if it's a power of 2. If
+  // it's not, then we assume it's TransitionPriority. This is a trick to avoid
+  // leaking the Lane type into the shared reconciler package. Consumers should
+  // be versioning the isomorphic package in lockstep with the renderer — but
+  // for those rare exceptions where there's a mismatch (like, say, a
+  // multi-renderer app), this is less likely to break than if we used an
+  // arbitrary lane bit that later changes meaning.
+  // TODO: Move this type conversion to the event priority module.
+  const maybePowerOfTwo = (eventPriority: any);
+  const isTransition =
+    maybePowerOfTwo !== 0 && (maybePowerOfTwo & (maybePowerOfTwo - 1)) !== 0;
   if (isTransition) {
     // The algorithm for assigning an update to a lane should be stable for all
     // updates at the same priority within the same event. To do this, the
@@ -417,13 +429,13 @@ export function requestUpdateLane(fiber: Fiber): Lane {
     return currentEventTransitionLane;
   }
 
+  // If this isn't a transition, the opaque type returned by the host config is
+  // internally a lane, so we can use that directly.
+  // TODO: Move this type conversion to the event priority module.
+  const updateLane: Lane = (eventPriority: any);
+
   // Updates originating inside certain React methods, like flushSync, have
   // their priority set by tracking it with a context variable.
-  //
-  // The opaque type returned by the host config is internally a lane, so we can
-  // use that directly.
-  // TODO: Move this type conversion to the event priority module.
-  const updateLane: Lane = (getCurrentUpdatePriority(): any);
   if (updateLane !== NoLane) {
     return updateLane;
   }
diff --git a/packages/react-reconciler/src/__tests__/ReactFlushSync-test.js b/packages/react-reconciler/src/__tests__/ReactFlushSync-test.js
@@ -3,6 +3,7 @@ let ReactNoop;
 let Scheduler;
 let useState;
 let useEffect;
+let startTransition;
 
 describe('ReactFlushSync', () => {
   beforeEach(() => {
@@ -13,6 +14,7 @@ describe('ReactFlushSync', () => {
     Scheduler = require('scheduler');
     useState = React.useState;
     useEffect = React.useEffect;
+    startTransition = React.startTransition;
   });
 
   function Text({text}) {
@@ -54,4 +56,45 @@ describe('ReactFlushSync', () => {
     });
     expect(root).toMatchRenderedOutput('1, 1');
   });
+
+  // @gate experimental
+  test('nested with startTransition', async () => {
+    let setSyncState;
+    let setState;
+    function App() {
+      const [syncState, _setSyncState] = useState(0);
+      const [state, _setState] = useState(0);
+      setSyncState = _setSyncState;
+      setState = _setState;
+      return <Text text={`${syncState}, ${state}`} />;
+    }
+
+    const root = ReactNoop.createRoot();
+    await ReactNoop.act(async () => {
+      root.render(<App />);
+    });
+    expect(Scheduler).toHaveYielded(['0, 0']);
+    expect(root).toMatchRenderedOutput('0, 0');
+
+    await ReactNoop.act(async () => {
+      ReactNoop.flushSync(() => {
+        startTransition(() => {
+          // This should be async even though flushSync is on the stack, because
+          // startTransition is closer.
+          setState(1);
+          ReactNoop.flushSync(() => {
+            // This should be async even though startTransition is on the stack,
+            // because flushSync is closer.
+            setSyncState(1);
+          });
+        });
+      });
+      // Only the sync update should have flushed
+      expect(Scheduler).toHaveYielded(['1, 0']);
+      expect(root).toMatchRenderedOutput('1, 0');
+    });
+    // Now the async update has flushed, too.
+    expect(Scheduler).toHaveYielded(['1, 1']);
+    expect(root).toMatchRenderedOutput('1, 1');
+  });
 });
diff --git a/packages/react/src/ReactCurrentBatchConfig.js b/packages/react/src/ReactCurrentBatchConfig.js
@@ -12,6 +12,8 @@
  * should suspend for if it needs to.
  */
 const ReactCurrentBatchConfig = {
+  // TODO: This now is used to track other types of event priorities, too, not
+  // just transitions. Consider renaming.
   transition: (0: number),
 };
 
diff --git a/packages/react/src/ReactStartTransition.js b/packages/react/src/ReactStartTransition.js
