[Segment Cache] Fix tests related to optimistic loading state reuse (#84498)

When navigating to a route that reads from search params, and there's no
matching route prefetch in the cache, we use a trick where we
optimistically assume that a route will not be rewritten or redirected
based on the search string. So, if we have a matching entry with the
same pathname but different search params, we can construct a route tree
based on that before making a new network request. Then, the router can
show the loading state.

We should be able to reuse any route entry that shares the same
pathname, but for now, we special case the empty search string. The plan
is to refactor the data structure that stores route entries to support
efficient lookups by pathname.

But for now, the workaround is: whenever prefetching a page with a
non-empty search string, also prefetch the route tree for the same
pathname and an empty search string. (Just the route tree, no data.)

Author: acdlite

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

@@ -21,6 +21,8 @@ export function createCacheKey(
   originalHref: string,
   nextUrl: string | null
 ): RouteCacheKey {
+  // TODO: We should remove the hash from the href and track that separately.
+  // There's no reason to vary route entries by hash.
   const originalUrl = new URL(originalHref)
   const cacheKey = {
     href: originalHref as NormalizedHref,

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

@@ -568,11 +568,11 @@ export function waitForSegmentCacheEntry(
  */
 export function readOrCreateRouteCacheEntry(
   now: number,
-  task: PrefetchTask
+  task: PrefetchTask,
+  key: RouteCacheKey
 ): RouteCacheEntry {
   attachInvalidationListener(task)
 
-  const key = task.key
   const existingEntry = readRouteCacheEntry(now, key)
   if (existingEntry !== null) {
     return existingEntry
@@ -649,27 +649,16 @@ export function requestOptimisticRouteCacheEntry(
     // string is in the cache. So we can bail out here.
     return null
   }
+  const urlWithoutSearchParams = new URL(requestedUrl)
+  urlWithoutSearchParams.search = ''
   const routeWithNoSearchParams = readRouteCacheEntry(
     now,
-    createPrefetchRequestKey(
-      requestedUrl.origin + requestedUrl.pathname,
-      nextUrl
-    )
+    createPrefetchRequestKey(urlWithoutSearchParams.href, nextUrl)
   )
 
   if (
     routeWithNoSearchParams === null ||
-    routeWithNoSearchParams.status !== EntryStatus.Fulfilled ||
-    // There's no point constructing an optimistic route tree if the metadata
-    // isn't fully available, because we'll have to do a blocking
-    // navigation anyway.
-    routeWithNoSearchParams.isHeadPartial ||
-    // We cannot reuse this route if it has dynamic metadata.
-    // TODO: Move the metadata out of the route cache entry so the route
-    // tree is reusable separately from the metadata. Then we can remove
-    // these checks.
-    routeWithNoSearchParams.TODO_metadataStatus !== EntryStatus.Empty ||
-    routeWithNoSearchParams.TODO_isHeadDynamic
+    routeWithNoSearchParams.status !== EntryStatus.Fulfilled
   ) {
     // Bail out of constructing an optimistic route tree. This will result in
     // a blocking, unprefetched navigation.
@@ -678,6 +667,36 @@ export function requestOptimisticRouteCacheEntry(
 
   // Now we have a base route tree we can "patch" with our optimistic values.
 
+  const TODO_isHeadDynamic = routeWithNoSearchParams.TODO_isHeadDynamic
+  let head
+  let isHeadPartial
+  let TODO_metadataStatus: EntryStatus.Empty | EntryStatus.Fulfilled
+  if (TODO_isHeadDynamic) {
+    // If the head was fetched via dynamic request, then we don't know
+    // whether it accessed search params. So we must be conservative — we
+    // cannot reuse it. The head will stream in during the dynamic navigation.
+    // TODO: When Cache Components is enabled, we should track whether the
+    // head varied on search params.
+    // TODO: Because we're rendering a `null` viewport as the partial state, the
+    // viewport will not block the navigation; it will stream in later,
+    // alongside the metadata. Viewport is supposed to be blocking. This is
+    // a subtle bug in the old implementation that we've preserved here. It's
+    // rare enough that we're not going to fix it for apps that don't enable
+    // Cache Components; when Cache Components is enabled, though, we should
+    // use an infinite promise here to block the navigation. But only if the
+    // entry actually varies on search params.
+    head = [null, null]
+    // Setting this to `true` ensures that on navigation, the head is requested.
+    isHeadPartial = true
+    TODO_metadataStatus = EntryStatus.Empty
+  } else {
+    // The head was fetched via a static/PPR request. So it's guaranteed to
+    // not contain search params. We can reuse it.
+    head = routeWithNoSearchParams.head
+    isHeadPartial = routeWithNoSearchParams.isHeadPartial
+    TODO_metadataStatus = EntryStatus.Empty
+  }
+
   // Optimistically assume that redirects for the requested pathname do
   // not vary on the search string. Therefore, if the base route was
   // redirected to a different search string, then the optimistic route
@@ -720,17 +739,17 @@ export function requestOptimisticRouteCacheEntry(
     // This isn't cloned because it's instance-specific
     blockedTasks: null,
     tree: routeWithNoSearchParams.tree,
-    head: routeWithNoSearchParams.head,
-    isHeadPartial: routeWithNoSearchParams.isHeadPartial,
+    head,
+    isHeadPartial,
     staleAt: routeWithNoSearchParams.staleAt,
     couldBeIntercepted: routeWithNoSearchParams.couldBeIntercepted,
     isPPREnabled: routeWithNoSearchParams.isPPREnabled,
 
     // Override the rendered search with the optimistic value.
     renderedSearch: optimisticRenderedSearch,
 
-    TODO_metadataStatus: routeWithNoSearchParams.TODO_metadataStatus,
-    TODO_isHeadDynamic: routeWithNoSearchParams.TODO_isHeadDynamic,
+    TODO_metadataStatus,
+    TODO_isHeadDynamic,
 
     // LRU-related fields
     keypath: null,
@@ -1288,13 +1307,13 @@ export function convertRouteTreeToFlightRouterState(
 
 export async function fetchRouteOnCacheMiss(
   entry: PendingRouteCacheEntry,
-  task: PrefetchTask
+  task: PrefetchTask,
+  key: RouteCacheKey
 ): Promise<PrefetchSubtaskResult<null> | null> {
   // This function is allowed to use async/await because it contains the actual
   // fetch that gets issued on a cache miss. Notice it writes the result to the
   // cache entry directly, rather than return data that is then written by
   // the caller.
-  const key = task.key
   const href = key.href
   const nextUrl = key.nextUrl
   const segmentPath = '/_tree' as SegmentRequestKey

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

@@ -142,35 +142,44 @@ export function navigate(
 
   // There was no matching route tree in the cache. Let's see if we can
   // construct an "optimistic" route tree.
-  const optimisticRoute = requestOptimisticRouteCacheEntry(now, url, nextUrl)
-  if (optimisticRoute !== null) {
-    // We have an optimistic route tree. Proceed with the normal flow.
-    const snapshot = readRenderSnapshotFromCache(
-      now,
-      optimisticRoute,
-      optimisticRoute.tree
-    )
-    const prefetchFlightRouterState = snapshot.flightRouterState
-    const prefetchSeedData = snapshot.seedData
-    const prefetchHead = optimisticRoute.head
-    const isPrefetchHeadPartial = optimisticRoute.isHeadPartial
-    const newCanonicalUrl = optimisticRoute.canonicalUrl
-    return navigateUsingPrefetchedRouteTree(
-      now,
-      url,
-      currentUrl,
-      nextUrl,
-      isSamePageNavigation,
-      currentCacheNode,
-      currentFlightRouterState,
-      prefetchFlightRouterState,
-      prefetchSeedData,
-      prefetchHead,
-      isPrefetchHeadPartial,
-      newCanonicalUrl,
-      shouldScroll,
-      url.hash
-    )
+  //
+  // Do not construct an optimistic route tree if there was a cache hit, but
+  // the entry has a rejected status, since it may have been rejected due to a
+  // rewrite or redirect based on the search params.
+  //
+  // TODO: There are multiple reasons a prefetch might be rejected; we should
+  // track them explicitly and choose what to do here based on that.
+  if (route === null || route.status !== EntryStatus.Rejected) {
+    const optimisticRoute = requestOptimisticRouteCacheEntry(now, url, nextUrl)
+    if (optimisticRoute !== null) {
+      // We have an optimistic route tree. Proceed with the normal flow.
+      const snapshot = readRenderSnapshotFromCache(
+        now,
+        optimisticRoute,
+        optimisticRoute.tree
+      )
+      const prefetchFlightRouterState = snapshot.flightRouterState
+      const prefetchSeedData = snapshot.seedData
+      const prefetchHead = optimisticRoute.head
+      const isPrefetchHeadPartial = optimisticRoute.isHeadPartial
+      const newCanonicalUrl = optimisticRoute.canonicalUrl
+      return navigateUsingPrefetchedRouteTree(
+        now,
+        url,
+        currentUrl,
+        nextUrl,
+        isSamePageNavigation,
+        currentCacheNode,
+        currentFlightRouterState,
+        prefetchFlightRouterState,
+        prefetchSeedData,
+        prefetchHead,
+        isPrefetchHeadPartial,
+        newCanonicalUrl,
+        shouldScroll,
+        url.hash
+      )
+    }
   }
 
   // There's no matching prefetch for this route in the cache.

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

@@ -28,6 +28,7 @@ import {
   canNewFetchStrategyProvideMoreContent,
 } from './cache'
 import type { RouteCacheKey } from './cache-key'
+import { createCacheKey } from './cache-key'
 import {
   FetchStrategy,
   type PrefetchTaskFetchStrategy,
@@ -439,8 +440,7 @@ function processQueueInMicrotask() {
   while (task !== null && hasNetworkBandwidth(task)) {
     task.cacheVersion = getCurrentCacheVersion()
 
-    const route = readOrCreateRouteCacheEntry(now, task)
-    const exitStatus = pingRootRouteTree(now, task, route)
+    const exitStatus = pingRoute(now, task)
 
     // These fields are only valid for a single attempt. Reset them after each
     // iteration of the task queue.
@@ -501,6 +501,57 @@ function background(task: PrefetchTask): boolean {
   return false
 }
 
+function pingRoute(now: number, task: PrefetchTask): PrefetchTaskExitStatus {
+  const key = task.key
+  const route = readOrCreateRouteCacheEntry(now, task, key)
+  const exitStatus = pingRootRouteTree(now, task, route)
+
+  if (exitStatus !== PrefetchTaskExitStatus.InProgress && key.search !== '') {
+    // If the URL has a non-empty search string, also prefetch the pathname
+    // without the search string. We use the searchless route tree as a base for
+    // optimistic routing; see requestOptimisticRouteCacheEntry for details.
+    //
+    // Note that we don't need to prefetch any of the segment data. Just the
+    // route tree.
+    //
+    // TODO: This is a temporary solution; the plan is to replace this by adding
+    // a wildcard lookup method to the TupleMap implementation. This is
+    // non-trivial to implement because it needs to account for things like
+    // fallback route entries, hence this temporary workaround.
+    const url = new URL(key.href)
+    url.search = ''
+    const keyWithoutSearch = createCacheKey(url.href, key.nextUrl)
+    const routeWithoutSearch = readOrCreateRouteCacheEntry(
+      now,
+      task,
+      keyWithoutSearch
+    )
+    switch (routeWithoutSearch.status) {
+      case EntryStatus.Empty: {
+        if (background(task)) {
+          routeWithoutSearch.status = EntryStatus.Pending
+          spawnPrefetchSubtask(
+            fetchRouteOnCacheMiss(routeWithoutSearch, task, keyWithoutSearch)
+          )
+        }
+        break
+      }
+      case EntryStatus.Pending:
+      case EntryStatus.Fulfilled:
+      case EntryStatus.Rejected: {
+        // Either the route tree is already cached, or there's already a
+        // request in progress. Since we don't need to fetch any segment data
+        // for this route, there's nothing left to do.
+        break
+      }
+      default:
+        routeWithoutSearch satisfies never
+    }
+  }
+
+  return exitStatus
+}
+
 function pingRootRouteTree(
   now: number,
   task: PrefetchTask,
@@ -522,7 +573,7 @@ function pingRootRouteTree(
       // behavior if PPR is disabled for a route (via the incremental opt-in).
       //
       // Those cases will be handled here.
-      spawnPrefetchSubtask(fetchRouteOnCacheMiss(route, task))
+      spawnPrefetchSubtask(fetchRouteOnCacheMiss(route, task, task.key))
 
       // If the request takes longer than a minute, a subsequent request should
       // retry instead of waiting for this one. When the response is received,

test/client-segment-cache-tests-manifest.json

@@ -25,18 +25,6 @@
       "failed": [
         "loading-tsx-no-partial-rendering when PPR is enabled, loading.tsx boundaries do not cause a partial prefetch"
       ]
-    },
-    "test/e2e/app-dir/searchparams-reuse-loading/searchparams-reuse-loading.test.ts": {
-      "failed": [
-        "searchparams-reuse-loading should re-use loading from \"full\" prefetch for param-full URL when navigating to param-full route",
-        "searchparams-reuse-loading should re-use loading from \"full\" prefetch for param-full URL when navigating to param-less route",
-        "searchparams-reuse-loading should re-use loading from \"full\" prefetch for param-less URL when navigating to param-full route"
-      ]
-    },
-    "test/e2e/next-form/default/next-form-prefetch.test.ts": {
-      "failed": [
-        "app dir - form prefetching should prefetch a loading state for the form's target"
-      ]
     }
   },
   "rules": {

test/e2e/app-dir/segment-cache/incremental-opt-in/segment-cache-incremental-opt-in.test.ts

@@ -485,6 +485,11 @@ describe('segment cache (incremental opt in)', () => {
       // prefetch the title, even though it's dynamic.
       {
         includes: 'Dynamic Title: yay',
+        // TODO: Due to a race condition, the metadata is sometimes fetched twice
+        // instead of once (though not more than that). The planned fix is to
+        // cache metadata separate from the route tree. Then we can remove
+        // this option.
+        allowMultipleResponses: true,
       }
     )
 

test/e2e/app-dir/segment-cache/router-act.ts

@@ -21,7 +21,11 @@ type PendingRSCRequest = {
 
 let currentBatch: Batch | null = null
 
-type ExpectedResponseConfig = { includes: string; block?: boolean | 'reject' }
+type ExpectedResponseConfig = {
+  includes: string
+  block?: boolean | 'reject'
+  allowMultipleResponses?: boolean
+}
 
 /**
  * Represents the expected responses sent by the server to fulfill requests
@@ -334,7 +338,8 @@ ${fulfilled.body}
                   // error message.
                   const otherResponse = alreadyMatched.get(includes)
                   if (otherResponse !== undefined) {
-                    error.message = `
+                    if (!expectedResponse.allowMultipleResponses) {
+                      error.message = `
 Received multiple responses containing the same expected substring.
 
 Expected substring:
@@ -348,13 +353,15 @@ ${fulfilled.body}
 
 Choose a more specific substring to assert on.
 `
-                    throw error
-                  }
-                  alreadyMatched.set(includes, fulfilled.body)
-                  if (actualResponses === null) {
-                    actualResponses = [expectedResponse]
+                      throw error
+                    }
                   } else {
-                    actualResponses.push(expectedResponse)
+                    alreadyMatched.set(includes, fulfilled.body)
+                    if (actualResponses === null) {
+                      actualResponses = [expectedResponse]
+                    } else {
+                      actualResponses.push(expectedResponse)
+                    }
                   }
                   if (block) {
                     shouldBlock = true