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

<!--
  Thanks for submitting a pull request!
We appreciate you spending the time to work on these changes. Please
provide enough information so that others can review your pull request.
The three fields below are mandatory.

Before submitting a pull request, please make sure the following is
done:

1. Fork [the repository](https://github.com/facebook/react) and create
your branch from `main`.
  2. Run `yarn` in the repository root.
3. If you've fixed a bug or added code that should be tested, add tests!
4. Ensure the test suite passes (`yarn test`). Tip: `yarn test --watch
TestName` is helpful in development.
5. Run `yarn test --prod` to test in the production environment. It
supports the same options as `yarn test`.
6. If you need a debugger, run `yarn debug-test --watch TestName`, open
`chrome://inspect`, and press "Inspect".
7. Format your code with
[prettier](https://github.com/prettier/prettier) (`yarn prettier`).
8. Make sure your code lints (`yarn lint`). Tip: `yarn linc` to only
check changed files.
  9. Run the [Flow](https://flowtype.org/) type checks (`yarn flow`).
  10. If you haven't already, complete the CLA.

Learn more about contributing:
https://reactjs.org/docs/how-to-contribute.html
-->

## 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

Author: tyao1

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;
 

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

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();
+  });
 });

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;
 

packages/react-native-renderer/src/ReactFiberConfigNative.js

@@ -265,6 +265,10 @@ export function getCurrentEventPriority(): * {
   return DefaultEventPriority;
 }
 
+export function shouldAttemptEagerTransition(): boolean {
+  return false;
+}
+
 // -------------------
 //     Mutation
 // -------------------

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,

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;
+}

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.');
   }

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;

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;