Fix issues with partial hydration combined with route.lazy (#11121)

.changeset/lazy-partial-hydration.md

@@ -0,0 +1,6 @@
+---
+"react-router": patch
+"@remix-run/router": patch
+---
+
+Fix bug with `route.lazy` not working correctly on initial SPA load when `v7_partialHydration` is specified

packages/react-router-dom/__tests__/partial-hydration-test.tsx

@@ -14,7 +14,7 @@ import {
 } from "react-router-dom";
 
 import getHtml from "../../react-router/__tests__/utils/getHtml";
-import { createDeferred } from "../../router/__tests__/utils/utils";
+import { createDeferred, tick } from "../../router/__tests__/utils/utils";
 
 let didAssertMissingHydrateFallback = false;
 
@@ -521,4 +521,133 @@ function testPartialHydration(
       </div>"
     `);
   });
+
+  it("supports partial hydration w/lazy initial routes (leaf fallback)", async () => {
+    let dfd = createDeferred();
+    let router = createTestRouter(
+      [
+        {
+          path: "/",
+          Component() {
+            return (
+              <>
+                <h1>Root</h1>
+                <Outlet />
+              </>
+            );
+          },
+          children: [
+            {
+              id: "index",
+              index: true,
+              HydrateFallback: () => <p>Index Loading...</p>,
+              async lazy() {
+                await tick();
+                return {
+                  loader: () => dfd.promise,
+                  Component() {
+                    let data = useLoaderData() as string;
+                    return <h2>{`Index - ${data}`}</h2>;
+                  },
+                };
+              },
+            },
+          ],
+        },
+      ],
+      {
+        future: {
+          v7_partialHydration: true,
+        },
+      }
+    );
+    let { container } = render(<RouterProvider router={router} />);
+
+    expect(getHtml(container)).toMatchInlineSnapshot(`
+      "<div>
+        <h1>
+          Root
+        </h1>
+        <p>
+          Index Loading...
+        </p>
+      </div>"
+    `);
+
+    dfd.resolve("INDEX DATA");
+    await waitFor(() => screen.getByText(/INDEX DATA/));
+    expect(getHtml(container)).toMatchInlineSnapshot(`
+      "<div>
+        <h1>
+          Root
+        </h1>
+        <h2>
+          Index - INDEX DATA
+        </h2>
+      </div>"
+    `);
+  });
+
+  it("supports partial hydration w/lazy initial routes (root fallback)", async () => {
+    let dfd = createDeferred();
+    let router = createTestRouter(
+      [
+        {
+          path: "/",
+          Component() {
+            return (
+              <>
+                <h1>Root</h1>
+                <Outlet />
+              </>
+            );
+          },
+          HydrateFallback: () => <p>Loading...</p>,
+          children: [
+            {
+              id: "index",
+              index: true,
+              async lazy() {
+                await tick();
+                return {
+                  loader: () => dfd.promise,
+                  Component() {
+                    let data = useLoaderData() as string;
+                    return <h2>{`Index - ${data}`}</h2>;
+                  },
+                };
+              },
+            },
+          ],
+        },
+      ],
+      {
+        future: {
+          v7_partialHydration: true,
+        },
+      }
+    );
+    let { container } = render(<RouterProvider router={router} />);
+
+    expect(getHtml(container)).toMatchInlineSnapshot(`
+      "<div>
+        <p>
+          Loading...
+        </p>
+      </div>"
+    `);
+
+    dfd.resolve("INDEX DATA");
+    await waitFor(() => screen.getByText(/INDEX DATA/));
+    expect(getHtml(container)).toMatchInlineSnapshot(`
+      "<div>
+        <h1>
+          Root
+        </h1>
+        <h2>
+          Index - INDEX DATA
+        </h2>
+      </div>"
+    `);
+  });
 }

packages/react-router/lib/hooks.tsx

@@ -438,7 +438,8 @@ export function useRoutesImpl(
     warning(
       matches == null ||
         matches[matches.length - 1].route.element !== undefined ||
-        matches[matches.length - 1].route.Component !== undefined,
+        matches[matches.length - 1].route.Component !== undefined ||
+        matches[matches.length - 1].route.lazy !== undefined,
       `Matched leaf route at location "${location.pathname}${location.search}${location.hash}" ` +
         `does not have an element or Component. This means it will render an <Outlet /> with a ` +
         `null value by default resulting in an "empty" page.`
@@ -704,23 +705,25 @@ export function _renderMatches(
       if (match.route.HydrateFallback || match.route.hydrateFallbackElement) {
         fallbackIndex = i;
       }
-      if (
-        match.route.loader &&
-        match.route.id &&
-        dataRouterState.loaderData[match.route.id] === undefined &&
-        (!dataRouterState.errors ||
-          dataRouterState.errors[match.route.id] === undefined)
-      ) {
-        // We found the first route without data/errors which means it's loader
-        // still needs to run.  Flag that we need to render a fallback and
-        // render up until the appropriate fallback
-        renderFallback = true;
-        if (fallbackIndex >= 0) {
-          renderedMatches = renderedMatches.slice(0, fallbackIndex + 1);
-        } else {
-          renderedMatches = [renderedMatches[0]];
+
+      if (match.route.id) {
+        let { loaderData, errors } = dataRouterState;
+        let needsToRunLoader =
+          match.route.loader &&
+          loaderData[match.route.id] === undefined &&
+          (!errors || errors[match.route.id] === undefined);
+        if (match.route.lazy || needsToRunLoader) {
+          // We found the first route that's not ready to render (waiting on
+          // lazy, or has a loader that hasn't run yet).  Flag that we need to
+          // render a fallback and render up until the appropriate fallback
+          renderFallback = true;
+          if (fallbackIndex >= 0) {
+            renderedMatches = renderedMatches.slice(0, fallbackIndex + 1);
+          } else {
+            renderedMatches = [renderedMatches[0]];
+          }
+          break;
         }
-        break;
       }
     }
   }

packages/router/__tests__/lazy-test.ts

@@ -43,6 +43,104 @@ describe("lazily loaded route modules", () => {
     },
   ];
 
+  describe("initialization", () => {
+    it("fetches lazy route modules on router initialization", async () => {
+      let dfd = createDeferred();
+      let router = createRouter({
+        routes: [
+          {
+            path: "/lazy",
+            lazy: () => dfd.promise,
+          },
+        ],
+        history: createMemoryHistory({ initialEntries: ["/lazy"] }),
+      });
+
+      expect(router.state.initialized).toBe(false);
+
+      router.initialize();
+
+      let route = { Component: () => null };
+      await dfd.resolve(route);
+
+      expect(router.state.location.pathname).toBe("/lazy");
+      expect(router.state.navigation.state).toBe("idle");
+      expect(router.state.initialized).toBe(true);
+      expect(router.state.matches[0].route).toMatchObject(route);
+    });
+
+    it("fetches lazy route modules and executes loaders on router initialization", async () => {
+      let dfd = createDeferred();
+      let router = createRouter({
+        routes: [
+          {
+            path: "/lazy",
+            lazy: () => dfd.promise,
+          },
+        ],
+        history: createMemoryHistory({ initialEntries: ["/lazy"] }),
+      });
+
+      expect(router.state.initialized).toBe(false);
+
+      router.initialize();
+
+      let loaderDfd = createDeferred();
+      let route = {
+        Component: () => null,
+        loader: () => loaderDfd.promise,
+      };
+      await dfd.resolve(route);
+      expect(router.state.initialized).toBe(false);
+
+      await loaderDfd.resolve("LOADER");
+      expect(router.state.location.pathname).toBe("/lazy");
+      expect(router.state.navigation.state).toBe("idle");
+      expect(router.state.initialized).toBe(true);
+      expect(router.state.loaderData).toEqual({
+        "0": "LOADER",
+      });
+      expect(router.state.matches[0].route).toMatchObject(route);
+    });
+
+    it("fetches lazy route modules and executes loaders with v7_partialHydration enabled", async () => {
+      let dfd = createDeferred();
+      let router = createRouter({
+        routes: [
+          {
+            path: "/lazy",
+            lazy: () => dfd.promise,
+          },
+        ],
+        history: createMemoryHistory({ initialEntries: ["/lazy"] }),
+        future: {
+          v7_partialHydration: true,
+        },
+      });
+
+      expect(router.state.initialized).toBe(false);
+
+      router.initialize();
+
+      let loaderDfd = createDeferred();
+      let route = {
+        Component: () => null,
+        loader: () => loaderDfd.promise,
+      };
+      await dfd.resolve(route);
+      expect(router.state.initialized).toBe(false);
+
+      await loaderDfd.resolve("LOADER");
+      expect(router.state.location.pathname).toBe("/lazy");
+      expect(router.state.navigation.state).toBe("idle");
+      expect(router.state.initialized).toBe(true);
+      expect(router.state.loaderData).toEqual({
+        "0": "LOADER",
+      });
+      expect(router.state.matches[0].route).toMatchObject(route);
+    });
+  });
+
   describe("happy path", () => {
     it("fetches lazy route modules on loading navigation", async () => {
       let t = setup({ routes: LAZY_ROUTES });

packages/router/router.ts

@@ -3655,21 +3655,27 @@ function getMatchesToLoad(
   let boundaryMatches = getLoaderMatchesUntilBoundary(matches, boundaryId);
 
   let navigationMatches = boundaryMatches.filter((match, index) => {
-    if (isInitialLoad) {
-      // On initial hydration we don't do any shouldRevalidate stuff - we just
-      // call the unhydrated loaders
-      return isUnhydratedRoute(state, match.route);
-    }
-
-    if (match.route.lazy) {
+    let { route } = match;
+    if (route.lazy) {
       // We haven't loaded this route yet so we don't know if it's got a loader!
       return true;
     }
 
-    if (match.route.loader == null) {
+    if (route.loader == null) {
       return false;
     }
 
+    if (isInitialLoad) {
+      if (route.loader.hydrate) {
+        return true;
+      }
+      return (
+        state.loaderData[route.id] === undefined &&
+        // Don't re-run if the loader ran and threw an error
+        (!state.errors || state.errors[route.id] === undefined)
+      );
+    }
+
     // Always call the loader on new route instances and pending defer cancellations
     if (
       isNewLoader(state.loaderData, state.matches[index], match) ||
@@ -3789,23 +3795,6 @@ function getMatchesToLoad(
   return [navigationMatches, revalidatingFetchers];
 }
 
-// Is this route unhydrated (when v7_partialHydration=true) such that we need
-// to call it's loader on the initial router creation
-function isUnhydratedRoute(state: RouterState, route: AgnosticDataRouteObject) {
-  if (!route.loader) {
-    return false;
-  }
-  if (route.loader.hydrate) {
-    return true;
-  }
-  return (
-    state.loaderData[route.id] === undefined &&
-    (!state.errors ||
-      // Loader ran but errored - don't re-run
-      state.errors[route.id] === undefined)
-  );
-}
-
 function isNewLoader(
   currentLoaderData: RouteData,
   currentMatch: AgnosticDataRouteMatch,