From 6dac8deb3b5aa2a1eb8b2e32329051f8724ff3ca Mon Sep 17 00:00:00 2001
From: Matt Brophy <matt@brophy.org>
Date: Wed, 17 May 2023 16:13:44 -0400
Subject: [PATCH] Wrap internal router state updates with React.startTransition
 (#10438)

---
 .changeset/start-transition-support.md        |   6 +
 package.json                                  |   8 +-
 .../__tests__/components/LazyComponent.tsx    |   5 +
 .../concurrent-mode-navigations-test.tsx      | 398 ++++++++++++++++++
 packages/react-router-dom/index.tsx           |  38 +-
 .../react-router/__tests__/navigate-test.tsx  |   9 +-
 packages/react-router/lib/components.tsx      |  22 +-
 7 files changed, 467 insertions(+), 19 deletions(-)
 create mode 100644 .changeset/start-transition-support.md
 create mode 100644 packages/react-router-dom/__tests__/components/LazyComponent.tsx
 create mode 100644 packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx

diff --git a/.changeset/start-transition-support.md b/.changeset/start-transition-support.md
new file mode 100644
index 0000000000..e4f3d54a48
--- /dev/null
+++ b/.changeset/start-transition-support.md
@@ -0,0 +1,6 @@
+---
+"react-router": minor
+"react-router-dom": minor
+---
+
+Wrap internal router state updates with `React.startTransition` if it exists
diff --git a/package.json b/package.json
index 8b95a2d15f..5e79f3be06 100644
--- a/package.json
+++ b/package.json
@@ -108,16 +108,16 @@
       "none": "45 kB"
     },
     "packages/react-router/dist/react-router.production.min.js": {
-      "none": "13.3 kB"
+      "none": "13.4 kB"
     },
     "packages/react-router/dist/umd/react-router.production.min.js": {
-      "none": "15.6 kB"
+      "none": "15.8 kB"
     },
     "packages/react-router-dom/dist/react-router-dom.production.min.js": {
-      "none": "11.8 kB"
+      "none": "12.0 kB"
     },
     "packages/react-router-dom/dist/umd/react-router-dom.production.min.js": {
-      "none": "17.7 kB"
+      "none": "17.9 kB"
     }
   }
 }
diff --git a/packages/react-router-dom/__tests__/components/LazyComponent.tsx b/packages/react-router-dom/__tests__/components/LazyComponent.tsx
new file mode 100644
index 0000000000..9edc993d5c
--- /dev/null
+++ b/packages/react-router-dom/__tests__/components/LazyComponent.tsx
@@ -0,0 +1,5 @@
+import * as React from "react";
+
+export default function LazyComponent() {
+  return <h1>Lazy</h1>;
+}
diff --git a/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx b/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx
new file mode 100644
index 0000000000..6e7839a012
--- /dev/null
+++ b/packages/react-router-dom/__tests__/concurrent-mode-navigations-test.tsx
@@ -0,0 +1,398 @@
+import * as React from "react";
+import {
+  MemoryRouter,
+  Routes,
+  Route,
+  useNavigate,
+  BrowserRouter,
+  HashRouter,
+  createMemoryRouter,
+  createRoutesFromElements,
+  RouterProvider,
+} from "react-router-dom";
+import {
+  act,
+  fireEvent,
+  prettyDOM,
+  render,
+  screen,
+  waitFor,
+} from "@testing-library/react";
+import { JSDOM } from "jsdom";
+import LazyComponent from "./components//LazyComponent";
+
+describe("Handles concurrent mode features during navigations", () => {
+  function getComponents() {
+    function Home() {
+      let navigate = useNavigate();
+      return (
+        <>
+          <h1>Home</h1>
+          <button onClick={() => navigate("/about")}>/about</button>
+          <button onClick={() => navigate("/lazy")}>/lazy</button>
+        </>
+      );
+    }
+
+    let resolved = false;
+    let dfd = createDeferred();
+    function resolve() {
+      resolved = true;
+      return dfd.resolve();
+    }
+
+    function About() {
+      let navigate = useNavigate();
+      if (!resolved) {
+        throw dfd.promise;
+      }
+      return (
+        <>
+          <h1>About</h1>
+          <button onClick={() => navigate(-1)}>back</button>
+        </>
+      );
+    }
+
+    let lazyDfd = createDeferred();
+
+    const LazyComponent = React.lazy(async () => {
+      await lazyDfd.promise;
+      return import("./components/LazyComponent");
+    });
+
+    return {
+      Home,
+      About,
+      LazyComponent,
+      resolve,
+      resolveLazy: lazyDfd.resolve,
+    };
+  }
+
+  describe("when the destination route suspends with a boundary", () => {
+    async function assertNavigation(
+      container: HTMLElement,
+      resolve: () => void,
+      resolveLazy: () => void
+    ) {
+      // Start on home
+      expect(getHtml(container)).toMatch("Home");
+
+      // Click to /about and should see Suspense boundary
+      await act(() => {
+        fireEvent.click(screen.getByText("/about"));
+      });
+      await waitFor(() => screen.getByText("Loading..."));
+      expect(getHtml(container)).toMatch("Loading...");
+
+      // Resolve the destination UI to clear the boundary
+      await act(() => resolve());
+      await waitFor(() => screen.getByText("About"));
+      expect(getHtml(container)).toMatch("About");
+
+      // Back to home
+      await act(() => {
+        fireEvent.click(screen.getByText("back"));
+      });
+      await waitFor(() => screen.getByText("Home"));
+      expect(getHtml(container)).toMatch("Home");
+
+      // Click to /lazy and should see Suspense boundary
+      await act(() => {
+        fireEvent.click(screen.getByText("/lazy"));
+      });
+      await waitFor(() => screen.getByText("Loading Lazy Component..."));
+      expect(getHtml(container)).toMatch("Loading Lazy Component...");
+
+      // Resolve the lazy component to clear the boundary
+      await act(() => resolveLazy());
+      await waitFor(() => screen.getByText("Lazy"));
+      expect(getHtml(container)).toMatch("Lazy");
+    }
+
+    // eslint-disable-next-line jest/expect-expect
+    it("MemoryRouter", async () => {
+      let { Home, About, LazyComponent, resolve, resolveLazy } =
+        getComponents();
+
+      let { container } = render(
+        <MemoryRouter>
+          <Routes>
+            <Route path="/" element={<Home />} />
+            <Route
+              path="/about"
+              element={
+                <React.Suspense fallback={<p>Loading...</p>}>
+                  <About />
+                </React.Suspense>
+              }
+            />
+            <Route
+              path="/lazy"
+              element={
+                <React.Suspense fallback={<p>Loading Lazy Component...</p>}>
+                  <LazyComponent />
+                </React.Suspense>
+              }
+            />
+          </Routes>
+        </MemoryRouter>
+      );
+
+      await assertNavigation(container, resolve, resolveLazy);
+    });
+
+    // eslint-disable-next-line jest/expect-expect
+    it("BrowserRouter", async () => {
+      let { Home, About, LazyComponent, resolve, resolveLazy } =
+        getComponents();
+
+      let { container } = render(
+        <BrowserRouter window={getWindowImpl("/", false)}>
+          <Routes>
+            <Route path="/" element={<Home />} />
+            <Route
+              path="/about"
+              element={
+                <React.Suspense fallback={<p>Loading...</p>}>
+                  <About />
+                </React.Suspense>
+              }
+            />
+            <Route
+              path="/lazy"
+              element={
+                <React.Suspense fallback={<p>Loading Lazy Component...</p>}>
+                  <LazyComponent />
+                </React.Suspense>
+              }
+            />
+          </Routes>
+        </BrowserRouter>
+      );
+
+      await assertNavigation(container, resolve, resolveLazy);
+    });
+
+    // eslint-disable-next-line jest/expect-expect
+    it("HashRouter", async () => {
+      let { Home, About, LazyComponent, resolve, resolveLazy } =
+        getComponents();
+
+      let { container } = render(
+        <HashRouter window={getWindowImpl("/", true)}>
+          <Routes>
+            <Route path="/" element={<Home />} />
+            <Route
+              path="/about"
+              element={
+                <React.Suspense fallback={<p>Loading...</p>}>
+                  <About />
+                </React.Suspense>
+              }
+            />
+            <Route
+              path="/lazy"
+              element={
+                <React.Suspense fallback={<p>Loading Lazy Component...</p>}>
+                  <LazyComponent />
+                </React.Suspense>
+              }
+            />
+          </Routes>
+        </HashRouter>
+      );
+
+      await assertNavigation(container, resolve, resolveLazy);
+    });
+
+    // eslint-disable-next-line jest/expect-expect
+    it("RouterProvider", async () => {
+      let { Home, About, LazyComponent, resolve, resolveLazy } =
+        getComponents();
+
+      let router = createMemoryRouter(
+        createRoutesFromElements(
+          <>
+            <Route path="/" element={<Home />} />
+            <Route
+              path="/about"
+              element={
+                <React.Suspense fallback={<p>Loading...</p>}>
+                  <About />
+                </React.Suspense>
+              }
+            />
+            <Route
+              path="/lazy"
+              element={
+                <React.Suspense fallback={<p>Loading Lazy Component...</p>}>
+                  <LazyComponent />
+                </React.Suspense>
+              }
+            />
+          </>
+        )
+      );
+      let { container } = render(<RouterProvider router={router} />);
+
+      await assertNavigation(container, resolve, resolveLazy);
+    });
+  });
+
+  describe("when the destination route suspends without a boundary", () => {
+    async function assertNavigation(
+      container: HTMLElement,
+      resolve: () => void,
+      resolveLazy: () => void
+    ) {
+      // Start on home
+      expect(getHtml(container)).toMatch("Home");
+
+      // Click to /about and should see the frozen current UI
+      await act(() => {
+        fireEvent.click(screen.getByText("/about"));
+      });
+      await waitFor(() => screen.getByText("Home"));
+      expect(getHtml(container)).toMatch("Home");
+
+      // Resolve the destination UI to clear the boundary
+      await act(() => resolve());
+      await waitFor(() => screen.getByText("About"));
+      expect(getHtml(container)).toMatch("About");
+
+      // Back to home
+      await act(() => {
+        fireEvent.click(screen.getByText("back"));
+      });
+      await waitFor(() => screen.getByText("Home"));
+      expect(getHtml(container)).toMatch("Home");
+
+      // Click to /lazy and should see the frozen current UI
+      await act(() => {
+        fireEvent.click(screen.getByText("/lazy"));
+      });
+      await waitFor(() => screen.getByText("Home"));
+      expect(getHtml(container)).toMatch("Home");
+
+      // Resolve the lazy component to clear the boundary
+      await act(() => resolveLazy());
+      await waitFor(() => screen.getByText("Lazy"));
+      expect(getHtml(container)).toMatch("Lazy");
+    }
+
+    // eslint-disable-next-line jest/expect-expect
+    it("MemoryRouter", async () => {
+      let { Home, About, resolve, LazyComponent, resolveLazy } =
+        getComponents();
+
+      let { container } = render(
+        <MemoryRouter>
+          <Routes>
+            <Route path="/" element={<Home />} />
+            <Route path="/about" element={<About />} />
+            <Route path="/lazy" element={<LazyComponent />} />
+          </Routes>
+        </MemoryRouter>
+      );
+
+      await assertNavigation(container, resolve, resolveLazy);
+    });
+
+    // eslint-disable-next-line jest/expect-expect
+    it("BrowserRouter", async () => {
+      let { Home, About, resolve, LazyComponent, resolveLazy } =
+        getComponents();
+
+      let { container } = render(
+        <BrowserRouter window-={getWindowImpl("/", true)}>
+          <Routes>
+            <Route path="/" element={<Home />} />
+            <Route path="/about" element={<About />} />
+            <Route path="/lazy" element={<LazyComponent />} />
+          </Routes>
+        </BrowserRouter>
+      );
+
+      await assertNavigation(container, resolve, resolveLazy);
+    });
+
+    // eslint-disable-next-line jest/expect-expect
+    it("HashRouter", async () => {
+      let { Home, About, resolve, LazyComponent, resolveLazy } =
+        getComponents();
+
+      let { container } = render(
+        <HashRouter window-={getWindowImpl("/", true)}>
+          <Routes>
+            <Route path="/" element={<Home />} />
+            <Route path="/about" element={<About />} />
+            <Route path="/lazy" element={<LazyComponent />} />
+          </Routes>
+        </HashRouter>
+      );
+
+      await assertNavigation(container, resolve, resolveLazy);
+    });
+
+    // eslint-disable-next-line jest/expect-expect
+    it("RouterProvider", async () => {
+      let { Home, About, resolve, LazyComponent, resolveLazy } =
+        getComponents();
+
+      let router = createMemoryRouter(
+        createRoutesFromElements(
+          <>
+            <Route path="/" element={<Home />} />
+            <Route path="/about" element={<About />} />
+            <Route path="/lazy" element={<LazyComponent />} />
+          </>
+        )
+      );
+      let { container } = render(<RouterProvider router={router} />);
+
+      await assertNavigation(container, resolve, resolveLazy);
+    });
+  });
+});
+
+function getWindowImpl(initialUrl: string, isHash = false): Window {
+  // Need to use our own custom DOM in order to get a working history
+  const dom = new JSDOM(`<!DOCTYPE html>`, { url: "http://localhost/" });
+  dom.window.history.replaceState(null, "", (isHash ? "#" : "") + initialUrl);
+  return dom.window as unknown as Window;
+}
+
+function getHtml(container: HTMLElement) {
+  return prettyDOM(container, undefined, {
+    highlight: false,
+  });
+}
+
+async function tick() {
+  await new Promise((r) => setTimeout(r, 0));
+}
+
+function createDeferred() {
+  let resolve: (val?: any) => Promise<void>;
+  let reject: (error?: Error) => Promise<void>;
+  let promise = new Promise((res, rej) => {
+    resolve = async (val: any) => {
+      res(val);
+      await tick();
+      await promise;
+    };
+    reject = async (error?: Error) => {
+      rej(error);
+      await promise.catch(() => tick());
+    };
+  });
+  return {
+    promise,
+    //@ts-ignore
+    resolve,
+    //@ts-ignore
+    reject,
+  };
+}
diff --git a/packages/react-router-dom/index.tsx b/packages/react-router-dom/index.tsx
index e835e8ef63..7495692dab 100644
--- a/packages/react-router-dom/index.tsx
+++ b/packages/react-router-dom/index.tsx
@@ -4,7 +4,9 @@
  */
 import * as React from "react";
 import type {
+  Location,
   NavigateOptions,
+  NavigationType,
   RelativeRoutingType,
   RouteObject,
   To,
@@ -312,12 +314,20 @@ export function BrowserRouter({
   }
 
   let history = historyRef.current;
-  let [state, setState] = React.useState({
+  let [state, setStateImpl] = React.useState({
     action: history.action,
     location: history.location,
   });
+  let setState = React.useCallback(
+    (newState: { action: NavigationType; location: Location }) => {
+      "startTransition" in React
+        ? React.startTransition(() => setStateImpl(newState))
+        : setStateImpl(newState);
+    },
+    [setStateImpl]
+  );
 
-  React.useLayoutEffect(() => history.listen(setState), [history]);
+  React.useLayoutEffect(() => history.listen(setState), [history, setState]);
 
   return (
     <Router
@@ -347,12 +357,20 @@ export function HashRouter({ basename, children, window }: HashRouterProps) {
   }
 
   let history = historyRef.current;
-  let [state, setState] = React.useState({
+  let [state, setStateImpl] = React.useState({
     action: history.action,
     location: history.location,
   });
+  let setState = React.useCallback(
+    (newState: { action: NavigationType; location: Location }) => {
+      "startTransition" in React
+        ? React.startTransition(() => setStateImpl(newState))
+        : setStateImpl(newState);
+    },
+    [setStateImpl]
+  );
 
-  React.useLayoutEffect(() => history.listen(setState), [history]);
+  React.useLayoutEffect(() => history.listen(setState), [history, setState]);
 
   return (
     <Router
@@ -378,12 +396,20 @@ export interface HistoryRouterProps {
  * version of the history library that React Router uses internally.
  */
 function HistoryRouter({ basename, children, history }: HistoryRouterProps) {
-  const [state, setState] = React.useState({
+  let [state, setStateImpl] = React.useState({
     action: history.action,
     location: history.location,
   });
+  let setState = React.useCallback(
+    (newState: { action: NavigationType; location: Location }) => {
+      "startTransition" in React
+        ? React.startTransition(() => setStateImpl(newState))
+        : setStateImpl(newState);
+    },
+    [setStateImpl]
+  );
 
-  React.useLayoutEffect(() => history.listen(setState), [history]);
+  React.useLayoutEffect(() => history.listen(setState), [history, setState]);
 
   return (
     <Router
diff --git a/packages/react-router/__tests__/navigate-test.tsx b/packages/react-router/__tests__/navigate-test.tsx
index 5342758b5e..7c12862284 100644
--- a/packages/react-router/__tests__/navigate-test.tsx
+++ b/packages/react-router/__tests__/navigate-test.tsx
@@ -639,9 +639,6 @@ describe("<Navigate>", () => {
           {
             index: true,
             Component() {
-              // When state managed by react and changes during render, we'll
-              // only "see" the value from the first pass through here in our
-              // effects
               let [count, setCount] = React.useState(0);
               React.useEffect(() => {
                 if (count === 0) {
@@ -684,11 +681,11 @@ describe("<Navigate>", () => {
           Page B
         </h1>
         <p>
-          0
+          1
         </p>
       </div>"
     `);
-    expect(navigateSpy).toHaveBeenCalledTimes(2);
+    expect(navigateSpy).toHaveBeenCalledTimes(3);
     expect(navigateSpy.mock.calls[0]).toMatchObject([
       { pathname: "/b" },
       { state: { count: 0 } },
@@ -697,7 +694,7 @@ describe("<Navigate>", () => {
       { pathname: "/b" },
       { state: { count: 0 } },
     ]);
-    expect(renders).toEqual([0, 0]);
+    expect(renders).toEqual([1, 1]);
   });
 
   it("handles setState in render in StrictMode using a data router (async loader)", async () => {
diff --git a/packages/react-router/lib/components.tsx b/packages/react-router/lib/components.tsx
index 56ef4c1507..184426804b 100644
--- a/packages/react-router/lib/components.tsx
+++ b/packages/react-router/lib/components.tsx
@@ -63,7 +63,15 @@ export function RouterProvider({
 }: RouterProviderProps): React.ReactElement {
   // Need to use a layout effect here so we are subscribed early enough to
   // pick up on any render-driven redirects/navigations (useEffect/<Navigate>)
-  let [state, setState] = React.useState(router.state);
+  let [state, setStateImpl] = React.useState(router.state);
+  let setState = React.useCallback(
+    (newState: RouterState) => {
+      "startTransition" in React
+        ? React.startTransition(() => setStateImpl(newState))
+        : setStateImpl(newState);
+    },
+    [setStateImpl]
+  );
   React.useLayoutEffect(() => router.subscribe(setState), [router, setState]);
 
   let navigator = React.useMemo((): Navigator => {
@@ -164,12 +172,20 @@ export function MemoryRouter({
   }
 
   let history = historyRef.current;
-  let [state, setState] = React.useState({
+  let [state, setStateImpl] = React.useState({
     action: history.action,
     location: history.location,
   });
+  let setState = React.useCallback(
+    (newState: { action: NavigationType; location: Location }) => {
+      "startTransition" in React
+        ? React.startTransition(() => setStateImpl(newState))
+        : setStateImpl(newState);
+    },
+    [setStateImpl]
+  );
 
-  React.useLayoutEffect(() => history.listen(setState), [history]);
+  React.useLayoutEffect(() => history.listen(setState), [history, setState]);
 
   return (
     <Router