[Per-Segment Prefetch] Add Route Tree prefetch

In the new design for prefetching, instead of prefetching all the data
for a page in a single request, the client issues a separate request for
each segment. The benefit of using separate requests is that the
responses can be separately cached. For example, when prefetching
multiple pages that share the same layout, the layout data only needs to
be fetched once. Simlarly, when revalidating a tag, only segments that
access that tag need to be refetched.

The client needs to know about the structure of the route tree in order
to determine which segments can be reused. So, before it prefetches any
actual segment data for a page, it will first fetch information about
the route tree. The route tree contains the necessary data for sending
and caching prefetch requests.

This PR updates the server to respond to route tree prefetches. To
prefetch the route tree, the client sets the
`Next-Router-Segment-Prefetch` to `/_tree`.

packages/next/src/server/app-render/app-render.tsx

@@ -2763,6 +2763,7 @@ async function prerenderToStream(
         metadata.segmentFlightData = await collectSegmentData(
           finalAttemptRSCPayload,
           flightData,
+          finalRenderPrerenderStore,
           ComponentMod,
           renderOpts
         )
@@ -3229,6 +3230,7 @@ async function prerenderToStream(
         metadata.segmentFlightData = await collectSegmentData(
           finalServerPayload,
           flightData,
+          finalClientPrerenderStore,
           ComponentMod,
           renderOpts
         )
@@ -3361,6 +3363,7 @@ async function prerenderToStream(
         metadata.segmentFlightData = await collectSegmentData(
           RSCPayload,
           flightData,
+          ssrPrerenderStore,
           ComponentMod,
           renderOpts
         )
@@ -3553,6 +3556,7 @@ async function prerenderToStream(
         metadata.segmentFlightData = await collectSegmentData(
           RSCPayload,
           flightData,
+          prerenderLegacyStore,
           ComponentMod,
           renderOpts
         )
@@ -3714,6 +3718,7 @@ async function prerenderToStream(
         metadata.segmentFlightData = await collectSegmentData(
           errorRSCPayload,
           flightData,
+          prerenderLegacyStore,
           ComponentMod,
           renderOpts
         )
@@ -3847,6 +3852,7 @@ const getGlobalErrorStyles = async (
 async function collectSegmentData(
   rscPayload: InitialRSCPayload,
   fullPageDataBuffer: Buffer,
+  prerenderStore: PrerenderStore,
   ComponentMod: AppPageModule,
   renderOpts: RenderOpts
 ): Promise<Map<string, Buffer> | undefined> {
@@ -3896,9 +3902,11 @@ async function collectSegmentData(
     serverModuleMap: null,
   }
 
+  const staleTime = prerenderStore.stale
   return await ComponentMod.collectSegmentData(
     routeTree,
     fullPageDataBuffer,
+    staleTime,
     clientReferenceManifest.clientModules as ManifestNode,
     serverConsumerManifest
   )

packages/next/src/server/app-render/collect-segment-data.tsx

@@ -19,9 +19,43 @@ import { UNDERSCORE_NOT_FOUND_ROUTE } from '../../api/constants'
 import { waitAtLeastOneReactRenderTask } from '../../lib/scheduler'
 import type { LoadingModuleData } from '../../shared/lib/app-router-context.shared-runtime'
 
+// Contains metadata about the route tree. The client must fetch this before
+// it can fetch any actual segment data.
+type RootTreePrefetch = {
+  tree: TreePrefetch
+  staleTime: number
+}
+
+type TreePrefetch = {
+  // The key to use when requesting the data for this segment (analogous to a
+  // URL). Also used as a cache key, although the server may specify a different
+  // cache key when it responds (analagous to a Vary header), like to omit
+  // params if they aren't used to compute the response. (This part not
+  // yet implemented)
+  key: string
+
+  // Child segments.
+  slots: null | {
+    [parallelRouteKey: string]: TreePrefetch
+  }
+
+  // Extra fields that only exist so we can reconstruct a FlightRouterState on
+  // the client. We may be able to unify TreePrefetch and FlightRouterState
+  // after some refactoring, but in the meantime it would be wasteful to add a
+  // bunch of new prefetch-only fields to FlightRouterState. So think of
+  // TreePrefetch as a superset of FlightRouterState.
+  extra: [segment: Segment, isRootLayout: boolean]
+}
+
+type SegmentPrefetch = {
+  rsc: React.ReactNode | null
+  loading: LoadingModuleData
+}
+
 export async function collectSegmentData(
-  routeTree: FlightRouterState,
+  flightRouterState: FlightRouterState,
   fullPageDataBuffer: Buffer,
+  staleTime: number,
   clientModules: ManifestNode,
   serverConsumerManifest: any
 ): Promise<Map<string, Buffer>> {
@@ -43,19 +77,19 @@ export async function collectSegmentData(
   } catch {}
 
   // A mutable map to collect the results as we traverse the route tree.
-  const segmentBufferMap = new Map<string, Buffer>()
+  const resultMap = new Map<string, Buffer>()
   // A mutable array to collect the promises for each segment stream, so that
   // they can run in parallel.
   const collectedTasks: Array<Promise<void>> = []
 
-  collectSegmentDataImpl(
-    routeTree,
+  const tree = collectSegmentDataImpl(
+    flightRouterState,
     fullPageDataBuffer,
     clientModules,
     serverConsumerManifest,
     [],
     '',
-    segmentBufferMap,
+    resultMap,
     collectedTasks
   )
 
@@ -64,7 +98,17 @@ export async function collectSegmentData(
   // to renderToReadableStream).
   await Promise.all(collectedTasks)
 
-  return segmentBufferMap
+  // Render the route tree to a special `/_tree` segment.
+  const treePrefetch: RootTreePrefetch = {
+    tree,
+    staleTime,
+  }
+  const treeStream = renderToReadableStream(treePrefetch, clientModules)
+  const routeBuffer = await streamToBuffer(treeStream)
+
+  resultMap.set('/_tree', routeBuffer)
+
+  return resultMap
 }
 
 function collectSegmentDataImpl(
@@ -76,7 +120,9 @@ function collectSegmentDataImpl(
   segmentPathStr: string,
   segmentBufferMap: Map<string, Buffer>,
   collectedTasks: Array<Promise<void>>
-): void {
+): TreePrefetch {
+  const slots: { [parallelRouteKey: string]: TreePrefetch } = {}
+
   const children = route[1]
   for (const parallelRouteKey in children) {
     const childRoute = children[parallelRouteKey]
@@ -88,7 +134,8 @@ function collectSegmentDataImpl(
       segmentPathStr +
       '/' +
       encodeChildSegmentAsFilesystemSafePathname(parallelRouteKey, childSegment)
-    collectSegmentDataImpl(
+
+    const childTree = collectSegmentDataImpl(
       childRoute,
       fullPageDataBuffer,
       clientModules,
@@ -98,6 +145,7 @@ function collectSegmentDataImpl(
       segmentBufferMap,
       collectedTasks
     )
+    slots[parallelRouteKey] = childTree
   }
 
   // Spawn a task to render the segment data to a stream.
@@ -111,6 +159,16 @@ function collectSegmentDataImpl(
       segmentBufferMap
     )
   )
+
+  // Metadata about the segment. Sent to the client as part of the
+  // tree prefetch.
+  const segment = route[0]
+  const isRootLayout = route[4]
+  return {
+    key: segmentPathStr,
+    slots,
+    extra: [segment, isRootLayout === true],
+  }
 }
 
 async function renderSegmentDataToStream(
@@ -161,11 +219,6 @@ async function renderSegmentDataToStream(
   }
 }
 
-type SegmentPrefetch = {
-  rsc: React.ReactNode | null
-  loading: LoadingModuleData
-}
-
 async function PickSegment({
   fullPageDataBuffer,
   serverConsumerManifest,

test/e2e/app-dir/ppr-navigations/simple/per-segment-prefetching.test.ts

@@ -14,6 +14,21 @@ describe('per segment prefetching', () => {
   // these types of requests. This tests that the server responds correctly.
   // TODO: Replace with e2e tests once more is implemented.
 
+  it('prefetch a route tree', async () => {
+    const response = await next.fetch('/en', {
+      headers: {
+        RSC: '1',
+        'Next-Router-Prefetch': '1',
+        'Next-Router-Segment-Prefetch': '/_tree',
+      },
+    })
+    expect(response.status).toBe(200)
+    const responseText = await response.text()
+    // This is a basic check to ensure that the name of an expected field is
+    // somewhere in the Flight stream.
+    expect(responseText).toInclude('"staleTime"')
+  })
+
   it('prefetch an individual segment', async () => {
     const response = await next.fetch('/en', {
       headers: {