Page data HMR (#3132)

Based on #2968 

This builds upon #2968 and @ForsakenHarmony's work on data routes to
enable page data HMR.

Page data HMR is a bit more clever than it is in Next.js as we won't
re-render a Node.js result for each page file update. Instead, thanks to
the `StripPageDefaultExport` transform, there are three versions of the
page chunks:
* client-side (strips page data exports);
* server-side (full);
* data server-side (strips page default export).

Instead of subscribing to the full server-side result, on hydration, the
client-side page separately subscribes to:
* client-side updates (already the case);
* data server-side updates (new).

This means that updating something that only affects the page component
will only cause a client-side update and **no Node.js re-rendering**,
while updating something that only affects the data will only cause a
server-side update.

~~I'm marking this as a draft for now as there are still a few areas to
test/investigate:~~
- [x] When something that is used in both the default page export and
data exports is changed, this will cause *two* HMR updates: one data
update, and one client-side chunk update. **The same case breaks in
Next.js, where we will receive a client-side update, but no server-side
update, ending up with an incorrect result.**
- [x] Differences between `getStaticProps/getServerSideProps`, as well
as `getInitialProps` (need to talk with @timneutkens about this) (see
vercel/next.js#44523)

Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>

crates/next-core/js/src/dev/hmr-client.ts

@@ -47,20 +47,25 @@ export function connect({ assetPrefix }: ClientOptions) {
   }
   globalThis.TURBOPACK_CHUNK_UPDATE_LISTENERS = {
     push: ([chunkPath, callback]: [ChunkPath, UpdateCallback]) => {
-      onChunkUpdate(chunkPath, callback);
+      subscribeToChunkUpdate(chunkPath, callback);
     },
   };
 
   if (Array.isArray(queued)) {
     for (const [chunkPath, callback] of queued) {
-      onChunkUpdate(chunkPath, callback);
+      subscribeToChunkUpdate(chunkPath, callback);
     }
   }
 
   subscribeToInitialCssChunksUpdates(assetPrefix);
 }
 
-const updateCallbacks: Map<ResourceKey, Set<UpdateCallback>> = new Map();
+type UpdateCallbackSet = {
+  callbacks: Set<UpdateCallback>;
+  unsubscribe: () => void;
+};
+
+const updateCallbackSets: Map<ResourceKey, UpdateCallbackSet> = new Map();
 
 function sendJSON(message: ClientMessage) {
   sendMessage(JSON.stringify(message));
@@ -75,15 +80,22 @@ function resourceKey(resource: ResourceIdentifier): ResourceKey {
   });
 }
 
-function subscribeToUpdates(resource: ResourceIdentifier) {
+function subscribeToUpdates(resource: ResourceIdentifier): () => void {
   sendJSON({
     type: "subscribe",
     ...resource,
   });
+
+  return () => {
+    sendJSON({
+      type: "unsubscribe",
+      ...resource,
+    });
+  };
 }
 
 function handleSocketConnected() {
-  for (const key of updateCallbacks.keys()) {
+  for (const key of updateCallbackSets.keys()) {
     subscribeToUpdates(JSON.parse(key));
   }
 }
@@ -249,42 +261,55 @@ function handleSocketMessage(msg: ServerMessage) {
   }
 }
 
-export function onChunkUpdate(chunkPath: ChunkPath, callback: UpdateCallback) {
-  onUpdate(
+export function subscribeToChunkUpdate(
+  chunkPath: ChunkPath,
+  callback: UpdateCallback
+): () => void {
+  return subscribeToUpdate(
     {
       path: chunkPath,
     },
     callback
   );
 }
 
-export function onUpdate(
+export function subscribeToUpdate(
   resource: ResourceIdentifier,
   callback: UpdateCallback
 ) {
   const key = resourceKey(resource);
-  let callbacks = updateCallbacks.get(key);
-  if (!callbacks) {
-    subscribeToUpdates(resource);
-    updateCallbacks.set(key, (callbacks = new Set([callback])));
+  let callbackSet: UpdateCallbackSet;
+  const existingCallbackSet = updateCallbackSets.get(key);
+  if (!existingCallbackSet) {
+    callbackSet = {
+      callbacks: new Set([callback]),
+      unsubscribe: subscribeToUpdates(resource),
+    };
+    updateCallbackSets.set(key, callbackSet);
   } else {
-    callbacks.add(callback);
+    existingCallbackSet.callbacks.add(callback);
+    callbackSet = existingCallbackSet;
   }
 
   return () => {
-    callbacks!.delete(callback);
+    callbackSet.callbacks.delete(callback);
+
+    if (callbackSet.callbacks.size === 0) {
+      callbackSet.unsubscribe();
+      updateCallbackSets.delete(key);
+    }
   };
 }
 
 function triggerUpdate(msg: ServerMessage) {
   const key = resourceKey(msg.resource);
-  const callbacks = updateCallbacks.get(key);
-  if (!callbacks) {
+  const callbackSet = updateCallbackSets.get(key);
+  if (!callbackSet) {
     return;
   }
 
   try {
-    for (const callback of callbacks) {
+    for (const callback of callbackSet.callbacks) {
       callback(msg);
     }
   } catch (err) {
@@ -324,7 +349,7 @@ export function subscribeToCssChunkUpdates(
   }
 
   const chunkPath = pathname.slice(cssChunkPrefix.length);
-  onChunkUpdate(chunkPath, (update) => {
+  subscribeToChunkUpdate(chunkPath, (update) => {
     switch (update.type) {
       case "restart": {
         console.info(`Reloading CSS chunk \`${chunkPath}\``);

crates/next-core/js/src/dev/hot-reloader.tsx

@@ -3,7 +3,7 @@
 import type React from "react";
 import { useRouter, usePathname } from "next/dist/client/components/navigation";
 import { useEffect } from "react";
-import { onUpdate } from "./hmr-client";
+import { subscribeToUpdate } from "./hmr-client";
 import { ReactDevOverlay } from "./client";
 
 type HotReloadProps = React.PropsWithChildren<{
@@ -15,7 +15,7 @@ export default function HotReload({ assetPrefix, children }: HotReloadProps) {
   const path = usePathname()!.slice(1);
 
   useEffect(() => {
-    const unsubscribe = onUpdate(
+    const unsubscribe = subscribeToUpdate(
       {
         path,
         headers: {

crates/next-core/js/src/entry/fallback.tsx

@@ -6,11 +6,11 @@ import {
   initializeHMR,
   ReactDevOverlay,
 } from "@vercel/turbopack-next/dev/client";
-import { onUpdate } from "@vercel/turbopack-next/dev/hmr-client";
+import { subscribeToUpdate } from "@vercel/turbopack-next/dev/hmr-client";
 
 const pageChunkPath = location.pathname.slice(1);
 
-onUpdate(
+subscribeToUpdate(
   {
     path: pageChunkPath,
     headers: {

crates/next-core/js/src/entry/next-hydrate.tsx

@@ -1,8 +1,17 @@
 import "@vercel/turbopack-next/internal/shims-client";
 
-import { initialize, hydrate } from "next/dist/client";
+import { initialize, hydrate, router } from "next/dist/client";
+import type { Router } from "next/dist/client/router";
+import {
+  assign,
+  urlQueryToSearchParams,
+} from "next/dist/shared/lib/router/utils/querystring";
+import { formatWithValidation } from "next/dist/shared/lib/router/utils/format-url";
 import { initializeHMR } from "@vercel/turbopack-next/dev/client";
-import { subscribeToCssChunkUpdates } from "@vercel/turbopack-next/dev/hmr-client";
+import {
+  subscribeToUpdate,
+  subscribeToCssChunkUpdates,
+} from "@vercel/turbopack-next/dev/hmr-client";
 
 import * as _app from "@vercel/turbopack-next/pages/_app";
 import * as page from ".";
@@ -66,5 +75,108 @@ async function loadPageChunk(assetPrefix: string, chunkPath: string) {
 
   await hydrate({});
 
+  // This needs to happen after hydration because the router is initialized
+  // during hydration. To make this dependency clearer, we pass `router` as an
+  // explicit argument instead of relying on the `router` import binding.
+  subscribeToCurrentPageData({ assetPrefix, router });
+
   console.debug("The page has been hydrated");
 })().catch((err) => console.error(err));
+
+/**
+ * Subscribes to the current page's data updates from the HMR server.
+ *
+ * Updates on route change.
+ */
+function subscribeToCurrentPageData({
+  router,
+  assetPrefix,
+}: {
+  router: Router;
+  assetPrefix: string;
+}) {
+  let dataPath = getCurrentPageDataHref();
+  let unsubscribe = subscribeToPageData({
+    router,
+    dataPath,
+    assetPrefix,
+  });
+
+  router.events.on("routeChangeComplete", () => {
+    const nextDataPath = getCurrentPageDataHref();
+    if (dataPath === nextDataPath) {
+      return;
+    }
+    dataPath = nextDataPath;
+
+    unsubscribe();
+    unsubscribe = subscribeToPageData({
+      router,
+      dataPath,
+      assetPrefix,
+    });
+  });
+}
+
+function getCurrentPageDataHref(): string {
+  return router.pageLoader.getDataHref({
+    asPath: router.asPath,
+    href: formatWithValidation({
+      // No need to pass `router.query` when `skipInterpolation` is true.
+      pathname: router.pathname,
+    }),
+    skipInterpolation: true,
+  });
+}
+
+/**
+ * TODO(alexkirsz): Handle assetPrefix/basePath.
+ */
+function subscribeToPageData({
+  router,
+  dataPath,
+  assetPrefix,
+}: {
+  router: Router;
+  dataPath: string;
+  assetPrefix: string;
+}): () => void {
+  return subscribeToUpdate(
+    {
+      // We need to remove the leading / from the data path as Turbopack
+      // resources are not prefixed with a /.
+      path: dataPath.slice(1),
+      headers: {
+        // This header is used by the Next.js server to determine whether this
+        // is a data request.
+        "x-nextjs-data": "1",
+      },
+    },
+    (update) => {
+      if (update.type !== "restart") {
+        return;
+      }
+
+      // This triggers a reload of the page data.
+      // Adapted from next.js/packages/next/client/next-dev.js.
+      router
+        .replace(
+          router.pathname +
+            "?" +
+            String(
+              assign(
+                urlQueryToSearchParams(router.query),
+                new URLSearchParams(location.search)
+              )
+            ),
+          router.asPath,
+          { scroll: false }
+        )
+        .catch(() => {
+          // trigger hard reload when failing to refresh data
+          // to show error overlay properly
+          location.reload();
+        });
+    }
+  );
+}

crates/next-core/js/src/entry/server-renderer.tsx

@@ -8,7 +8,6 @@ import "@vercel/turbopack-next/internal/shims";
 import type { IncomingMessage, ServerResponse } from "node:http";
 
 import { renderToHTML, RenderOpts } from "next/dist/server/render";
-import RenderResult from "next/dist/server/render-result";
 import type { BuildManifest } from "next/dist/server/get-page-files";
 
 import { ServerResponseShim } from "@vercel/turbopack-next/internal/http";
@@ -102,16 +101,25 @@ async function runOperation(
     ampFirstPages: [],
   };
 
+  // When rendering a data request, the default component export is eliminated
+  // by the Next.js strip export transform. The following checks for this case
+  // and replaces the default export with a dummy component instead.
+  const comp =
+    typeof Component === "undefined" ||
+    (typeof Component === "object" && Object.keys(Component).length === 0)
+      ? () => {}
+      : Component;
+
   const renderOpts: RenderOpts = {
     /* LoadComponentsReturnType */
-    Component,
+    Component: comp,
     App,
     Document,
     pageConfig: {},
     buildManifest,
     reactLoadableManifest: {},
     ComponentMod: {
-      default: Component,
+      default: comp,
       ...otherExports,
     },
     pathname: renderData.path,

crates/next-core/src/lib.rs

@@ -15,15 +15,17 @@ mod next_font_google;
 pub mod next_image;
 mod next_import_map;
 pub mod next_server;
+pub mod next_shared;
 mod page_loader;
+mod page_source;
 pub mod react_refresh;
 mod runtime;
-mod server_rendered_source;
 mod util;
 mod web_entry_source;
 
 pub use app_source::create_app_source;
-pub use server_rendered_source::create_server_rendered_source;
+pub use page_source::create_page_source;
+pub use turbopack_node::source_map;
 pub use web_entry_source::create_web_entry_source;
 
 pub fn register() {

crates/next-core/src/next_client/context.rs

@@ -136,35 +136,6 @@ pub async fn get_client_module_options_context(
     Ok(add_next_font_transform(module_options_context.cell()))
 }
 
-#[turbo_tasks::function]
-pub async fn add_next_transforms_to_pages(
-    module_options_context: ModuleOptionsContextVc,
-    pages_dir: FileSystemPathVc,
-) -> Result<ModuleOptionsContextVc> {
-    let mut module_options_context = module_options_context.await?.clone_value();
-    // Apply the Next SSG tranform to all pages.
-    module_options_context.custom_rules.push(ModuleRule::new(
-        ModuleRuleCondition::all(vec![
-            ModuleRuleCondition::ResourcePathInExactDirectory(pages_dir.await?),
-            ModuleRuleCondition::not(ModuleRuleCondition::ReferenceType(ReferenceType::Url(
-                UrlReferenceSubType::Undefined,
-            ))),
-            ModuleRuleCondition::any(vec![
-                ModuleRuleCondition::ResourcePathEndsWith(".js".to_string()),
-                ModuleRuleCondition::ResourcePathEndsWith(".jsx".to_string()),
-                ModuleRuleCondition::ResourcePathEndsWith(".ts".to_string()),
-                ModuleRuleCondition::ResourcePathEndsWith(".tsx".to_string()),
-            ]),
-        ]),
-        vec![ModuleRuleEffect::AddEcmascriptTransforms(
-            EcmascriptInputTransformsVc::cell(vec![
-                EcmascriptInputTransform::NextJsStripPageDataExports,
-            ]),
-        )],
-    ));
-    Ok(module_options_context.cell())
-}
-
 #[turbo_tasks::function]
 pub async fn add_next_font_transform(
     module_options_context: ModuleOptionsContextVc,

crates/next-core/src/next_shared/mod.rs

@@ -0,0 +1 @@
+pub mod transforms;