[Segment Cache] Search param fallback handling

This implements search param handling in the Segment Cache.

When a cache entry is fetched via PPR, search params are treated as
dynamic data. We can share the same cache entry for every set of search
params for a given page segment.

By contrast, when a cache entry is fetched using a dynamic request (e.g.
using `<Link prefetch={true}`), the result may vary by the search param
values. So we must include the search params in the cache key.

During a navigation, we will first check for the more specific cache
entry (i.e. one that includes search params), and then fallback to the
PPR version if no such entry exists.

Author: acdlite

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

@@ -3,13 +3,17 @@ type Opaque<K, T> = T & { __brand: K }
 
 // Only functions in this module should be allowed to create CacheKeys.
 export type NormalizedHref = Opaque<'NormalizedHref', string>
+export type NormalizedSearch = Opaque<'NormalizedSearch', string>
 export type NormalizedNextUrl = Opaque<'NormalizedNextUrl', string>
 
 export type RouteCacheKey = Opaque<
   'RouteCacheKey',
   {
     href: NormalizedHref
+    search: NormalizedSearch
     nextUrl: NormalizedNextUrl | null
+
+    // TODO: Eventually the dynamic params will be added here, too.
   }
 >
 
@@ -18,20 +22,10 @@ export function createCacheKey(
   nextUrl: string | null
 ): RouteCacheKey {
   const originalUrl = new URL(originalHref)
-
-  // TODO: As of now, we never include search params in the cache key because
-  // per-segment prefetch requests are always static, and cannot contain search
-  // params. But to support <Link prefetch={true}>, we will sometimes populate
-  // the cache with dynamic data, so this will have to change.
-  originalUrl.search = ''
-
-  const normalizedHref = originalUrl.href as NormalizedHref
-  const normalizedNextUrl = nextUrl as NormalizedNextUrl | null
-
   const cacheKey = {
-    href: normalizedHref,
-    nextUrl: normalizedNextUrl,
+    href: originalHref as NormalizedHref,
+    search: originalUrl.search as NormalizedSearch,
+    nextUrl: nextUrl as NormalizedNextUrl | null,
   } as RouteCacheKey
-
   return cacheKey
 }

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

@@ -37,6 +37,7 @@ import { createHrefFromUrl } from '../router-reducer/create-href-from-url'
 import type {
   NormalizedHref,
   NormalizedNextUrl,
+  NormalizedSearch,
   RouteCacheKey,
 } from './cache-key'
 import { createTupleMap, type TupleMap, type Prefix } from './tuple-map'
@@ -52,6 +53,7 @@ import type {
 } from '../../../server/app-render/types'
 import { normalizeFlightData } from '../../flight-data-helpers'
 import { STATIC_STALETIME_MS } from '../router-reducer/prefetch-cache-utils'
+import { PAGE_SEGMENT_KEY } from '../../../shared/lib/segment'
 
 // A note on async/await when working in the prefetch cache:
 //
@@ -156,9 +158,9 @@ type SegmentCacheEntryShared = {
   revalidating: SegmentCacheEntry | null
 
   // LRU-related fields
-  key: null | string
-  next: null | RouteCacheEntry
-  prev: null | RouteCacheEntry
+  keypath: null | Prefix<SegmentCacheKeypath>
+  next: null | SegmentCacheEntry
+  prev: null | SegmentCacheEntry
   size: number
 }
 
@@ -225,9 +227,9 @@ let routeCacheLru = createLRU<RouteCacheEntry>(
   onRouteLRUEviction
 )
 
-// TODO: We may eventually store segment entries in a tuple map, too, to
-// account for search params.
-let segmentCacheMap = new Map<string, SegmentCacheEntry>()
+type SegmentCacheKeypath = [string, NormalizedSearch]
+let segmentCacheMap: TupleMap<SegmentCacheKeypath, SegmentCacheEntry> =
+  createTupleMap()
 // NOTE: Segments and Route entries are managed by separate LRUs. We could
 // combine them into a single LRU, but because they are separate types, we'd
 // need to wrap each one in an extra LRU node (to maintain monomorphism, at the
@@ -252,7 +254,7 @@ export function revalidateEntireCache() {
   // correctly: background revalidations. See note in `upsertSegmentEntry`.
   routeCacheMap = createTupleMap()
   routeCacheLru = createLRU(maxRouteLruSize, onRouteLRUEviction)
-  segmentCacheMap = new Map()
+  segmentCacheMap = createTupleMap()
   segmentCacheLru = createLRU(maxSegmentLruSize, onSegmentLRUEviction)
 }
 
@@ -296,12 +298,59 @@ export function readRouteCacheEntry(
   return readExactRouteCacheEntry(now, key.href, key.nextUrl)
 }
 
+export function getSegmentKeypathForTask(
+  task: PrefetchTask,
+  path: string
+): Prefix<SegmentCacheKeypath> {
+  // When a prefetch includes dynamic data, the search params are included
+  // in the result, so we must include the search string in the segment
+  // cache key. (Note that this is true even if the search string is empty.)
+  //
+  // If we're fetching using PPR, we do not need to include the search params in
+  // the cache key, because the search params are treated as dynamic data. The
+  // cache entry is valid for all possible search param values.
+  return task.includeDynamicData && path.endsWith('/' + PAGE_SEGMENT_KEY)
+    ? [path, task.key.search]
+    : [path]
+}
+
 export function readSegmentCacheEntry(
   now: number,
+  routeCacheKey: RouteCacheKey,
   path: string
 ): SegmentCacheEntry | null {
-  const existingEntry = segmentCacheMap.get(path)
-  if (existingEntry !== undefined) {
+  if (!path.endsWith('/' + PAGE_SEGMENT_KEY)) {
+    // Fast path. Search params only exist on page segments.
+    return readExactSegmentCacheEntry(now, [path])
+  }
+
+  // Page segments may or may not contain search params. If they were prefetched
+  // using a dynamic request, then we will have an entry with search params.
+  // Check for that case first.
+  const entryWithSearchParams = readExactSegmentCacheEntry(now, [
+    path,
+    routeCacheKey.search,
+  ])
+  if (entryWithSearchParams !== null) {
+    return entryWithSearchParams
+  }
+
+  // If we did not find an entry with the given search params, check for a
+  // "fallback" entry, where the search params are treated as dynamic data. This
+  // is the common case because PPR/static prerenders always treat search params
+  // as dynamic.
+  //
+  // See corresponding logic in `getSegmentKeypathForTask`.
+  const entryWithoutSearchParams = readExactSegmentCacheEntry(now, [path])
+  return entryWithoutSearchParams
+}
+
+function readExactSegmentCacheEntry(
+  now: number,
+  keypath: Prefix<SegmentCacheKeypath>
+): SegmentCacheEntry | null {
+  const existingEntry = segmentCacheMap.get(keypath)
+  if (existingEntry !== null) {
     // Check if the entry is stale
     if (existingEntry.staleAt > now) {
       // Reuse the existing entry.
@@ -315,14 +364,18 @@ export function readSegmentCacheEntry(
       const revalidatingEntry = existingEntry.revalidating
       if (revalidatingEntry !== null) {
         // There's a revalidation in progress. Upsert it.
-        const upsertedEntry = upsertSegmentEntry(now, path, revalidatingEntry)
+        const upsertedEntry = upsertSegmentEntry(
+          now,
+          keypath,
+          revalidatingEntry
+        )
         if (upsertedEntry !== null && upsertedEntry.staleAt > now) {
           // We can use the upserted revalidation entry.
           return upsertedEntry
         }
       } else {
         // Evict the stale entry from the cache.
-        deleteSegmentFromCache(existingEntry, path)
+        deleteSegmentFromCache(existingEntry, keypath)
       }
     }
   }
@@ -415,20 +468,22 @@ export function readOrCreateRouteCacheEntry(
  */
 export function readOrCreateSegmentCacheEntry(
   now: number,
+  task: PrefetchTask,
   // TODO: Don't need to pass the whole route. Just `staleAt`.
   route: FulfilledRouteCacheEntry,
   path: string
 ): SegmentCacheEntry {
-  const existingEntry = readSegmentCacheEntry(now, path)
+  const keypath = getSegmentKeypathForTask(task, path)
+  const existingEntry = readExactSegmentCacheEntry(now, keypath)
   if (existingEntry !== null) {
     return existingEntry
   }
   // Create a pending entry and add it to the cache.
   const pendingEntry = createDetachedSegmentCacheEntry(route.staleAt)
-  segmentCacheMap.set(path, pendingEntry)
+  segmentCacheMap.set(keypath, pendingEntry)
   // Stash the keypath on the entry so we know how to remove it from the map
   // if it gets evicted from the LRU.
-  pendingEntry.key = path
+  pendingEntry.keypath = keypath
   segmentCacheLru.put(pendingEntry)
   return pendingEntry
 }
@@ -460,7 +515,7 @@ export function readOrCreateRevalidatingSegmentEntry(
 
 export function upsertSegmentEntry(
   now: number,
-  segmentKeyPath: string,
+  keypath: Prefix<SegmentCacheKeypath>,
   candidateEntry: SegmentCacheEntry
 ): SegmentCacheEntry | null {
   // We have a new entry that has not yet been inserted into the cache. Before
@@ -469,7 +524,7 @@ export function upsertSegmentEntry(
   // TODO: We should not upsert an entry if its key was invalidated in the time
   // since the request was made. We can do that by passing the "owner" entry to
   // this function and confirming it's the same as `existingEntry`.
-  const existingEntry = readSegmentCacheEntry(now, segmentKeyPath)
+  const existingEntry = readExactSegmentCacheEntry(now, keypath)
   if (existingEntry !== null) {
     if (candidateEntry.isPartial && !existingEntry.isPartial) {
       // Don't replace a full segment with a partial one. A case where this
@@ -488,12 +543,12 @@ export function upsertSegmentEntry(
       return null
     }
     // Evict the existing entry from the cache.
-    deleteSegmentFromCache(existingEntry, segmentKeyPath)
+    deleteSegmentFromCache(existingEntry, keypath)
   }
-  segmentCacheMap.set(segmentKeyPath, candidateEntry)
+  segmentCacheMap.set(keypath, candidateEntry)
   // Stash the keypath on the entry so we know how to remove it from the map
   // if it gets evicted from the LRU.
-  candidateEntry.key = segmentKeyPath
+  candidateEntry.keypath = keypath
   segmentCacheLru.put(candidateEntry)
   return candidateEntry
 }
@@ -514,7 +569,7 @@ export function createDetachedSegmentCacheEntry(
     promise: null,
 
     // LRU-related fields
-    key: null,
+    keypath: null,
     next: null,
     prev: null,
     size: 0,
@@ -541,9 +596,12 @@ function deleteRouteFromCache(
   routeCacheLru.delete(entry)
 }
 
-function deleteSegmentFromCache(entry: SegmentCacheEntry, key: string): void {
+function deleteSegmentFromCache(
+  entry: SegmentCacheEntry,
+  keypath: Prefix<SegmentCacheKeypath>
+): void {
   cancelEntryListeners(entry)
-  segmentCacheMap.delete(key)
+  segmentCacheMap.delete(keypath)
   segmentCacheLru.delete(entry)
   clearRevalidatingSegmentFromOwner(entry)
 }
@@ -581,11 +639,11 @@ function onRouteLRUEviction(entry: RouteCacheEntry): void {
 
 function onSegmentLRUEviction(entry: SegmentCacheEntry): void {
   // The LRU evicted this entry. Remove it from the map.
-  const key = entry.key
-  if (key !== null) {
-    entry.key = null
+  const keypath = entry.keypath
+  if (keypath !== null) {
+    entry.keypath = null
     cancelEntryListeners(entry)
-    segmentCacheMap.delete(key)
+    segmentCacheMap.delete(keypath)
   }
 }
 
@@ -765,9 +823,25 @@ function convertFlightRouterStateToRouteTree(
     }
   }
 
+  // The navigation implementation expects the search params to be included
+  // in the segment. However, in the case of a static response, the search
+  // params are omitted. So the client needs to add them back in when reading
+  // from the Segment Cache.
+  //
+  // For consistency, we'll do this for dynamic responses, too.
+  //
+  // TODO: We should move search params out of FlightRouterState and handle them
+  // entirely on the client, similar to our plan for dynamic params.
+  const originalSegment = flightRouterState[0]
+  const segmentWithoutSearchParams =
+    typeof originalSegment === 'string' &&
+    originalSegment.startsWith(PAGE_SEGMENT_KEY)
+      ? PAGE_SEGMENT_KEY
+      : originalSegment
+
   return {
     key,
-    segment: flightRouterState[0],
+    segment: segmentWithoutSearchParams,
     slots,
     isRootLayout: flightRouterState[4] === true,
   }
@@ -1093,6 +1167,7 @@ export async function fetchSegmentPrefetchesUsingDynamicRequest(
     // in the LRU as more data comes in.
     fulfilledEntries = writeDynamicRenderResponseIntoCache(
       Date.now(),
+      task,
       response,
       serverData,
       route,
@@ -1181,6 +1256,7 @@ function rejectSegmentEntriesIfStillPending(
 
 function writeDynamicRenderResponseIntoCache(
   now: number,
+  task: PrefetchTask,
   response: Response,
   serverData: NavigationFlightResponse,
   route: FulfilledRouteCacheEntry,
@@ -1231,6 +1307,7 @@ function writeDynamicRenderResponseIntoCache(
           : STATIC_STALETIME_MS
       writeSeedDataIntoCache(
         now,
+        task,
         route,
         now + staleTimeMs,
         seedData,
@@ -1256,6 +1333,7 @@ function writeDynamicRenderResponseIntoCache(
 
 function writeSeedDataIntoCache(
   now: number,
+  task: PrefetchTask,
   route: FulfilledRouteCacheEntry,
   staleAt: number,
   seedData: CacheNodeSeedData,
@@ -1279,7 +1357,12 @@ function writeSeedDataIntoCache(
     fulfillSegmentCacheEntry(ownedEntry, rsc, loading, staleAt, isPartial)
   } else {
     // There's no matching entry. Attempt to create a new one.
-    const possiblyNewEntry = readOrCreateSegmentCacheEntry(now, route, key)
+    const possiblyNewEntry = readOrCreateSegmentCacheEntry(
+      now,
+      task,
+      route,
+      key
+    )
     if (possiblyNewEntry.status === EntryStatus.Empty) {
       // Confirmed this is a new entry. We can fulfill it.
       const newEntry = possiblyNewEntry
@@ -1294,7 +1377,7 @@ function writeSeedDataIntoCache(
         staleAt,
         isPartial
       )
-      upsertSegmentEntry(now, key, newEntry)
+      upsertSegmentEntry(now, getSegmentKeypathForTask(task, key), newEntry)
     }
   }
   // Recursively write the child data into the cache.
@@ -1306,6 +1389,7 @@ function writeSeedDataIntoCache(
         const childSegment = childSeedData[0]
         writeSeedDataIntoCache(
           now,
+          task,
           route,
           staleAt,
           childSeedData,