Add onInvalidate option to router.prefetch

packages/next/src/client/components/app-router-instance.ts

@@ -321,7 +321,8 @@ export const publicAppRouterInstance: AppRouterInstance = {
           href,
           actionQueue.state.nextUrl,
           actionQueue.state.tree,
-          options?.kind === PrefetchKind.FULL
+          options?.kind === PrefetchKind.FULL,
+          options?.onInvalidate ?? null
         )
       }
     : (href: string, options?: PrefetchOptions) => {

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

@@ -3,7 +3,7 @@ import type { AppRouterInstance } from '../../shared/lib/app-router-context.shar
 import { getCurrentAppRouterState } from './app-router-instance'
 import { createPrefetchURL } from './app-router'
 import { PrefetchKind } from './router-reducer/router-reducer-types'
-import { getCurrentCacheVersion } from './segment-cache'
+import { isPrefetchTaskDirty } from './segment-cache'
 import { createCacheKey } from './segment-cache'
 import {
   type PrefetchTask,
@@ -28,13 +28,9 @@ type LinkOrFormInstanceShared = {
   wasHoveredOrTouched: boolean
 
   // The most recently initiated prefetch task. It may or may not have
-  // already completed.  The same prefetch task object can be reused across
+  // already completed. The same prefetch task object can be reused across
   // multiple prefetches of the same link.
   prefetchTask: PrefetchTask | null
-
-  // The cache version at the time the task was initiated. This is used to
-  // determine if the cache was invalidated since the task was initiated.
-  cacheVersion: number
 }
 
 export type FormInstance = LinkOrFormInstanceShared & {
@@ -158,7 +154,6 @@ export function mountLinkInstance(
         isVisible: false,
         wasHoveredOrTouched: false,
         prefetchTask: null,
-        cacheVersion: -1,
         prefetchHref: prefetchURL.href,
         setOptimisticLinkStatus,
       }
@@ -176,7 +171,6 @@ export function mountLinkInstance(
     isVisible: false,
     wasHoveredOrTouched: false,
     prefetchTask: null,
-    cacheVersion: -1,
     prefetchHref: null,
     setOptimisticLinkStatus,
   }
@@ -203,7 +197,6 @@ export function mountFormInstance(
     isVisible: false,
     wasHoveredOrTouched: false,
     prefetchTask: null,
-    cacheVersion: -1,
     prefetchHref: prefetchURL.href,
     setOptimisticLinkStatus: null,
   }
@@ -322,7 +315,8 @@ function rescheduleLinkPrefetch(instance: PrefetchableInstance) {
         cacheKey,
         treeAtTimeOfPrefetch,
         instance.kind === PrefetchKind.FULL,
-        priority
+        priority,
+        null
       )
     } else {
       // We already have an old task object that we can reschedule. This is
@@ -334,10 +328,6 @@ function rescheduleLinkPrefetch(instance: PrefetchableInstance) {
         priority
       )
     }
-
-    // Keep track of the cache version at the time the prefetch was requested.
-    // This is used to check if the prefetch is stale.
-    instance.cacheVersion = getCurrentCacheVersion()
   }
 }
 
@@ -352,15 +342,9 @@ export function pingVisibleLinks(
   // This is called when the Next-Url or the base tree changes, since those
   // may affect the result of a prefetch task. It's also called after a
   // cache invalidation.
-  const currentCacheVersion = getCurrentCacheVersion()
   for (const instance of prefetchableAndVisible) {
     const task = instance.prefetchTask
-    if (
-      task !== null &&
-      instance.cacheVersion === currentCacheVersion &&
-      task.key.nextUrl === nextUrl &&
-      task.treeAtTimeOfPrefetch === tree
-    ) {
+    if (task !== null && !isPrefetchTaskDirty(task, nextUrl, tree)) {
       // The cache has not been invalidated, and none of the inputs have
       // changed. Bail out.
       continue
@@ -378,9 +362,9 @@ export function pingVisibleLinks(
       cacheKey,
       tree,
       instance.kind === PrefetchKind.FULL,
-      priority
+      priority,
+      null
     )
-    instance.cacheVersion = getCurrentCacheVersion()
   }
 }
 

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

@@ -28,6 +28,7 @@ import {
 } from '../router-reducer/fetch-server-response'
 import {
   pingPrefetchTask,
+  isPrefetchTaskDirty,
   type PrefetchTask,
   type PrefetchSubtaskResult,
 } from './scheduler'
@@ -245,6 +246,14 @@ let segmentCacheLru = createLRU<SegmentCacheEntry>(
   onSegmentLRUEviction
 )
 
+// All invalidation listeners for the whole cache are tracked in single set.
+// Since we don't yet support tag or path-based invalidation, there's no point
+// tracking them any more granularly than this. Once we add granular
+// invalidation, that may change, though generally the model is to just notify
+// the listeners and allow the caller to poll the prefetch cache with a new
+// prefetch task if desired.
+let invalidationListeners: Set<PrefetchTask> | null = null
+
 // Incrementing counter used to track cache invalidations.
 let currentCacheVersion = 0
 
@@ -276,6 +285,65 @@ export function revalidateEntireCache(
 
   // Prefetch all the currently visible links again, to re-fill the cache.
   pingVisibleLinks(nextUrl, tree)
+
+  // Similarly, notify all invalidation listeners (i.e. those passed to
+  // `router.prefetch(onInvalidate)`), so they can trigger a new prefetch
+  // if needed.
+  pingInvalidationListeners(nextUrl, tree)
+}
+
+function attachInvalidationListener(task: PrefetchTask): void {
+  // This function is called whenever a prefetch task reads a cache entry. If
+  // the task has an onInvalidate function associated with it — i.e. the one
+  // optionally passed to router.prefetch(onInvalidate) — then we attach that
+  // listener to the every cache entry that the task reads. Then, if an entry
+  // is invalidated, we call the function.
+  if (task.onInvalidate !== null) {
+    if (invalidationListeners === null) {
+      invalidationListeners = new Set([task])
+    } else {
+      invalidationListeners.add(task)
+    }
+  }
+}
+
+function notifyInvalidationListener(task: PrefetchTask): void {
+  const onInvalidate = task.onInvalidate
+  if (onInvalidate !== null) {
+    // Clear the callback from the task object to guarantee it's not called more
+    // than once.
+    task.onInvalidate = null
+
+    // This is a user-space function, so we must wrap in try/catch.
+    try {
+      onInvalidate()
+    } catch (error) {
+      if (typeof reportError === 'function') {
+        reportError(error)
+      } else {
+        console.error(error)
+      }
+    }
+  }
+}
+
+export function pingInvalidationListeners(
+  nextUrl: string | null,
+  tree: FlightRouterState
+): void {
+  // The rough equivalent of pingVisibleLinks, but for onInvalidate callbacks.
+  // This is called when the Next-Url or the base tree changes, since those
+  // may affect the result of a prefetch task. It's also called after a
+  // cache invalidation.
+  if (invalidationListeners !== null) {
+    const tasks = invalidationListeners
+    invalidationListeners = null
+    for (const task of tasks) {
+      if (isPrefetchTaskDirty(task, nextUrl, tree)) {
+        notifyInvalidationListener(task)
+      }
+    }
+  }
 }
 
 export function readExactRouteCacheEntry(
@@ -445,6 +513,8 @@ export function readOrCreateRouteCacheEntry(
   now: number,
   task: PrefetchTask
 ): RouteCacheEntry {
+  attachInvalidationListener(task)
+
   const key = task.key
   const existingEntry = readRouteCacheEntry(now, key)
   if (existingEntry !== null) {

packages/next/src/client/components/segment-cache-impl/prefetch.ts

@@ -9,17 +9,27 @@ import { PrefetchPriority } from '../segment-cache'
  * @param href - The URL to prefetch. Typically this will come from a <Link>,
  * or router.prefetch. It must be validated before we attempt to prefetch it.
  * @param nextUrl - A special header used by the server for interception routes.
- * Roughly  corresponds to the current URL.
+ * Roughly corresponds to the current URL.
  * @param treeAtTimeOfPrefetch - The FlightRouterState at the time the prefetch
  * was requested. This is only used when PPR is disabled.
  * @param includeDynamicData - Whether to prefetch dynamic data, in addition to
  * static data. This is used by <Link prefetch={true}>.
+ * @param onInvalidate - A callback that will be called when the prefetch cache
+ * When called, it signals to the listener that the data associated with the
+ * prefetch may have been invalidated from the cache. This is not a live
+ * subscription — it's called at most once per `prefetch` call. The only
+ * supported use case is to trigger a new prefetch inside the listener, if
+ * desired. It also may be called even in cases where the associated data is
+ * still cached. Prefetching is a poll-based (pull) operation, not an event-
+ * based (push) one. Rather than subscribe to specific cache entries, you
+ * occasionally poll the prefetch cache to check if anything is missing.
  */
 export function prefetch(
   href: string,
   nextUrl: string | null,
   treeAtTimeOfPrefetch: FlightRouterState,
-  includeDynamicData: boolean
+  includeDynamicData: boolean,
+  onInvalidate: null | (() => void)
 ) {
   const url = createPrefetchURL(href)
   if (url === null) {
@@ -31,6 +41,7 @@ export function prefetch(
     cacheKey,
     treeAtTimeOfPrefetch,
     includeDynamicData,
-    PrefetchPriority.Default
+    PrefetchPriority.Default,
+    onInvalidate
   )
 }

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

@@ -26,7 +26,7 @@ import {
   getSegmentKeypathForTask,
 } from './cache'
 import type { RouteCacheKey } from './cache-key'
-import { PrefetchPriority } from '../segment-cache'
+import { getCurrentCacheVersion, PrefetchPriority } from '../segment-cache'
 
 const scheduleMicrotask =
   typeof queueMicrotask === 'function'
@@ -50,6 +50,12 @@ export type PrefetchTask = {
    */
   treeAtTimeOfPrefetch: FlightRouterState
 
+  /**
+   * The cache version at the time the task was initiated. This is used to
+   * determine if the cache was invalidated since the task was initiated.
+   */
+  cacheVersion: number
+
   /**
    * Whether to prefetch dynamic data, in addition to static data. This is
    * used by <Link prefetch={true}>.
@@ -105,6 +111,11 @@ export type PrefetchTask = {
    */
   isCanceled: boolean
 
+  /**
+   * The callback passed to `router.prefetch`, if given.
+   */
+  onInvalidate: null | (() => void)
+
   /**
    * 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
@@ -182,18 +193,21 @@ export function schedulePrefetchTask(
   key: RouteCacheKey,
   treeAtTimeOfPrefetch: FlightRouterState,
   includeDynamicData: boolean,
-  priority: PrefetchPriority
+  priority: PrefetchPriority,
+  onInvalidate: null | (() => void)
 ): PrefetchTask {
   // Spawn a new prefetch task
   const task: PrefetchTask = {
     key,
     treeAtTimeOfPrefetch,
+    cacheVersion: getCurrentCacheVersion(),
     priority,
     phase: PrefetchPhase.RouteTree,
     hasBackgroundWork: false,
     includeDynamicData,
     sortId: sortIdCounter++,
     isCanceled: false,
+    onInvalidate,
     _heapIndex: -1,
   }
   heapPush(taskHeap, task)
@@ -254,6 +268,24 @@ export function reschedulePrefetchTask(
   ensureWorkIsScheduled()
 }
 
+export function isPrefetchTaskDirty(
+  task: PrefetchTask,
+  nextUrl: string | null,
+  tree: FlightRouterState
+): boolean {
+  // This is used to quickly bail out of a prefetch task if the result is
+  // guaranteed to not have changed since the task was initiated. This is
+  // strictly an optimization — theoretically, if it always returned true, no
+  // behavior should change because a full prefetch task will effectively
+  // perform the same checks.
+  const currentCacheVersion = getCurrentCacheVersion()
+  return (
+    task.cacheVersion !== currentCacheVersion ||
+    task.treeAtTimeOfPrefetch !== tree ||
+    task.key.nextUrl !== nextUrl
+  )
+}
+
 function ensureWorkIsScheduled() {
   if (didScheduleMicrotask || !hasNetworkBandwidth()) {
     // Either we already scheduled a task to process the queue, or there are
@@ -344,6 +376,8 @@ function processQueueInMicrotask() {
   // Process the task queue until we run out of network bandwidth.
   let task = heapPeek(taskHeap)
   while (task !== null && hasNetworkBandwidth()) {
+    task.cacheVersion = getCurrentCacheVersion()
+
     const route = readOrCreateRouteCacheEntry(now, task)
     const exitStatus = pingRootRouteTree(now, task, route)
 

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

@@ -85,6 +85,15 @@ export const reschedulePrefetchTask: typeof import('./segment-cache-impl/schedul
       }
     : notEnabled
 
+export const isPrefetchTaskDirty: typeof import('./segment-cache-impl/scheduler').isPrefetchTaskDirty =
+  process.env.__NEXT_CLIENT_SEGMENT_CACHE
+    ? function (...args) {
+        return require('./segment-cache-impl/scheduler').isPrefetchTaskDirty(
+          ...args
+        )
+      }
+    : notEnabled
+
 export const createCacheKey: typeof import('./segment-cache-impl/cache-key').createCacheKey =
   process.env.__NEXT_CLIENT_SEGMENT_CACHE
     ? function (...args) {

packages/next/src/shared/lib/app-router-context.shared-runtime.ts

@@ -122,6 +122,7 @@ export interface NavigateOptions {
 
 export interface PrefetchOptions {
   kind: PrefetchKind
+  onInvalidate?: () => void
 }
 
 export interface AppRouterInstance {

test/e2e/app-dir/segment-cache/revalidation/app/page.tsx

@@ -1,5 +1,9 @@
 import { revalidatePath, revalidateTag } from 'next/cache'
-import { LinkAccordion, FormAccordion } from '../components/link-accordion'
+import {
+  LinkAccordion,
+  FormAccordion,
+  ManualPrefetchLinkAccordion,
+} from '../components/link-accordion'
 import Link from 'next/link'
 
 export default async function Page() {
@@ -36,6 +40,12 @@ export default async function Page() {
             Form pointing to target page with prefetching enabled
           </FormAccordion>
         </li>
+        <li>
+          <ManualPrefetchLinkAccordion href="/greeting">
+            Manual link (router.prefetch) to target page with prefetching
+            enabled
+          </ManualPrefetchLinkAccordion>
+        </li>
         <li>
           <Link prefetch={false} href="/greeting">
             Link to target with prefetching disabled