Experimental bfcache: Restore state w/ <Activity> (#77992)

> [!NOTE]  
> This feature does not affect the caching of dynamic data — only UI
state. Dynamic data is cached during back/forward navigations using a
different, preexisting mechanism; during normal navigations, dynamic
data is never cached, unless you opt in via
`experimental.staleTimes.dynamic`.

When navigating to a route that has already been visited, we should
restore as much of its previous state as we reasonably can — not only
state tracked by the browser, like scroll position and form inputs, but
also state owned by React, like `useState`.

The browser's native bfcache provides some of this automatically when
you use the back/forward buttons, or the history API. But it has
limitations, mainly that it can't restore state controlled by React. It
also doesn't work with navigations by regular links.

React has an experimental API called `<Activity>` that is designed for
this purpose. Content inside an Activity boundary can be hidden from the
UI without unmounting it; later, it can be restored to the UI without
having lost any state. React optimizes hidden Activity boundaries by
skipping over them during rendering, similar to how the browser handles
`content-visibility`. Later, during idle time, React will prerender the
boundaries so they are ready by the time they are revealed again.

In this PR, I've added an Activity boundary around every route segment
(i.e. layout or page). For each level of the route tree, the router will
render the N most recently active routes, and automatically toggle their
visibility as the user navigates, regardless of whether it's via the
back/forward buttons or regular links.

Aside from preserving state, keeping the inactive routes mounted also
makes navigations faster, since by the time you navigate the next screen
has already been prerendered.

In the future, we'll use this same mechanism to speculatively/
optimistically prerender routes that haven't yet been visited.

This is a nested bfcache — we track the history separately at each level
of the route tree. The lifetime of the bfcache is tied to the lifetime
of the React tree.

For now, the maximum number of entries per level is hardcoded to 3.
Eventually this will likely need to configurable per segment, but the
plan is for the default to be some smallish non-zero number. We'll
tinker with the heuristic before the feature is shipped to stable.

Author: acdlite

packages/next/src/client/components/bfcache.ts

@@ -0,0 +1,112 @@
+import type { FlightRouterState } from '../../server/app-render/types'
+import { useState } from 'react'
+
+// When the flag is disabled, only track the currently active tree
+const MAX_BF_CACHE_ENTRIES = process.env.__NEXT_ROUTER_BF_CACHE ? 3 : 1
+
+export type RouterBFCacheEntry = {
+  tree: FlightRouterState
+  stateKey: string
+  // The entries form a linked list, sorted in order of most recently active.
+  next: RouterBFCacheEntry | null
+}
+
+/**
+ * Keeps track of the most recent N trees (FlightRouterStates) that were active
+ * at a certain segment level. E.g. for a segment "/a/b/[param]", this hook
+ * tracks the last N param values that the router rendered for N.
+ *
+ * The result of this hook precisely determines the number and order of
+ * trees that are rendered in parallel at their segment level.
+ *
+ * The purpose of this cache is to we can preserve the React and DOM state of
+ * some number of inactive trees, by rendering them in an <Activity> boundary.
+ * That means it would not make sense for the the lifetime of the cache to be
+ * any longer than the lifetime of the React tree; e.g. if the hook were
+ * unmounted, then the React tree would be, too. So, we use React state to
+ * manage it.
+ *
+ * Note that we don't store the RSC data for the cache entries in this hook —
+ * the data for inactive segments is stored in the parent CacheNode, which
+ * *does* have a longer lifetime than the React tree. This hook only determines
+ * which of those trees should have their *state* preserved, by <Activity>.
+ */
+export function useRouterBFCache(
+  activeTree: FlightRouterState,
+  activeStateKey: string
+): RouterBFCacheEntry {
+  // The currently active entry. The entries form a linked list, sorted in
+  // order of most recently active. This allows us to reuse parts of the list
+  // without cloning, unless there's a reordering or removal.
+  // TODO: Once we start tracking back/forward history at each route level,
+  // we should use the history order instead. In other words, when traversing
+  // to an existing entry as a result of a popstate event, we should maintain
+  // the existing order instead of moving it to the front of the list. I think
+  // an initial implementation of this could be to pass an incrementing id
+  // to history.pushState/replaceState, then use that here for ordering.
+  const [prevActiveEntry, setPrevActiveEntry] = useState<RouterBFCacheEntry>(
+    () => {
+      const initialEntry: RouterBFCacheEntry = {
+        tree: activeTree,
+        stateKey: activeStateKey,
+        next: null,
+      }
+      return initialEntry
+    }
+  )
+
+  if (prevActiveEntry.tree === activeTree) {
+    // Fast path. The active tree hasn't changed, so we can reuse the
+    // existing state.
+    return prevActiveEntry
+  }
+
+  // The route tree changed. Note that this doesn't mean that the tree changed
+  // *at this level* — the change may be due to a child route. Either way, we
+  // need to either add or update the router tree in the bfcache.
+  //
+  // The rest of the code looks more complicated than it actually is because we
+  // can't mutate the state in place; we have to copy-on-write.
+
+  // Create a new entry for the active cache key. This is the head of the new
+  // linked list.
+  const newActiveEntry: RouterBFCacheEntry = {
+    tree: activeTree,
+    stateKey: activeStateKey,
+    next: null,
+  }
+
+  // We need to append the old list onto the new list. If the head of the new
+  // list was already present in the cache, then we'll need to clone everything
+  // that came before it. Then we can reuse the rest.
+  let n = 1
+  let oldEntry: RouterBFCacheEntry | null = prevActiveEntry
+  let clonedEntry: RouterBFCacheEntry = newActiveEntry
+  while (oldEntry !== null && n < MAX_BF_CACHE_ENTRIES) {
+    if (oldEntry.stateKey === activeStateKey) {
+      // Fast path. This entry in the old list that corresponds to the key that
+      // is now active. We've already placed a clone of this entry at the front
+      // of the new list. We can reuse the rest of the old list without cloning.
+      // NOTE: We don't need to worry about eviction in this case because we
+      // haven't increased the size of the cache, and we assume the max size
+      // is constant across renders. If we were to change it to a dynamic limit,
+      // then the implementation would need to account for that.
+      clonedEntry.next = oldEntry.next
+      break
+    } else {
+      // Clone the entry and append it to the list.
+      n++
+      const entry: RouterBFCacheEntry = {
+        tree: oldEntry.tree,
+        stateKey: oldEntry.stateKey,
+        next: null,
+      }
+      clonedEntry.next = entry
+      clonedEntry = entry
+    }
+    oldEntry = oldEntry.next
+  }
+
+  setPrevActiveEntry(newActiveEntry)
+  return newActiveEntry
+}

packages/next/src/client/components/layout-router.tsx

@@ -39,6 +39,11 @@ import { HTTPAccessFallbackBoundary } from './http-access-fallback/error-boundar
 import { createRouterCacheKey } from './router-reducer/create-router-cache-key'
 import { hasInterceptionRouteInCurrentTree } from './router-reducer/reducers/has-interception-route-in-current-tree'
 import { dispatchAppRouterAction } from './use-action-queue'
+import { useRouterBFCache, type RouterBFCacheEntry } from './bfcache'
+
+const Activity = process.env.__NEXT_ROUTER_BF_CACHE
+  ? require('react').unstable_Activity
+  : null
 
 /**
  * Add refetch marker to router state at the point of the current layout segment.
@@ -526,13 +531,7 @@ export default function OuterLayoutRouter({
     segmentMap = new Map()
     parentParallelRoutes.set(parallelRouterKey, segmentMap)
   }
-
-  // Get the active segment in the tree
-  // The reason arrays are used in the data format is that these are transferred from the server to the browser so it's optimized to save bytes.
   const parentTreeSegment = parentTree[0]
-  const tree = parentTree[1][parallelRouterKey]
-  const treeSegment = tree[0]
-
   const segmentPath =
     parentSegmentPath === null
       ? // TODO: The root segment value is currently omitted from the segment
@@ -551,31 +550,49 @@ export default function OuterLayoutRouter({
   // it's possible that the segment accessed the search params on the server.
   // (This only applies to page segments; layout segments cannot access search
   // params on the server.)
-  const cacheKey = createRouterCacheKey(treeSegment)
-  const stateKey = createRouterCacheKey(treeSegment, true) // no search params
-
-  // Read segment path from the parallel router cache node.
-  let cacheNode = segmentMap.get(cacheKey)
-  if (cacheNode === undefined) {
-    // When data is not available during rendering client-side we need to fetch
-    // it from the server.
-    const newLazyCacheNode: LazyCacheNode = {
-      lazyData: null,
-      rsc: null,
-      prefetchRsc: null,
-      head: null,
-      prefetchHead: null,
-      parallelRoutes: new Map(),
-      loading: null,
-      navigatedAt: -1,
-    }
+  const activeTree = parentTree[1][parallelRouterKey]
+  const activeSegment = activeTree[0]
+  const activeStateKey = createRouterCacheKey(activeSegment, true) // no search params
+
+  // At each level of the route tree, not only do we render the currently
+  // active segment — we also render the last N segments that were active at
+  // this level inside a hidden <Activity> boundary, to preserve their state
+  // if or when the user navigates to them again.
+  //
+  // bfcacheEntry is a linked list of FlightRouterStates.
+  let bfcacheEntry: RouterBFCacheEntry | null = useRouterBFCache(
+    activeTree,
+    activeStateKey
+  )
+  let children: Array<React.ReactNode> = []
+  do {
+    const tree = bfcacheEntry.tree
+    const stateKey = bfcacheEntry.stateKey
+    const segment = tree[0]
+    const cacheKey = createRouterCacheKey(segment)
+
+    // Read segment path from the parallel router cache node.
+    let cacheNode = segmentMap.get(cacheKey)
+    if (cacheNode === undefined) {
+      // When data is not available during rendering client-side we need to fetch
+      // it from the server.
+      const newLazyCacheNode: LazyCacheNode = {
+        lazyData: null,
+        rsc: null,
+        prefetchRsc: null,
+        head: null,
+        prefetchHead: null,
+        parallelRoutes: new Map(),
+        loading: null,
+        navigatedAt: -1,
+      }
 
-    // Flight data fetch kicked off during render and put into the cache.
-    cacheNode = newLazyCacheNode
-    segmentMap.set(cacheKey, newLazyCacheNode)
-  }
+      // Flight data fetch kicked off during render and put into the cache.
+      cacheNode = newLazyCacheNode
+      segmentMap.set(cacheKey, newLazyCacheNode)
+    }
 
-  /*
+    /*
     - Error boundary
       - Only renders error boundary if error component is provided.
       - Rendered for each segment to ensure they have their own error state.
@@ -585,49 +602,66 @@ export default function OuterLayoutRouter({
       - Passed to the router during rendering to ensure it can be immediately rendered when suspending on a Flight fetch.
   */
 
-  // TODO: The loading module data for a segment is stored on the parent, then
-  // applied to each of that parent segment's parallel route slots. In the
-  // simple case where there's only one parallel route (the `children` slot),
-  // this is no different from if the loading module data where stored on the
-  // child directly. But I'm not sure this actually makes sense when there are
-  // multiple parallel routes. It's not a huge issue because you always have
-  // the option to define a narrower loading boundary for a particular slot. But
-  // this sort of smells like an implementation accident to me.
-  const loadingModuleData = parentCacheNode.loading
+    // TODO: The loading module data for a segment is stored on the parent, then
+    // applied to each of that parent segment's parallel route slots. In the
+    // simple case where there's only one parallel route (the `children` slot),
+    // this is no different from if the loading module data where stored on the
+    // child directly. But I'm not sure this actually makes sense when there are
+    // multiple parallel routes. It's not a huge issue because you always have
+    // the option to define a narrower loading boundary for a particular slot. But
+    // this sort of smells like an implementation accident to me.
+    const loadingModuleData = parentCacheNode.loading
+    let child = (
+      <TemplateContext.Provider
+        key={stateKey}
+        value={
+          <ScrollAndFocusHandler segmentPath={segmentPath}>
+            <ErrorBoundary
+              errorComponent={error}
+              errorStyles={errorStyles}
+              errorScripts={errorScripts}
+            >
+              <LoadingBoundary loading={loadingModuleData}>
+                <HTTPAccessFallbackBoundary
+                  notFound={notFound}
+                  forbidden={forbidden}
+                  unauthorized={unauthorized}
+                >
+                  <RedirectBoundary>
+                    <InnerLayoutRouter
+                      url={url}
+                      tree={tree}
+                      cacheNode={cacheNode}
+                      segmentPath={segmentPath}
+                    />
+                  </RedirectBoundary>
+                </HTTPAccessFallbackBoundary>
+              </LoadingBoundary>
+            </ErrorBoundary>
+          </ScrollAndFocusHandler>
+        }
+      >
+        {templateStyles}
+        {templateScripts}
+        {template}
+      </TemplateContext.Provider>
+    )
 
-  return (
-    <TemplateContext.Provider
-      key={stateKey}
-      value={
-        <ScrollAndFocusHandler segmentPath={segmentPath}>
-          <ErrorBoundary
-            errorComponent={error}
-            errorStyles={errorStyles}
-            errorScripts={errorScripts}
-          >
-            <LoadingBoundary loading={loadingModuleData}>
-              <HTTPAccessFallbackBoundary
-                notFound={notFound}
-                forbidden={forbidden}
-                unauthorized={unauthorized}
-              >
-                <RedirectBoundary>
-                  <InnerLayoutRouter
-                    url={url}
-                    tree={tree}
-                    cacheNode={cacheNode}
-                    segmentPath={segmentPath}
-                  />
-                </RedirectBoundary>
-              </HTTPAccessFallbackBoundary>
-            </LoadingBoundary>
-          </ErrorBoundary>
-        </ScrollAndFocusHandler>
-      }
-    >
-      {templateStyles}
-      {templateScripts}
-      {template}
-    </TemplateContext.Provider>
-  )
+    if (process.env.__NEXT_ROUTER_BF_CACHE) {
+      child = (
+        <Activity
+          key={stateKey}
+          mode={stateKey === activeStateKey ? 'visible' : 'hidden'}
+        >
+          {child}
+        </Activity>
+      )
+    }
+
+    children.push(child)
+
+    bfcacheEntry = bfcacheEntry.next
+  } while (bfcacheEntry !== null)
+
+  return children
 }

test/e2e/app-dir/back-forward-cache/app/layout.tsx

@@ -1,11 +1,29 @@
+import Link from 'next/link'
+
 export default function RootLayout({
   children,
 }: {
   children: React.ReactNode
 }) {
+  const links = []
+  for (let n = 1; n <= 5; n++) {
+    links.push(
+      <li key={n}>
+        <Link href={`/page/${n}`}>Page {n}</Link>
+      </li>
+    )
+    links.push(
+      <li key={n + '-with-search-param'}>
+        <Link href={`/page/${n}?param=true`}>Page {n} (with search param)</Link>
+      </li>
+    )
+  }
   return (
     <html lang="en">
-      <body>{children}</body>
+      <body>
+        <ul>{links}</ul>
+        <div>{children}</div>
+      </body>
     </html>
   )
 }

test/e2e/app-dir/back-forward-cache/app/page/[n]/page.tsx

@@ -0,0 +1,20 @@
+import { Suspense } from 'react'
+import { StatefulClientComponent } from './stateful-client-component'
+
+export default async function Page({
+  params,
+}: {
+  params: Promise<{ n: string }>
+}) {
+  const { n } = await params
+  return (
+    <>
+      <h2>Page {n}</h2>
+      <div>
+        <Suspense fallback={<div>Loading...</div>}>
+          <StatefulClientComponent n={n} />
+        </Suspense>
+      </div>
+    </>
+  )
+}