Fix console replaying and React.cache usage in "use cache" functions (#75520)

Console logs in server components are replayed in the browser. For
example, when you run `console.log('foo')`, the log will be replayed in
the browser as `[Server] foo`. When the component is inside of a `"use
cache"` scope, it's replayed as `[Cache] foo`.

However, when logging directly in the function body of a `"use cache"`
function, we are currently not replaying the log in the browser.

The reason for that is that the function is called outside of React's
rendering, before handing the result promise over to React for
serialization. Since the function is called without React's request
storage, no console chunks are emitted.

We can work around this by invoking the function lazily when React calls
`.then()` on the promise. This ensures that the function is run inside
of React's request storage and console chunks can be emitted.

In addition, this also unlocks that `React.cache` can be used in a `"use
cache"` function to dedupe other function calls.

closes NAR-83

Author: unstubbable

packages/next/src/server/use-cache/use-cache-wrapper.ts

@@ -60,7 +60,7 @@ function generateCacheEntry(
   outerWorkUnitStore: WorkUnitStore | undefined,
   clientReferenceManifest: DeepReadonly<ClientReferenceManifestForRsc>,
   encodedArguments: FormData | string,
-  fn: any,
+  fn: (...args: unknown[]) => Promise<unknown>,
   timeoutError: UseCacheTimeoutError
 ): Promise<[ReadableStream, Promise<CacheEntry>]> {
   // We need to run this inside a clean AsyncLocalStorage snapshot so that the cache
@@ -84,7 +84,7 @@ function generateCacheEntryWithRestoredWorkStore(
   outerWorkUnitStore: WorkUnitStore | undefined,
   clientReferenceManifest: DeepReadonly<ClientReferenceManifestForRsc>,
   encodedArguments: FormData | string,
-  fn: any,
+  fn: (...args: unknown[]) => Promise<unknown>,
   timeoutError: UseCacheTimeoutError
 ) {
   // Since we cleared the AsyncLocalStorage we need to restore the workStore.
@@ -111,7 +111,7 @@ function generateCacheEntryWithCacheContext(
   outerWorkUnitStore: WorkUnitStore | undefined,
   clientReferenceManifest: DeepReadonly<ClientReferenceManifestForRsc>,
   encodedArguments: FormData | string,
-  fn: any,
+  fn: (...args: unknown[]) => Promise<unknown>,
   timeoutError: UseCacheTimeoutError
 ) {
   if (!workStore.cacheLifeProfiles) {
@@ -291,7 +291,7 @@ async function generateCacheEntryImpl(
   innerCacheStore: UseCacheStore,
   clientReferenceManifest: DeepReadonly<ClientReferenceManifestForRsc>,
   encodedArguments: FormData | string,
-  fn: any,
+  fn: (...args: unknown[]) => Promise<unknown>,
   timeoutError: UseCacheTimeoutError
 ): Promise<[ReadableStream, Promise<CacheEntry>]> {
   const temporaryReferences = createServerTemporaryReferenceSet()
@@ -335,24 +335,28 @@ async function generateCacheEntryImpl(
 
   // Track the timestamp when we started computing the result.
   const startTime = performance.timeOrigin + performance.now()
-  // Invoke the inner function to load a new result.
-  const result = fn.apply(null, args)
+
+  // Invoke the inner function to load a new result. We delay the invocation
+  // though, until React awaits the promise so that React's request store (ALS)
+  // is available when the function is invoked. This allows us, for example, to
+  // capture logs so that we can later replay them.
+  const resultPromise = createLazyResult(() => fn.apply(null, args))
 
   let errors: Array<unknown> = []
 
   let timer = undefined
   const controller = new AbortController()
   if (outerWorkUnitStore?.type === 'prerender') {
-    // If we're prerendering, we give you 50 seconds to fill a cache entry. Otherwise
-    // we assume you stalled on hanging input and deopt. This needs to be lower than
-    // just the general timeout of 60 seconds.
+    // If we're prerendering, we give you 50 seconds to fill a cache entry.
+    // Otherwise we assume you stalled on hanging input and de-opt. This needs
+    // to be lower than just the general timeout of 60 seconds.
     timer = setTimeout(() => {
       controller.abort(timeoutError)
     }, 50000)
   }
 
   const stream = renderToReadableStream(
-    result,
+    resultPromise,
     clientReferenceManifest.clientModules,
     {
       environmentName: 'Cache',
@@ -488,7 +492,7 @@ export function cache(
   kind: string,
   id: string,
   boundArgsLength: number,
-  fn: any
+  fn: (...args: unknown[]) => Promise<unknown>
 ) {
   for (const [key, value] of Object.entries(
     cacheHandlerGlobal.__nextCacheHandlers || {}
@@ -819,3 +823,22 @@ export function cache(
   }[name]
   return 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)
+    },
+  }
+}

test/e2e/app-dir/use-cache/app/(no-suspense)/logs/page.tsx

@@ -0,0 +1,18 @@
+async function Bar() {
+  'use cache'
+  const date = new Date().toLocaleTimeString()
+  console.log('deep inside', date)
+  return <p>{date}</p>
+}
+
+async function Foo() {
+  'use cache'
+  console.log('inside')
+  return <Bar />
+}
+
+export default async function Page() {
+  console.log('outside')
+
+  return <Foo />
+}

test/e2e/app-dir/use-cache/use-cache.test.ts

@@ -1,5 +1,8 @@
 import { nextTestSetup } from 'e2e-utils'
 import { retry, waitFor } from 'next-test-utils'
+import stripAnsi from 'strip-ansi'
+import { format } from 'util'
+import { BrowserInterface } from 'next-webdriver'
 
 const GENERIC_RSC_ERROR =
   'An error occurred in the Server Components render. The specific message is omitted in production builds to avoid leaking sensitive details. A digest property is included on this error instance which may provide additional details about the nature of the error.'
@@ -93,8 +96,7 @@ describe('use-cache', () => {
     const browser = await next.browser('/react-cache')
     const a = await browser.waitForElementByCss('#a').text()
     const b = await browser.waitForElementByCss('#b').text()
-    // TODO: This is broken. It is expected to pass once we fix it.
-    expect(a).not.toBe(b)
+    expect(a).toBe(b)
   })
 
   it('should error when cookies/headers/draftMode is used inside "use cache"', async () => {
@@ -313,30 +315,33 @@ describe('use-cache', () => {
     })
   })
 
-  it('should be able to revalidate a page using unstable_expireTag', async () => {
-    const browser = await next.browser(`/form`)
-    const time1 = await browser.waitForElementByCss('#t').text()
+  // TODO(useCache): Re-activate for deploy tests when NAR-85 is resolved.
+  if (!isNextDeploy) {
+    it('should be able to revalidate a page using unstable_expireTag', async () => {
+      const browser = await next.browser(`/form`)
+      const time1 = await browser.waitForElementByCss('#t').text()
 
-    await browser.loadPage(new URL(`/form`, next.url).toString())
+      await browser.loadPage(new URL(`/form`, next.url).toString())
 
-    const time2 = await browser.waitForElementByCss('#t').text()
+      const time2 = await browser.waitForElementByCss('#t').text()
 
-    expect(time1).toBe(time2)
+      expect(time1).toBe(time2)
 
-    await browser.elementByCss('#refresh').click()
+      await browser.elementByCss('#refresh').click()
 
-    await waitFor(500)
+      await waitFor(500)
 
-    const time3 = await browser.waitForElementByCss('#t').text()
+      const time3 = await browser.waitForElementByCss('#t').text()
 
-    expect(time3).not.toBe(time2)
+      expect(time3).not.toBe(time2)
 
-    // Reloading again should ideally be the same value but because the Action seeds
-    // the cache with real params as the argument it has a different cache key.
-    // await browser.loadPage(new URL(`/form?c`, next.url).toString())
-    // const time4 = await browser.waitForElementByCss('#t').text()
-    // expect(time4).toBe(time3);
-  })
+      // Reloading again should ideally be the same value but because the Action seeds
+      // the cache with real params as the argument it has a different cache key.
+      // await browser.loadPage(new URL(`/form?c`, next.url).toString())
+      // const time4 = await browser.waitForElementByCss('#t').text()
+      // expect(time4).toBe(time3);
+    })
+  }
 
   it('should use revalidate config in fetch', async () => {
     const browser = await next.browser('/fetch-revalidate')
@@ -475,5 +480,61 @@ describe('use-cache', () => {
       )
       expect(next.cliOutput).not.toContain('HANGING_PROMISE_REJECTION')
     })
+
+    it('replays logs from "use cache" functions', async () => {
+      const browser = await next.browser('/logs')
+      const initialLogs = await getSanitizedLogs(browser)
+
+      // We ignore the logged time string at the end of this message:
+      const logMessageWithDateRegexp =
+        /^ Server {3}Cache {3}Cache {2}deep inside /
+
+      let logMessageWithCachedDate: string | undefined
+
+      await retry(async () => {
+        // TODO(veil): We might want to show only the original (right-most)
+        // environment badge when caches are nested.
+        expect(initialLogs).toMatchObject(
+          expect.arrayContaining([
+            ' Server  outside',
+            ' Server   Cache  inside',
+            expect.stringMatching(logMessageWithDateRegexp),
+          ])
+        )
+
+        logMessageWithCachedDate = initialLogs.find((log) =>
+          logMessageWithDateRegexp.test(log)
+        )
+
+        expect(logMessageWithCachedDate).toBeDefined()
+      })
+
+      // Load the page again and expect the cached logs to be replayed again.
+      // We're using an explicit `loadPage` instead of `refresh` here, to start
+      // with an empty set of logs.
+      await browser.loadPage(await browser.url())
+
+      await retry(async () => {
+        const newLogs = await getSanitizedLogs(browser)
+
+        expect(newLogs).toMatchObject(
+          expect.arrayContaining([
+            ' Server  outside',
+            ' Server   Cache  inside',
+            logMessageWithCachedDate,
+          ])
+        )
+      })
+    })
   }
 })
+
+async function getSanitizedLogs(browser: BrowserInterface): Promise<string[]> {
+  const logs = await browser.log({ includeArgs: true })
+
+  return logs.map(({ args }) =>
+    format(
+      ...args.map((arg) => (typeof arg === 'string' ? stripAnsi(arg) : arg))
+    )
+  )
+}

test/lib/browsers/base.ts

@@ -119,7 +119,13 @@ export abstract class BrowserInterface<TCurrent = any> {
   abstract text(): Promise<string>
   abstract getComputedCss(prop: string): Promise<string>
   abstract hasElementByCssSelector(selector: string): Promise<boolean>
-  abstract log(): Promise<{ source: string; message: string }[]>
+  abstract log<T extends boolean = false>(options?: {
+    includeArgs?: T
+  }): Promise<
+    T extends true
+      ? { source: string; message: string; args: unknown[] }[]
+      : { source: string; message: string }[]
+  >
   abstract websocketFrames(): Promise<any[]>
   abstract url(): Promise<string>
   abstract waitForIdleNetwork(): Promise<void>

test/lib/browsers/playwright.ts

@@ -12,11 +12,13 @@ import {
 } from 'playwright'
 import path from 'path'
 
+type PageLog = { source: string; message: string; args: unknown[] }
+
 let page: Page
 let browser: Browser
 let context: BrowserContext
 let contextHasJSEnabled: boolean = true
-let pageLogs: Array<{ source: string; message: string }> = []
+let pageLogs: Array<Promise<PageLog> | PageLog> = []
 let websocketFrames: Array<{ payload: string | Buffer }> = []
 
 const tracePlaywright = process.env.TRACE_PLAYWRIGHT
@@ -223,7 +225,12 @@ export class Playwright extends BrowserInterface {
 
     page.on('console', (msg) => {
       console.log('browser log:', msg)
-      pageLogs.push({ source: msg.type(), message: msg.text() })
+
+      pageLogs.push(
+        Promise.all(
+          msg.args().map((handle) => handle.jsonValue().catch(() => {}))
+        ).then((args) => ({ source: msg.type(), message: msg.text(), args }))
+      )
     })
     page.on('crash', () => {
       console.error('page crashed')
@@ -232,7 +239,7 @@ export class Playwright extends BrowserInterface {
       console.error('page error', error)
 
       if (opts?.pushErrorAsConsoleLog) {
-        pageLogs.push({ source: 'error', message: error.message })
+        pageLogs.push({ source: 'error', message: error.message, args: [] })
       }
     })
     page.on('request', (req) => {
@@ -474,8 +481,26 @@ export class Playwright extends BrowserInterface {
     return page.evaluate<T>(fn).catch(() => null)
   }
 
-  async log() {
-    return this.chain(() => pageLogs)
+  async log<T extends boolean = false>(options?: {
+    includeArgs?: T
+  }): Promise<
+    T extends true
+      ? { source: string; message: string; args: unknown[] }[]
+      : { source: string; message: string }[]
+  > {
+    return this.chain(
+      () =>
+        options?.includeArgs
+          ? Promise.all(pageLogs)
+          : Promise.all(pageLogs).then((logs) =>
+              logs.map(({ source, message }) => ({ source, message }))
+            )
+      // TODO: Starting with TypeScript 5.8 we might not need this type cast.
+    ) as Promise<
+      T extends true
+        ? { source: string; message: string; args: unknown[] }[]
+        : { source: string; message: string }[]
+    >
   }
 
   async websocketFrames() {