Lazily call refreshTags and getExpiration (#77779)

When `"use cache"` is not used on the current route, we don't need to call `refreshTags` for the configured cache handlers. So instead of calling it at the beginning of the request for every cache handler, we now call it lazily right before the first cache entry is retrieved for the respective cache handler (once per request).

Similarly, we now call `getExpiration` for the implicit tags of the current route lazily (also once per request) after an existing cache entry has been retrieved, and its timestamp needs to be compared with the expiration of the implicit tags.

Author: unstubbable

Diff for: packages/next/src/server/app-render/work-async-storage.external.ts

@@ -54,15 +54,26 @@ export interface WorkStore {
   nextFetchId?: number
   pathWasRevalidated?: boolean
 
-  // Tags that were revalidated during the current request. They need to be sent
-  // to cache handlers to propagate their revalidation.
+  /**
+   * Tags that were revalidated during the current request. They need to be sent
+   * to cache handlers to propagate their revalidation.
+   */
   pendingRevalidatedTags?: string[]
 
-  // Tags that were previously revalidated (e.g. by a redirecting server action)
-  // and have already been sent to cache handlers. Retrieved cache entries that
-  // include any of these tags must be discarded.
+  /**
+   * Tags that were previously revalidated (e.g. by a redirecting server action)
+   * and have already been sent to cache handlers. Retrieved cache entries that
+   * include any of these tags must be discarded.
+   */
   readonly previouslyRevalidatedTags: readonly string[]
 
+  /**
+   * This map contains promise-like values so that we can evaluate them lazily
+   * when a cache entry is read. It allows us to skip refreshing tags if no
+   * caches are read at all.
+   */
+  readonly refreshTagsByCacheKind: Map<string, PromiseLike<void>>
+
   fetchMetrics?: FetchMetrics
 
   isDraftMode?: boolean

Diff for: packages/next/src/server/async-storage/work-store.ts

@@ -10,6 +10,8 @@ import type { CacheLife } from '../use-cache/cache-life'
 import { AfterContext } from '../after/after-context'
 
 import { normalizeAppPath } from '../../shared/lib/router/utils/app-paths'
+import { createLazyResult } from '../lib/lazy-result'
+import { getCacheHandlerEntries } from '../use-cache/handlers'
 
 export type WorkStoreContext = {
   /**
@@ -135,6 +137,7 @@ export function createWorkStore({
     dynamicIOEnabled: renderOpts.experimental.dynamicIO,
     dev: renderOpts.dev ?? false,
     previouslyRevalidatedTags,
+    refreshTagsByCacheKind: createRefreshTagsByCacheKind(),
   }
 
   // TODO: remove this when we resolve accessing the store outside the execution context
@@ -151,3 +154,25 @@ function createAfterContext(renderOpts: RequestLifecycleOpts): AfterContext {
     onTaskError: onAfterTaskError,
   })
 }
+
+/**
+ * Creates a map with promise-like objects, that refresh tags for the given
+ * cache kind when they're awaited for the first time.
+ */
+function createRefreshTagsByCacheKind(): Map<string, PromiseLike<void>> {
+  const refreshTagsByCacheKind = new Map<string, PromiseLike<void>>()
+  const cacheHandlers = getCacheHandlerEntries()
+
+  if (cacheHandlers) {
+    for (const [kind, cacheHandler] of cacheHandlers) {
+      if ('refreshTags' in cacheHandler) {
+        refreshTagsByCacheKind.set(
+          kind,
+          createLazyResult(async () => cacheHandler.refreshTags())
+        )
+      }
+    }
+  }
+
+  return refreshTagsByCacheKind
+}

Diff for: packages/next/src/server/base-server.ts

@@ -1440,7 +1440,9 @@ export default abstract class Server<
         await Promise.all(
           [...cacheHandlers].map(async (cacheHandler) => {
             if ('refreshTags' in cacheHandler) {
-              await cacheHandler.refreshTags()
+              // Note: cacheHandler.refreshTags() is called lazily before the
+              // first cache entry is retrieved. It allows us to skip the
+              // refresh request if no caches are read at all.
             } else {
               const previouslyRevalidatedTags = getPreviouslyRevalidatedTags(
                 req.headers,

Diff for: packages/next/src/server/lib/implicit-tags.ts

@@ -1,6 +1,7 @@
 import { NEXT_CACHE_IMPLICIT_TAG_ID } from '../../lib/constants'
 import type { FallbackRouteParams } from '../request/fallback-params'
 import { getCacheHandlers } from '../use-cache/handlers'
+import { createLazyResult } from './lazy-result'
 
 export interface ImplicitTags {
   /**
@@ -9,11 +10,13 @@ export interface ImplicitTags {
    */
   readonly tags: string[]
   /**
-   * Modern cache handlers don't receive implicit tags. Instead, the
-   * implicit tags' expiration is stored in the work unit store, and used to
-   * compare with a cache entry's timestamp.
+   * Modern cache handlers don't receive implicit tags. Instead, the implicit
+   * tags' expiration is stored in the work unit store, and used to compare with
+   * a cache entry's timestamp. Note: This is a promise-like value so that we
+   * can evaluate it lazily when a cache entry is read. It allows us to skip
+   * fetching the expiration value if no caches are read at all.
    */
-  readonly expiration: number
+  readonly expiration: PromiseLike<number>
 }
 
 const getDerivedTags = (pathname: string): string[] => {
@@ -98,7 +101,9 @@ export async function getImplicitTags(
     tags.push(tag)
   }
 
-  const expiration = await getImplicitTagsExpiration(tags)
+  const expiration = createLazyResult(async () =>
+    getImplicitTagsExpiration(tags)
+  )
 
   return { tags, expiration }
 }

Diff for: packages/next/src/server/lib/lazy-result.ts

@@ -0,0 +1,19 @@
+/**
+ * Calls the given async function only when the returned promise-like object is
+ * awaited.
+ */
+export function createLazyResult<TResult>(
+  fn: () => Promise<TResult>
+): PromiseLike<TResult> {
+  let pendingResult: Promise<TResult> | undefined
+
+  return {
+    then(onfulfilled, onrejected) {
+      if (!pendingResult) {
+        pendingResult = fn()
+      }
+
+      return pendingResult.then(onfulfilled, onrejected)
+    },
+  }
+}

Diff for: packages/next/src/server/use-cache/handlers.ts

@@ -78,7 +78,8 @@ export function initializeCacheHandlers(): boolean {
 /**
  * Get a cache handler by kind.
  * @param kind - The kind of cache handler to get.
- * @returns The cache handler, or `undefined` if it is not initialized or does not exist.
+ * @returns The cache handler, or `undefined` if it does not exist.
+ * @throws If the cache handlers are not initialized.
  */
 export function getCacheHandler(kind: string): CacheHandlerCompat | undefined {
   // This should never be called before initializeCacheHandlers.
@@ -90,8 +91,9 @@ export function getCacheHandler(kind: string): CacheHandlerCompat | undefined {
 }
 
 /**
- * Get an iterator over the cache handlers.
- * @returns An iterator over the cache handlers, or `undefined` if they are not initialized.
+ * Get a set iterator over the cache handlers.
+ * @returns An iterator over the cache handlers, or `undefined` if they are not
+ * initialized.
  */
 export function getCacheHandlers():
   | SetIterator<CacheHandlerCompat>
@@ -103,6 +105,22 @@ export function getCacheHandlers():
   return reference[handlersSetSymbol].values()
 }
 
+/**
+ * Get a map iterator over the cache handlers (keyed by kind).
+ * @returns An iterator over the cache handler entries, or `undefined` if they
+ * are not initialized.
+ * @throws If the cache handlers are not initialized.
+ */
+export function getCacheHandlerEntries():
+  | MapIterator<[string, CacheHandlerCompat]>
+  | undefined {
+  if (!reference[handlersMapSymbol]) {
+    return undefined
+  }
+
+  return reference[handlersMapSymbol].entries()
+}
+
 /**
  * Set a cache handler by kind.
  * @param kind - The kind of cache handler to set.

Diff for: packages/next/src/server/use-cache/use-cache-wrapper.ts

@@ -52,6 +52,7 @@ import {
 import type { Params } from '../request/params'
 import React from 'react'
 import type { ImplicitTags } from '../lib/implicit-tags'
+import { createLazyResult } from '../lib/lazy-result'
 
 type CacheKeyParts =
   | [buildId: string, id: string, args: unknown[]]
@@ -694,6 +695,11 @@ export function cache(
         const implicitTags = workUnitStore?.implicitTags
         const forceRevalidate = shouldForceRevalidate(workStore, workUnitStore)
 
+        // Lazily refresh the tags for the cache handler that's associated with
+        // this cache function. This is only done once per request and cache
+        // handler, when it's awaited for the first time.
+        await workStore.refreshTagsByCacheKind.get(kind)
+
         let entry = forceRevalidate
           ? undefined
           : 'getExpiration' in cacheHandler
@@ -706,7 +712,10 @@ export function cache(
                 implicitTags?.tags ?? []
               )
 
-        if (entry && shouldDiscardCacheEntry(entry, workStore, implicitTags)) {
+        if (
+          entry &&
+          (await shouldDiscardCacheEntry(entry, workStore, implicitTags))
+        ) {
           entry = undefined
         }
 
@@ -880,25 +889,6 @@ export function cache(
   return React.cache(cachedFn)
 }
 
-/**
- * Calls the given function only when the returned promise is awaited.
- */
-function createLazyResult<TResult>(
-  fn: () => Promise<TResult>
-): PromiseLike<TResult> {
-  let pendingResult: Promise<TResult> | undefined
-
-  return {
-    then(onfulfilled, onrejected) {
-      if (!pendingResult) {
-        pendingResult = fn()
-      }
-
-      return pendingResult.then(onfulfilled, onrejected)
-    },
-  }
-}
-
 function isPageComponent(
   args: any[]
 ): args is [UseCachePageComponentProps, undefined] {
@@ -937,11 +927,11 @@ function shouldForceRevalidate(
   return false
 }
 
-function shouldDiscardCacheEntry(
+async function shouldDiscardCacheEntry(
   entry: CacheEntry,
   workStore: WorkStore,
   implicitTags: ImplicitTags | undefined
-): boolean {
+): Promise<boolean> {
   // If the cache entry contains revalidated tags that the cache handler might
   // not know about yet, we need to discard it.
   if (entry.tags.some((tag) => isRecentlyRevalidatedTag(tag, workStore))) {
@@ -951,7 +941,7 @@ function shouldDiscardCacheEntry(
   if (implicitTags) {
     // If the cache entry was created before any of the implicit tags were
     // revalidated last, we also need to discard it.
-    if (entry.timestamp <= implicitTags.expiration) {
+    if (entry.timestamp <= (await implicitTags.expiration)) {
       return true
     }
 

Diff for: test/e2e/app-dir/use-cache-custom-handler/app/no-cache/page.tsx

@@ -0,0 +1,3 @@
+export default function Page() {
+  return <p>This page does not use "use cache".</p>
+}

Diff for: test/e2e/app-dir/use-cache-custom-handler/use-cache-custom-handler.test.ts

@@ -23,13 +23,9 @@ describe('use-cache-custom-handler', () => {
     const initialData = await browser.elementById('data').text()
     expect(initialData).toMatch(isoDateRegExp)
 
-    expect(next.cliOutput.slice(outputIndex)).toContain(
-      'ModernCustomCacheHandler::refreshTags'
-    )
+    const cliOutput = next.cliOutput.slice(outputIndex)
 
-    expect(next.cliOutput.slice(outputIndex)).toContain(
-      `ModernCustomCacheHandler::getExpiration ["_N_T_/layout","_N_T_/page","_N_T_/"]`
-    )
+    expect(cliOutput).toContain('ModernCustomCacheHandler::refreshTags')
 
     expect(next.cliOutput.slice(outputIndex)).toMatch(
       /ModernCustomCacheHandler::get \["(development|[A-Za-z0-9_-]{21})","([0-9a-f]{2})+",\[\]\]/
@@ -39,13 +35,26 @@ describe('use-cache-custom-handler', () => {
       /ModernCustomCacheHandler::set \["(development|[A-Za-z0-9_-]{21})","([0-9a-f]{2})+",\[\]\]/
     )
 
+    // Since no existing cache entry was retrieved, we don't need to call
+    // getExpiration() to compare the cache entries timestamp with the
+    // expiration of the implicit tags.
+    expect(cliOutput).not.toContain(`ModernCustomCacheHandler::getExpiration`)
+
     // The data should be cached initially.
 
+    outputIndex = next.cliOutput.length
     await browser.refresh()
     let data = await browser.elementById('data').text()
     expect(data).toMatch(isoDateRegExp)
     expect(data).toEqual(initialData)
 
+    // Now that a cache entry exists, we expect that getExpiration() is called
+    // to compare the cache entries timestamp with the expiration of the
+    // implicit tags.
+    expect(next.cliOutput.slice(outputIndex)).toContain(
+      `ModernCustomCacheHandler::getExpiration ["_N_T_/layout","_N_T_/page","_N_T_/"]`
+    )
+
     // Because we use a low `revalidate` value for the "use cache" function, new
     // data should be returned eventually.
 
@@ -57,6 +66,19 @@ describe('use-cache-custom-handler', () => {
     }, 5000)
   })
 
+  it('calls neither refreshTags nor getExpiration if "use cache" is not used', async () => {
+    await next.fetch(`/no-cache`)
+    const cliOutput = next.cliOutput.slice(outputIndex)
+
+    expect(cliOutput).not.toContain('ModernCustomCacheHandler::refreshTags')
+    expect(cliOutput).not.toContain(`ModernCustomCacheHandler::getExpiration`)
+
+    // We don't optimize for legacy cache handlers though:
+    expect(cliOutput).toContain(
+      `LegacyCustomCacheHandler::receiveExpiredTags []`
+    )
+  })
+
   it('should use a legacy custom cache handler if provided', async () => {
     const browser = await next.browser(`/legacy`)
     const initialData = await browser.elementById('data').text()
@@ -142,29 +164,20 @@ describe('use-cache-custom-handler', () => {
     await retry(async () => {
       const cliOutput = next.cliOutput.slice(outputIndex)
       expect(cliOutput).toInclude('ModernCustomCacheHandler::refreshTags')
-      expect(cliOutput).toInclude('ModernCustomCacheHandler::getExpiration')
       expect(cliOutput).not.toInclude('ModernCustomCacheHandler::expireTags')
     })
   })
 
-  it('should not call getExpiration again after an action', async () => {
+  it('should not call getExpiration after an action', async () => {
     const browser = await next.browser(`/`)
 
-    await retry(async () => {
-      const cliOutput = next.cliOutput.slice(outputIndex)
-      expect(cliOutput).toInclude('ModernCustomCacheHandler::getExpiration')
-    })
-
     outputIndex = next.cliOutput.length
 
     await browser.elementById('revalidate-tag').click()
 
     await retry(async () => {
       const cliOutput = next.cliOutput.slice(outputIndex)
-      expect(cliOutput).toIncludeRepeated(
-        'ModernCustomCacheHandler::getExpiration',
-        1
-      )
+      expect(cliOutput).not.toInclude('ModernCustomCacheHandler::getExpiration')
       expect(cliOutput).toIncludeRepeated(
         `ModernCustomCacheHandler::expireTags`,
         1