[Segment Cache] Prioritize route trees over segments (#75213)

When processing the prefetch queue, we should prioritize fetching the
route trees for all the links before we fetch any of the segments.

The reason the route tree is more important is because it tells us the
structure of the target page; from this alone, we can determine which
segments are already cached and skip them during the navigation. Whereas
if the route tree is missing at the time of the navigation, we cannot
instruct the server to skip over prefetched segments, because we don't
know which ones to skip.

To implement this, I added a `phase` property to the PrefetchTask type.
This is essentially a secondary priority field that changes as the task
progresses. There are two phases: `RouteTree` and `Segments`. Any task
in the `RouteTree` phase has higher priority than tasks in the
`Segments` phase.

Another way to implement this would be to add a secondary queue for
segment prefetches. This would require reference counting to clean up
segments when there are no more tasks that depend on them. The ref count
could also be used as an additional prioritization heuristic: shared
layouts with many dependent tasks could be prioritized over segments
with fewer. However, I'm not sure the benefits are worth the extra
code/complexity, so I'm starting with this simpler approach.

packages/next/src/client/components/segment-cache/scheduler.ts

@@ -83,6 +83,14 @@ export type PrefetchTask = {
    */
   priority: PrefetchPriority
 
+  /**
+   * The phase of the task. Tasks are split into multiple phases so that their
+   * priority can be adjusted based on what kind of work they're doing.
+   * Concretely, prefetching the route tree is higher priority than prefetching
+   * segment data.
+   */
+  phase: PrefetchPhase
+
   /**
    * Temporary state for tracking the currently running task. This is currently
    * used to track whether a task deferred some work to run background at
@@ -149,6 +157,16 @@ export const enum PrefetchPriority {
   Background = 0,
 }
 
+/**
+ * Prefetch tasks are processed in two phases: first the route tree is fetched,
+ * then the segments. We use this to priortize tasks that have not yet fetched
+ * the route tree.
+ */
+const enum PrefetchPhase {
+  RouteTree = 1,
+  Segments = 0,
+}
+
 export type PrefetchSubtaskResult<T> = {
   /**
    * A promise that resolves when the network connection is closed.
@@ -190,6 +208,7 @@ export function schedulePrefetchTask(
     key,
     treeAtTimeOfPrefetch,
     priority,
+    phase: PrefetchPhase.RouteTree,
     hasBackgroundWork: false,
     includeDynamicData,
     sortId: sortIdCounter++,
@@ -359,7 +378,12 @@ function processQueueInMicrotask() {
         task = heapPeek(taskHeap)
         continue
       case PrefetchTaskExitStatus.Done:
-        if (hasBackgroundWork) {
+        if (task.phase === PrefetchPhase.RouteTree) {
+          // Finished prefetching the route tree. Proceed to prefetching
+          // the segments.
+          task.phase = PrefetchPhase.Segments
+          heapResift(taskHeap, task)
+        } else if (hasBackgroundWork) {
           // The task spawned additional background work. Reschedule the task
           // at background priority.
           task.priority = PrefetchPriority.Background
@@ -446,6 +470,10 @@ function pingRootRouteTree(
       return PrefetchTaskExitStatus.Done
     }
     case EntryStatus.Fulfilled: {
+      if (task.phase !== PrefetchPhase.Segments) {
+        // Do not prefetch segment data until we've entered the segment phase.
+        return PrefetchTaskExitStatus.Done
+      }
       // Recursively fill in the segment tree.
       if (!hasNetworkBandwidth()) {
         // Stop prefetching segments until there's more bandwidth.
@@ -1058,8 +1086,15 @@ function compareQueuePriority(a: PrefetchTask, b: PrefetchTask) {
     return priorityDiff
   }
 
-  // sortId is an incrementing counter assigned to prefetches. We want to
-  // process the newest prefetches first.
+  // If the priority is the same, check which phase the prefetch is in — is it
+  // prefetching the route tree, or the segments? Route trees are prioritized.
+  const phaseDiff = b.phase - a.phase
+  if (phaseDiff !== 0) {
+    return phaseDiff
+  }
+
+  // Finally, check the insertion order. `sortId` is an incrementing counter
+  // assigned to prefetches. We want to process the newest prefetches first.
   return b.sortId - a.sortId
 }
 

test/e2e/app-dir/segment-cache/prefetch-scheduling/app/cancellation/[pageNumber]/page.tsx

@@ -4,6 +4,19 @@ type Params = {
   pageNumber: string
 }
 
+export async function generateViewport({
+  params,
+}: {
+  params: Promise<Params>
+}) {
+  const { pageNumber } = await params
+  return {
+    // Put the page number into the media query. This is just a trick to allow
+    // the test to detect when the viewport for this page has been prefetched.
+    themeColor: [{ media: `(min-width: ${pageNumber}px)`, color: 'light' }],
+  }
+}
+
 async function Content({ params }: { params: Promise<Params> }) {
   const { pageNumber } = await params
   return 'Content of page ' + pageNumber
@@ -23,7 +36,7 @@ export default async function LinkCancellationTargetPage({
 
 export async function generateStaticParams(): Promise<Array<Params>> {
   const result: Array<Params> = []
-  for (let n = 1; n <= 1; n++) {
+  for (let n = 1; n <= 10; n++) {
     result.push({ pageNumber: n.toString() })
   }
   return result

test/e2e/app-dir/segment-cache/prefetch-scheduling/prefetch-scheduling.test.ts

@@ -52,6 +52,48 @@ describe('segment cache prefetch scheduling', () => {
     )
   })
 
+  it('prioritizes prefetching the route trees before the segments', async () => {
+    let act: ReturnType<typeof createRouterAct>
+    const browser = await next.browser('/cancellation', {
+      beforePageLoad(p: Playwright.Page) {
+        act = createRouterAct(p)
+      },
+    })
+
+    const checkbox = await browser.elementByCss('input[type="checkbox"]')
+
+    await act(async () => {
+      // Reveal the links to start prefetching
+      await checkbox.click()
+    }, [
+      // Assert on the order that the prefetches requests are
+      // initiated. We don't need to assert on every single prefetch response;
+      // this will only check the order of the ones that we've listed.
+      //
+      // To detect when the route tree is prefetched, we check for a string
+      // that is known to be present in the target page's viewport config
+      // (which is included in the route tree response). In this test app, the
+      // page number is used in the media query of the theme color. E.g. for
+      // page 1, the viewport includes:
+      //
+      // <meta name="theme-color" media="(min-width: 1px)" content="light"/>
+
+      // First we should prefetch all the route trees:
+      { includes: '(min-width: 7px)' },
+      { includes: '(min-width: 6px)' },
+      { includes: '(min-width: 5px)' },
+      { includes: '(min-width: 4px)' },
+      { includes: '(min-width: 3px)' },
+
+      // Then we should prefetch the segments:
+      { includes: 'Content of page 7' },
+      { includes: 'Content of page 6' },
+      { includes: 'Content of page 5' },
+      { includes: 'Content of page 4' },
+      { includes: 'Content of page 3' },
+    ])
+  })
+
   it(
     'even on mouseexit, any link that was previously hovered is prioritized ' +
       'over links that were never hovered at all',