Synchronously flush the transition lane scheduled in a popstate event (#26025)

tyao1 · web-flow · commit d121c67004a2 · 2023-04-13T15:21:19.000-04:00
## Summary Browsers restore state like forms and scroll position right after the popstate event. To make sure the page work as expected on back or forward button, we need to flush transitions scheduled in a popstate synchronously, and only yields if it suspends. This PR adds a new HostConfig method to check if `window.event === 'popstate'`, and `scheduleMicrotask` if a transition is scheduled in a `PopStateEvent`. ## How did you test this change? yarn test
diff --git a/packages/react-art/src/ReactFiberConfigART.js b/packages/react-art/src/ReactFiberConfigART.js
@@ -346,6 +346,10 @@ export function getCurrentEventPriority() {
   return DefaultEventPriority;
 }
 
+export function shouldAttemptEagerTransition() {
+  return false;
+}
+
 // The ART renderer is secondary to the React DOM renderer.
 export const isPrimaryRenderer = false;
 
diff --git a/packages/react-dom-bindings/src/client/ReactFiberConfigDOM.js b/packages/react-dom-bindings/src/client/ReactFiberConfigDOM.js
@@ -527,6 +527,10 @@ export function getCurrentEventPriority(): EventPriority {
   return getEventPriority(currentEvent.type);
 }
 
+export function shouldAttemptEagerTransition(): boolean {
+  return window.event && window.event.type === 'popstate';
+}
+
 export const isPrimaryRenderer = true;
 export const warnsIfNotActing = true;
 // This initialization code may run even on server environments
diff --git a/packages/react-dom/src/__tests__/ReactDOMFiberAsync-test.js b/packages/react-dom/src/__tests__/ReactDOMFiberAsync-test.js
@@ -27,7 +27,6 @@ describe('ReactDOMFiberAsync', () => {
   let container;
 
   beforeEach(() => {
-    jest.resetModules();
     container = document.createElement('div');
     React = require('react');
     ReactDOM = require('react-dom');
@@ -40,6 +39,7 @@ describe('ReactDOMFiberAsync', () => {
     assertLog = InternalTestUtils.assertLog;
 
     document.body.appendChild(container);
+    window.event = undefined;
   });
 
   afterEach(() => {
@@ -566,4 +566,139 @@ describe('ReactDOMFiberAsync', () => {
 
     expect(container.textContent).toBe('new');
   });
+
+  it('should synchronously render the transition lane scheduled in a popState', async () => {
+    function App() {
+      const [syncState, setSyncState] = React.useState(false);
+      const [hasNavigated, setHasNavigated] = React.useState(false);
+      function onPopstate() {
+        Scheduler.log(`popState`);
+        React.startTransition(() => {
+          setHasNavigated(true);
+        });
+        setSyncState(true);
+      }
+      React.useEffect(() => {
+        window.addEventListener('popstate', onPopstate);
+        return () => {
+          window.removeEventListener('popstate', onPopstate);
+        };
+      }, []);
+      Scheduler.log(`render:${hasNavigated}/${syncState}`);
+      return null;
+    }
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(<App />);
+    });
+    assertLog(['render:false/false']);
+
+    await act(async () => {
+      const popStateEvent = new Event('popstate');
+      // Jest is not emulating window.event correctly in the microtask
+      window.event = popStateEvent;
+      window.dispatchEvent(popStateEvent);
+      queueMicrotask(() => {
+        window.event = undefined;
+      });
+    });
+
+    assertLog(['popState', 'render:true/true']);
+    await act(() => {
+      root.unmount();
+    });
+  });
+
+  it('Should not flush transition lanes if there is no transition scheduled in popState', async () => {
+    let setHasNavigated;
+    function App() {
+      const [syncState, setSyncState] = React.useState(false);
+      const [hasNavigated, _setHasNavigated] = React.useState(false);
+      setHasNavigated = _setHasNavigated;
+      function onPopstate() {
+        setSyncState(true);
+      }
+
+      React.useEffect(() => {
+        window.addEventListener('popstate', onPopstate);
+        return () => {
+          window.removeEventListener('popstate', onPopstate);
+        };
+      }, []);
+
+      Scheduler.log(`render:${hasNavigated}/${syncState}`);
+      return null;
+    }
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(<App />);
+    });
+    assertLog(['render:false/false']);
+
+    React.startTransition(() => {
+      setHasNavigated(true);
+    });
+    await act(async () => {
+      const popStateEvent = new Event('popstate');
+      // Jest is not emulating window.event correctly in the microtask
+      window.event = popStateEvent;
+      window.dispatchEvent(popStateEvent);
+      queueMicrotask(() => {
+        window.event = undefined;
+      });
+    });
+    assertLog(['render:false/true', 'render:true/true']);
+    await act(() => {
+      root.unmount();
+    });
+  });
+
+  it('transition lane in popState should yield if it suspends', async () => {
+    const never = {then() {}};
+    let _setText;
+
+    function App() {
+      const [shouldSuspend, setShouldSuspend] = React.useState(false);
+      const [text, setText] = React.useState('0');
+      _setText = setText;
+      if (shouldSuspend) {
+        Scheduler.log('Suspend!');
+        throw never;
+      }
+      function onPopstate() {
+        React.startTransition(() => {
+          setShouldSuspend(val => !val);
+        });
+      }
+      React.useEffect(() => {
+        window.addEventListener('popstate', onPopstate);
+        return () => window.removeEventListener('popstate', onPopstate);
+      }, []);
+      Scheduler.log(`Child:${shouldSuspend}/${text}`);
+      return text;
+    }
+
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(<App />);
+    });
+    assertLog(['Child:false/0']);
+
+    await act(() => {
+      const popStateEvent = new Event('popstate');
+      window.event = popStateEvent;
+      window.dispatchEvent(popStateEvent);
+      queueMicrotask(() => {
+        window.event = undefined;
+      });
+    });
+    assertLog(['Suspend!']);
+
+    await act(async () => {
+      _setText('1');
+    });
+    assertLog(['Child:false/1', 'Suspend!']);
+
+    root.unmount();
+  });
 });
diff --git a/packages/react-native-renderer/src/ReactFiberConfigFabric.js b/packages/react-native-renderer/src/ReactFiberConfigFabric.js
@@ -334,6 +334,10 @@ export function getCurrentEventPriority(): * {
   return DefaultEventPriority;
 }
 
+export function shouldAttemptEagerTransition(): boolean {
+  return false;
+}
+
 // The Fabric renderer is secondary to the existing React Native renderer.
 export const isPrimaryRenderer = false;
 
diff --git a/packages/react-native-renderer/src/ReactFiberConfigNative.js b/packages/react-native-renderer/src/ReactFiberConfigNative.js
@@ -265,6 +265,10 @@ export function getCurrentEventPriority(): * {
   return DefaultEventPriority;
 }
 
+export function shouldAttemptEagerTransition(): boolean {
+  return false;
+}
+
 // -------------------
 //     Mutation
 // -------------------
diff --git a/packages/react-noop-renderer/src/createReactNoop.js b/packages/react-noop-renderer/src/createReactNoop.js
@@ -526,6 +526,10 @@ function createReactNoop(reconciler: Function, useMutation: boolean) {
       return currentEventPriority;
     },
 
+    shouldAttemptEagerTransition(): boolean {
+      return false;
+    },
+
     now: Scheduler.unstable_now,
 
     isPrimaryRenderer: true,
diff --git a/packages/react-reconciler/src/ReactFiberRootScheduler.js b/packages/react-reconciler/src/ReactFiberRootScheduler.js
@@ -20,6 +20,8 @@ import {
   getNextLanes,
   includesSyncLane,
   markStarvedLanesAsExpired,
+  markRootEntangled,
+  mergeLanes,
 } from './ReactFiberLane';
 import {
   CommitContext,
@@ -49,7 +51,11 @@ import {
   IdleEventPriority,
   lanesToEventPriority,
 } from './ReactEventPriorities';
-import {supportsMicrotasks, scheduleMicrotask} from './ReactFiberConfig';
+import {
+  supportsMicrotasks,
+  scheduleMicrotask,
+  shouldAttemptEagerTransition,
+} from './ReactFiberConfig';
 
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 const {ReactCurrentActQueue} = ReactSharedInternals;
@@ -72,6 +78,8 @@ let mightHavePendingSyncWork: boolean = false;
 
 let isFlushingWork: boolean = false;
 
+let currentEventTransitionLane: Lane = NoLanes;
+
 export function ensureRootIsScheduled(root: FiberRoot): void {
   // This function is called whenever a root receives an update. It does two
   // things 1) it ensures the root is in the root schedule, and 2) it ensures
@@ -238,6 +246,14 @@ function processRootScheduleInMicrotask() {
   let root = firstScheduledRoot;
   while (root !== null) {
     const next = root.next;
+
+    if (
+      currentEventTransitionLane !== NoLane &&
+      shouldAttemptEagerTransition()
+    ) {
+      markRootEntangled(root, mergeLanes(currentEventTransitionLane, SyncLane));
+    }
+
     const nextLanes = scheduleTaskForRootDuringMicrotask(root, currentTime);
     if (nextLanes === NoLane) {
       // This root has no more pending work. Remove it from the schedule. To
@@ -267,6 +283,8 @@ function processRootScheduleInMicrotask() {
     root = next;
   }
 
+  currentEventTransitionLane = NoLane;
+
   // At the end of the microtask, flush any pending synchronous work. This has
   // to come at the end, because it does actual rendering work that might throw.
   flushSyncWorkOnAllRoots();
@@ -472,3 +490,11 @@ function scheduleImmediateTask(cb: () => mixed) {
     Scheduler_scheduleCallback(ImmediateSchedulerPriority, cb);
   }
 }
+
+export function getCurrentEventTransitionLane(): Lane {
+  return currentEventTransitionLane;
+}
+
+export function setCurrentEventTransitionLane(lane: Lane): void {
+  currentEventTransitionLane = lane;
+}
diff --git a/packages/react-reconciler/src/ReactFiberWorkLoop.js b/packages/react-reconciler/src/ReactFiberWorkLoop.js
@@ -278,6 +278,8 @@ import {
   flushSyncWorkOnAllRoots,
   flushSyncWorkOnLegacyRootsOnly,
   getContinuationForRoot,
+  getCurrentEventTransitionLane,
+  setCurrentEventTransitionLane,
 } from './ReactFiberRootScheduler';
 import {getMaskedContext, getUnmaskedContext} from './ReactFiberContext';
 
@@ -582,8 +584,6 @@ const NESTED_PASSIVE_UPDATE_LIMIT = 50;
 let nestedPassiveUpdateCount: number = 0;
 let rootWithPassiveNestedUpdates: FiberRoot | null = null;
 
-let currentEventTransitionLane: Lanes = NoLanes;
-
 let isRunningInsertionEffect = false;
 
 export function getWorkInProgressRoot(): FiberRoot | null {
@@ -640,11 +640,11 @@ export function requestUpdateLane(fiber: Fiber): Lane {
     // The trick we use is to cache the first of each of these inputs within an
     // event. Then reset the cached values once we can be sure the event is
     // over. Our heuristic for that is whenever we enter a concurrent work loop.
-    if (currentEventTransitionLane === NoLane) {
+    if (getCurrentEventTransitionLane() === NoLane) {
       // All transitions within the same event are assigned the same lane.
-      currentEventTransitionLane = claimNextTransitionLane();
+      setCurrentEventTransitionLane(claimNextTransitionLane());
     }
-    return currentEventTransitionLane;
+    return getCurrentEventTransitionLane();
   }
 
   // Updates originating inside certain React methods, like flushSync, have
@@ -848,8 +848,6 @@ export function performConcurrentWorkOnRoot(
     resetNestedUpdateFlag();
   }
 
-  currentEventTransitionLane = NoLanes;
-
   if ((executionContext & (RenderContext | CommitContext)) !== NoContext) {
     throw new Error('Should not already be working.');
   }
diff --git a/packages/react-reconciler/src/__tests__/ReactFiberHostContext-test.internal.js b/packages/react-reconciler/src/__tests__/ReactFiberHostContext-test.internal.js
@@ -70,6 +70,9 @@ describe('ReactFiberHostContext', () => {
       getCurrentEventPriority: function () {
         return DefaultEventPriority;
       },
+      shouldAttemptEagerTransition() {
+        return false;
+      },
       requestPostPaintCallback: function () {},
       maySuspendCommit(type, props) {
         return false;
diff --git a/packages/react-reconciler/src/forks/ReactFiberConfig.custom.js b/packages/react-reconciler/src/forks/ReactFiberConfig.custom.js
@@ -66,6 +66,8 @@ export const preparePortalMount = $$$config.preparePortalMount;
 export const prepareScopeUpdate = $$$config.prepareScopeUpdate;
 export const getInstanceFromScope = $$$config.getInstanceFromScope;
 export const getCurrentEventPriority = $$$config.getCurrentEventPriority;
+export const shouldAttemptEagerTransition =
+  $$$config.shouldAttemptEagerTransition;
 export const detachDeletedInstance = $$$config.detachDeletedInstance;
 export const requestPostPaintCallback = $$$config.requestPostPaintCallback;
 export const maySuspendCommit = $$$config.maySuspendCommit;
diff --git a/packages/react-test-renderer/src/ReactFiberConfigTestHost.js b/packages/react-test-renderer/src/ReactFiberConfigTestHost.js
@@ -213,6 +213,9 @@ export function createTextInstance(
 export function getCurrentEventPriority(): * {
   return DefaultEventPriority;
 }
+export function shouldAttemptEagerTransition(): boolean {
+  return false;
+}
 
 export const isPrimaryRenderer = false;
 export const warnsIfNotActing = true;
