[Segment Cache] Cancel prefetch on viewport exit (#74671)

We use an IntersectionObserver to prefetch links when they enter the
viewport. This updates the behavior to cancel the prefetch if the link
exits the viewport before it completes.

This can greatly reduce the amount of data transfer caused by
prefetching, however, the impact of this change will depend on the
user's network conditions. The faster the network conditions, the more
likely the link will have already been prefetched by the time the link
exits the screen.

We'll need a different strategy for limiting prefetch data transfer in
fast network conditions, perhaps by tracking and throttling the overall
bitrate.

packages/next/src/client/app-dir/link.tsx

@@ -15,6 +15,7 @@ import { warnOnce } from '../../shared/lib/utils/warn-once'
 import {
   type PrefetchTask,
   schedulePrefetchTask as scheduleSegmentPrefetchTask,
+  cancelPrefetchTask,
   bumpPrefetchTask,
 } from '../components/segment-cache/scheduler'
 import { getCurrentAppRouterState } from '../../shared/lib/router/action-queue'
@@ -201,8 +202,7 @@ export function unmountLinkInstance(element: HTMLAnchorElement | SVGAElement) {
     links.delete(element)
     const prefetchTask = instance.prefetchTask
     if (prefetchTask !== null) {
-      // TODO: In the Segment Cache implementation, cancel the prefetch task
-      // when the link is unmounted.
+      cancelPrefetchTask(prefetchTask)
     }
   }
   if (observer !== null) {
@@ -255,8 +255,15 @@ function rescheduleLinkPrefetch(instance: LinkInstance) {
   const existingPrefetchTask = instance.prefetchTask
 
   if (!instance.isVisible) {
-    // TODO: In the Segment Cache implementation, cancel the prefetch task when
-    // the link leaves the viewport.
+    // Cancel any in-progress prefetch task. (If it already finished then this
+    // is a no-op.)
+    if (existingPrefetchTask !== null) {
+      cancelPrefetchTask(existingPrefetchTask)
+    }
+    // We don't need to reset the prefetchTask to null upon cancellation; an
+    // old task object can be rescheduled with bumpPrefetchTask. This is a
+    // micro-optimization but also makes the code simpler (don't need to
+    // worry about whether an old task object is stale).
     return
   }
 

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

@@ -90,6 +90,11 @@ export type PrefetchTask = {
    */
   hasBackgroundWork: boolean
 
+  /**
+   * True if the prefetch was cancelled.
+   */
+  isCanceled: boolean
+
   /**
    * The index of the task in the heap's backing array. Used to efficiently
    * change the priority of a task by re-sifting it, which requires knowing
@@ -174,6 +179,7 @@ export function schedulePrefetchTask(
     hasBackgroundWork: false,
     includeDynamicData,
     sortId: sortIdCounter++,
+    isCanceled: false,
     _heapIndex: -1,
   }
   heapPush(taskHeap, task)
@@ -190,6 +196,16 @@ export function schedulePrefetchTask(
   return task
 }
 
+export function cancelPrefetchTask(task: PrefetchTask): void {
+  // Remove the prefetch task from the queue. If the task already completed,
+  // then this is a no-op.
+  //
+  // We must also explicitly mark the task as canceled so that a blocked task
+  // does not get added back to the queue when it's pinged by the network.
+  task.isCanceled = true
+  heapDelete(taskHeap, task)
+}
+
 export function bumpPrefetchTask(task: PrefetchTask): void {
   // Bump the prefetch task to the top of the queue, as if it were a fresh
   // task. This is essentially the same as canceling the task and scheduling
@@ -198,6 +214,9 @@ export function bumpPrefetchTask(task: PrefetchTask): void {
   // The primary use case is to increase the relative priority of a Link-
   // initated prefetch on hover.
 
+  // Un-cancel the task, in case it was previously canceled.
+  task.isCanceled = false
+
   // Assign a new sort ID. Higher sort IDs are higher priority.
   task.sortId = sortIdCounter++
   if (task._heapIndex !== -1) {
@@ -276,8 +295,9 @@ function onPrefetchConnectionClosed(): void {
 export function pingPrefetchTask(task: PrefetchTask) {
   // "Ping" a prefetch that's already in progress to notify it of new data.
   if (
+    // Check if prefetch was canceled.
+    task.isCanceled ||
     // Check if prefetch is already queued.
-    // TODO: Check if task was canceled, too
     task._heapIndex !== -1
   ) {
     return
@@ -1068,6 +1088,21 @@ function heapPop(heap: Array<PrefetchTask>): PrefetchTask | null {
   return first
 }
 
+function heapDelete(heap: Array<PrefetchTask>, node: PrefetchTask): void {
+  const index = node._heapIndex
+  if (index !== -1) {
+    node._heapIndex = -1
+    if (heap.length !== 0) {
+      const last = heap.pop() as PrefetchTask
+      if (last !== node) {
+        heap[index] = last
+        last._heapIndex = index
+        heapSiftDown(heap, last, index)
+      }
+    }
+  }
+}
+
 function heapResift(heap: Array<PrefetchTask>, node: PrefetchTask): void {
   const index = node._heapIndex
   if (index !== -1) {

test/e2e/app-dir/segment-cache/basic/segment-cache-basic.test.ts

@@ -47,15 +47,11 @@ describe('segment cache (basic tests)', () => {
           `"<div><div data-streaming-text-static="Static in nav">Static in nav</div><div data-streaming-text-dynamic="Dynamic in nav">Dynamic in nav</div></div>"`
         )
       },
-      // When the outer act scope exits, the blocked prefetches are allowed
-      // to continue.
-      // TODO: As an optimization, in the case where a fully static page is
-      // returned during a dynamic response, we should populate the prefetch
-      // cache with the static data. Then we wouldn't have to prefetch it again
-      // here, because it would already be cached. Only works for fully static
-      // pages because in a partial response we don't know which parts are
-      // static versus dynamic.
-      { includes: 'Static in nav' }
+      // Although the blocked prefetches are allowed to continue when we exit
+      // the outer `act` scope, they were canceled when we navigated to the new
+      // page. So there should be no additional requests in the outer
+      // `act` scope.
+      'no-requests'
     )
   })
 

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

@@ -58,4 +58,95 @@ describe('segment cache prefetch scheduling', () => {
       ]
     )
   })
+
+  it(
+    'cancels a viewport-initiated prefetch if the link leaves the viewport ' +
+      'before it finishes',
+    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, but block the responses from
+          // reaching the client. Because the router limits the number of
+          // concurrent prefetches, not all the links will start prefetching —
+          // some of them will remain in the queue, waiting for additional
+          // network bandwidth. This test demonstrates that those prefetches
+          // will be canceled on viewport exit, too.
+          await act(async () => {
+            await checkbox.click()
+          }, 'block')
+
+          // Before the prefetch finishes, click the checkbox again to hide
+          // the link.
+          await checkbox.click()
+        },
+        // When the outer `act` scope finishes, the route tree prefetch will
+        // continue. Normally when the router is done prefetching the route
+        // tree, it will proceed to prefetching the segments. However, since
+        // the link is no longer visible, it should stop prefetching.
+        //
+        // Assert that no additional network requests are initiated in this
+        // outer scope. If this fails, it suggests that the prefetches were not
+        // canceled when the links left the viewport.
+        'no-requests'
+      )
+    }
+  )
+
+  it("reschedules a link's prefetch when it re-enters the viewport", 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, but block the responses from
+        // reaching the client. Because the router limits the number of
+        // concurrent prefetches, not all the links will start prefetching —
+        // some of them will remain in the queue, waiting for additional
+        // network bandwidth. This test demonstrates that those prefetches
+        // will be canceled on viewport exit, too.
+        await act(async () => {
+          await checkbox.click()
+        }, 'block')
+
+        // Before the prefetch finishes, click the checkbox again to hide
+        // the link.
+        await checkbox.click()
+      },
+      // When the outer `act` scope finishes, the route tree prefetch will
+      // continue. Normally when the router is done prefetching the route
+      // tree, it will proceed to prefetching the segments. However, since
+      // the link is no longer visible, it should stop prefetching.
+      //
+      // Assert that no additional network requests are initiated in this
+      // outer scope. If this fails, it suggests that the prefetches were not
+      // canceled when the links left the viewport.
+      'no-requests'
+    )
+
+    // Now we'll reveal the links again to verify that the prefetch tasks are
+    // rescheduled, after having been canceled.
+    await act(
+      async () => {
+        await checkbox.click()
+      },
+      // Don't need to assert on all the prefetch responses. I picked an
+      // arbitrary one.
+      { includes: 'Content of page 5' }
+    )
+  })
 })