[PPR Nav] Fix flash of loading state during back/forward (vercel#60578)

### Depends on

- vercel#60577 

---

A popstate navigation reads data from the local cache. It does not issue
new network requests (unless the cache entries have been evicted). So,
when navigating with back/forward, we should not switch back to the PPR
loading state. We should render the full, cached dynamic data
immediately.

To implement this, on a popstate navigation, we update the cache to drop
the prefetch data for any segment whose dynamic data was already
received. We clone the entire cache node tree and set the `prefetchRsc`
field to `null` to prevent it from being rendered. (We can't mutate the
node in place because Cache Node is a concurrent data structure.)

Technically, what we're actually checking is whether the dynamic network
response was received. But since it's a streaming response, this does
not mean that all the dynamic data has fully streamed in. It just means
that _some_ of the dynamic data was received. But as a heuristic, we
assume that the rest dynamic data will stream in quickly, so it's still
better to skip the prefetch state.

Closes NEXT-2084

packages/next/src/client/components/router-reducer/ppr-navigations.ts

@@ -727,6 +727,67 @@ function abortPendingCacheNode(
   }
 }
 
+export function updateCacheNodeOnPopstateRestoration(
+  oldCacheNode: CacheNode,
+  routerState: FlightRouterState
+) {
+  // A popstate navigation reads data from the local cache. It does not issue
+  // new network requests (unless the cache entries have been evicted). So, we
+  // update the cache to drop the prefetch  data for any segment whose dynamic
+  // data was already received. This prevents an unnecessary flash back to PPR
+  // state during a back/forward navigation.
+  //
+  // This function clones the entire cache node tree and sets the `prefetchRsc`
+  // field to `null` to prevent it from being rendered. We can't mutate the node
+  // in place because this is a concurrent data structure.
+
+  const routerStateChildren = routerState[1]
+  const oldParallelRoutes = oldCacheNode.parallelRoutes
+  const newParallelRoutes = new Map(oldParallelRoutes)
+  for (let parallelRouteKey in routerStateChildren) {
+    const routerStateChild: FlightRouterState =
+      routerStateChildren[parallelRouteKey]
+    const segmentChild = routerStateChild[0]
+    const segmentKeyChild = createRouterCacheKey(segmentChild)
+    const oldSegmentMapChild = oldParallelRoutes.get(parallelRouteKey)
+    if (oldSegmentMapChild !== undefined) {
+      const oldCacheNodeChild = oldSegmentMapChild.get(segmentKeyChild)
+      if (oldCacheNodeChild !== undefined) {
+        const newCacheNodeChild = updateCacheNodeOnPopstateRestoration(
+          oldCacheNodeChild,
+          routerStateChild
+        )
+        const newSegmentMapChild = new Map(oldSegmentMapChild)
+        newSegmentMapChild.set(segmentKeyChild, newCacheNodeChild)
+        newParallelRoutes.set(parallelRouteKey, newSegmentMapChild)
+      }
+    }
+  }
+
+  // Only show prefetched data if the dynamic data is still pending.
+  //
+  // Tehnically, what we're actually checking is whether the dynamic network
+  // response was received. But since it's a streaming response, this does not
+  // mean that all the dynamic data has fully streamed in. It just means that
+  // _some_ of the dynamic data was received. But as a heuristic, we assume that
+  // the rest dynamic data will stream in quickly, so it's still better to skip
+  // the prefetch state.
+  const rsc = oldCacheNode.rsc
+  const shouldUsePrefetch = isDeferredRsc(rsc) && rsc.status === 'pending'
+
+  return {
+    lazyData: null,
+    rsc,
+    head: oldCacheNode.head,
+
+    prefetchHead: shouldUsePrefetch ? oldCacheNode.prefetchHead : null,
+    prefetchRsc: shouldUsePrefetch ? oldCacheNode.prefetchRsc : null,
+
+    // These are the cloned children we computed above
+    parallelRoutes: newParallelRoutes,
+  }
+}
+
 const DEFERRED = Symbol()
 
 type PendingDeferredRsc = Promise<React.ReactNode> & {

packages/next/src/client/components/router-reducer/reducers/restore-reducer.ts

@@ -5,6 +5,7 @@ import type {
   RestoreAction,
 } from '../router-reducer-types'
 import { extractPathFromFlightRouterState } from '../compute-changed-path'
+import { updateCacheNodeOnPopstateRestoration } from '../ppr-navigations'
 
 export function restoreReducer(
   state: ReadonlyReducerState,
@@ -13,6 +14,15 @@ export function restoreReducer(
   const { url, tree } = action
   const href = createHrefFromUrl(url)
 
+  const oldCache = state.cache
+  const newCache = process.env.__NEXT_PPR
+    ? // When PPR is enabled, we update the cache to drop the prefetch
+      // data for any segment whose dynamic data was already received. This
+      // prevents an unnecessary flash back to PPR state during a
+      // back/forward navigation.
+      updateCacheNodeOnPopstateRestoration(oldCache, tree)
+    : oldCache
+
   return {
     buildId: state.buildId,
     // Set canonical url
@@ -24,7 +34,7 @@ export function restoreReducer(
       preserveCustomHistoryState: true,
     },
     focusAndScrollRef: state.focusAndScrollRef,
-    cache: state.cache,
+    cache: newCache,
     prefetchCache: state.prefetchCache,
     // Restore provided tree
     tree: tree,

test/e2e/app-dir/ppr-navigations/avoid-popstate-flash/app/layout.tsx

@@ -0,0 +1,11 @@
+import { ShouldFallbackThrowContainer } from './some-page/client'
+
+export default function Root({ children }) {
+  return (
+    <html>
+      <body>
+        <ShouldFallbackThrowContainer>{children}</ShouldFallbackThrowContainer>
+      </body>
+    </html>
+  )
+}

test/e2e/app-dir/ppr-navigations/avoid-popstate-flash/app/page.tsx

@@ -0,0 +1,9 @@
+import Link from 'next/link'
+
+export default function Page() {
+  return (
+    <div>
+      <Link href="/some-page">Some page</Link>
+    </div>
+  )
+}

test/e2e/app-dir/ppr-navigations/avoid-popstate-flash/app/some-page/client.tsx

@@ -0,0 +1,45 @@
+'use client'
+
+import { Component, createContext, use, useState } from 'react'
+
+const ShouldFallbackThrowContext = createContext(false)
+
+export function ShouldFallbackThrowContainer({ children }) {
+  const [shouldFallbackThrow, setShouldFallbackThrow] = useState(false)
+  return (
+    <>
+      <label>
+        Throw if fallback appears
+        <input
+          id="should-fallback-throw"
+          type="checkbox"
+          checked={shouldFallbackThrow}
+          onChange={(e) => setShouldFallbackThrow(e.target.checked)}
+        />
+      </label>
+      <ShouldFallbackThrowContext.Provider value={shouldFallbackThrow}>
+        <ErrorBoundary>{children}</ErrorBoundary>
+      </ShouldFallbackThrowContext.Provider>
+    </>
+  )
+}
+
+export function Fallback({ children }) {
+  if (use(ShouldFallbackThrowContext)) {
+    throw new Error('Unexpected fallback')
+  }
+  return children
+}
+
+class ErrorBoundary extends Component<{ children: React.ReactNode }> {
+  state = { error: null }
+  static getDerivedStateFromError(error) {
+    return { error }
+  }
+  render() {
+    if (this.state.error) {
+      return <div id="error">{this.state.error.message}</div>
+    }
+    return this.props.children
+  }
+}

test/e2e/app-dir/ppr-navigations/avoid-popstate-flash/app/some-page/page.tsx

@@ -0,0 +1,24 @@
+import React, { Suspense } from 'react'
+import { getDynamicTestData, getStaticTestData } from '../test-data-service'
+import { Fallback } from './client'
+
+async function Dynamic() {
+  return <div id="dynamic">{await getDynamicTestData('Dynamic')}</div>
+}
+
+async function Static() {
+  return <div id="static">{await getStaticTestData('Static')}</div>
+}
+
+export default async function Page() {
+  return (
+    <div id="container">
+      <Suspense fallback={<Fallback>Loading dynamic...</Fallback>}>
+        <Dynamic />
+      </Suspense>
+      <Suspense fallback={<Fallback>Loading static...</Fallback>}>
+        <Static />
+      </Suspense>
+    </div>
+  )
+}

test/e2e/app-dir/ppr-navigations/avoid-popstate-flash/app/test-data-service.ts

@@ -0,0 +1,47 @@
+import 'server-only'
+
+import { unstable_noStore } from 'next/cache'
+
+// NOTE: I've intentionally not yet moved these helpers into a shared module, to
+// avoid early abstraction. I will if/when we start using them for other tests.
+// They are based on the testing patterns we use all over the React codebase, so
+// I'm reasonably confident in them.
+const TEST_DATA_SERVICE_URL = process.env.TEST_DATA_SERVICE_URL
+const ARTIFICIAL_DELAY = 3000
+
+async function getTestData(key: string, isStatic: boolean): Promise<string> {
+  const searchParams = new URLSearchParams({
+    key,
+  })
+  if (!TEST_DATA_SERVICE_URL) {
+    // If environment variable is not set, resolve automatically after a delay.
+    // This is so you can run the test app locally without spinning up a
+    // data server.
+    await new Promise<void>((resolve) =>
+      setTimeout(() => resolve(), ARTIFICIAL_DELAY)
+    )
+    if (!isStatic) {
+      unstable_noStore()
+    }
+    return key
+  }
+  const response = await fetch(
+    TEST_DATA_SERVICE_URL + '?' + searchParams.toString(),
+    {
+      cache: isStatic ? 'force-cache' : 'no-store',
+    }
+  )
+  const text = await response.text()
+  if (response.status !== 200) {
+    throw new Error(text)
+  }
+  return text
+}
+
+export async function getStaticTestData(key: string): Promise<string> {
+  return getTestData(key, true)
+}
+
+export async function getDynamicTestData(key: string): Promise<string> {
+  return getTestData(key, false)
+}

test/e2e/app-dir/ppr-navigations/avoid-popstate-flash/avoid-popstate-flash.test.ts

@@ -0,0 +1,89 @@
+import { createNext } from 'e2e-utils'
+import { findPort } from 'next-test-utils'
+import { createTestDataServer } from 'test-data-service/writer'
+import { createTestLog } from 'test-log'
+
+describe('avoid-popstate-flash', () => {
+  if ((global as any).isNextDev) {
+    test('ppr is disabled in dev', () => {})
+    return
+  }
+
+  let server
+  let next
+  afterEach(async () => {
+    await next?.destroy()
+    server?.close()
+  })
+
+  test('does not flash back to partial PPR data during back/forward navigation', async () => {
+    const TestLog = createTestLog()
+    let autoresolveRequests = true
+    let pendingRequests = new Map()
+    server = createTestDataServer(async (key, res) => {
+      TestLog.log('REQUEST: ' + key)
+      if (autoresolveRequests) {
+        res.resolve()
+        return
+      }
+      if (pendingRequests.has(key)) {
+        throw new Error('Request already pending for ' + key)
+      }
+      pendingRequests.set(key, res)
+    })
+    const port = await findPort()
+    server.listen(port)
+    next = await createNext({
+      files: __dirname,
+      env: { TEST_DATA_SERVICE_URL: `http://localhost:${port}` },
+    })
+    TestLog.assert(['REQUEST: Static'])
+    autoresolveRequests = false
+
+    const browser = await next.browser('/')
+
+    // Navigate to the target page.
+    const link = await browser.elementByCss('a[href="/some-page"]')
+    await link.click()
+
+    // The static UI appears immediately because it was prerendered at
+    // build time.
+    const staticContainer = await browser.elementById('static')
+    expect(await staticContainer.innerText()).toBe('Static')
+
+    await TestLog.waitFor(['REQUEST: Dynamic'])
+    pendingRequests.get('Dynamic').resolve()
+
+    // Now the dynamic data appears.
+    const dynamic = await browser.elementById('dynamic')
+    expect(await dynamic.innerText()).toBe('Dynamic')
+
+    // At this point all the data has been loaded into the cache. We're going
+    // to test what happens during a back/forward navigation.
+
+    // Set a global state that causes Suspense fallbacks to throw.
+    const checkbox = await browser.elementById('should-fallback-throw')
+    await checkbox.click()
+    const checked = await checkbox.getProperty('checked')
+    expect(await checked.jsonValue()).toBe(true)
+
+    // Navigate using back/forward using the browser's history stack. This
+    // should not trigger a fresh navigation, nor any network requests. We
+    // should read the data from the cache. And we should not render the
+    // partially complete PPR data that was used during the initial navigation.
+    //
+    // If the data is not read from cache, or if partial data is shown, it will
+    // trigger a fallback, which will throw an error because of the state we
+    // set above.
+    await browser.back()
+    await browser.forward()
+
+    // Confirm that the dynamic data is visible. This implies that the fallback
+    // did not throw.
+    const dynamic2 = await browser.elementById('dynamic')
+    expect(await dynamic2.innerText()).toBe('Dynamic')
+
+    // There should have been no additional requests.
+    TestLog.assert([])
+  })
+})

test/e2e/app-dir/ppr-navigations/avoid-popstate-flash/next.config.js

@@ -0,0 +1,10 @@
+/**
+ * @type {import('next').NextConfig}
+ */
+const nextConfig = {
+  experimental: {
+    ppr: true,
+  },
+}
+
+module.exports = nextConfig