Repository: vercel/next.js
Pull request: Use React's `prerender` function for `"use cache"` with Dynamic IO

During prerendering, we don't need to return a stream early for a `"use cache"` function. We can instead prerender the cache function fully, before returning the prelude as the stream. This simplifies the timeout handling a bit, when `dynamicIO` is enabled, and more importantly, it prepares us for handling cookies access in `"use cache"` functions during prerendering.

- [x] depends on facebook/react#32953

Changed files:
filename: packages/next/src/server/use-cache/use-cache-wrapper.ts
patch: @@ -6,12 +6,13 @@ import {
   decodeReplyFromAsyncIterable,
   createTemporaryReferenceSet as createServerTemporaryReferenceSet,
 } from 'react-server-dom-webpack/server.edge'
-/* eslint-disable import/no-extraneous-dependencies */
 import {
   createFromReadableStream,
   encodeReply,
   createTemporaryReferenceSet as createClientTemporaryReferenceSet,
 } from 'react-server-dom-webpack/client.edge'
+import { unstable_prerender as prerender } from 'react-server-dom-webpack/static.edge'
+/* eslint-enable import/no-extraneous-dependencies */
 
 import type { WorkStore } from '../app-render/work-async-storage.external'
 import { workAsyncStorage } from '../app-render/work-async-storage.external'
@@ -224,8 +225,7 @@ async function collectResult(
   outerWorkUnitStore: WorkUnitStore | undefined,
   innerCacheStore: UseCacheStore,
   startTime: number,
-  errors: Array<unknown>, // This is a live array that gets pushed into.,
-  timer: any
+  errors: Array<unknown> // This is a live array that gets pushed into.
 ): Promise<CacheEntry> {
   // We create a buffered stream that collects all chunks until the end to
   // ensure that RSC has finished rendering and therefore we have collected
@@ -242,8 +242,13 @@ async function collectResult(
 
   const buffer: any[] = []
   const reader = savedStream.getReader()
-  for (let entry; !(entry = await reader.read()).done; ) {
-    buffer.push(entry.value)
+
+  try {
+    for (let entry; !(entry = await reader.read()).done; ) {
+      buffer.push(entry.value)
+    }
+  } catch (error) {
+    errors.push(error)
   }
 
   let idx = 0
@@ -287,21 +292,19 @@ async function collectResult(
     stale: collectedStale,
     tags: collectedTags === null ? [] : collectedTags,
   }
+
   // Propagate tags/revalidate to the parent context.
   propagateCacheLifeAndTags(outerWorkUnitStore, entry)
 
   const cacheSignal =
     outerWorkUnitStore && outerWorkUnitStore.type === 'prerender'
       ? outerWorkUnitStore.cacheSignal
       : null
+
   if (cacheSignal) {
     cacheSignal.endRead()
   }
 
-  if (timer !== undefined) {
-    clearTimeout(timer)
-  }
-
   return entry
 }
 
@@ -363,53 +366,98 @@ async function generateCacheEntryImpl(
   const resultPromise = createLazyResult(() => fn.apply(null, args))
 
   let errors: Array<unknown> = []
+  let timeoutErrorHandled = false
+
+  // In the "Cache" environment, we only need to make sure that the error
+  // digests are handled correctly. Error formatting and reporting is not
+  // necessary here; the errors are encoded in the stream, and will be reported
+  // in the "Server" environment.
+  const handleError = (error: unknown): string | undefined => {
+    const digest = getDigestForWellKnownError(error)
+
+    if (digest) {
+      return digest
+    }
+
+    if (process.env.NODE_ENV !== 'development') {
+      // TODO: For now we're also reporting the error here, because in
+      // production, the "Server" environment will only get the obfuscated
+      // error (created by the Flight Client in the cache wrapper).
+      console.error(error)
+    }
+
+    if (error === timeoutError) {
+      timeoutErrorHandled = true
+      // The timeout error already aborted the whole stream. We don't need
+      // to also push this error into the `errors` array.
+      return timeoutError.digest
+    }
+
+    errors.push(error)
+  }
+
+  let stream: ReadableStream<Uint8Array>
 
-  let timer = undefined
-  const controller = new AbortController()
   if (outerWorkUnitStore?.type === 'prerender') {
+    const timeoutAbortController = new AbortController()
+
     // 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)
+    const timer = setTimeout(() => {
+      timeoutAbortController.abort(timeoutError)
     }, 50000)
-  }
 
-  const stream = renderToReadableStream(
-    resultPromise,
-    clientReferenceManifest.clientModules,
-    {
-      environmentName: 'Cache',
-      signal: controller.signal,
-      temporaryReferences,
-      // In the "Cache" environment, we only need to make sure that the error
-      // digests are handled correctly. Error formatting and reporting is not
-      // necessary here; the errors are encoded in the stream, and will be
-      // reported in the "Server" environment.
-      onError: (error) => {
-        const digest = getDigestForWellKnownError(error)
-
-        if (digest) {
-          return digest
-        }
+    const { renderSignal } = outerWorkUnitStore
 
-        if (process.env.NODE_ENV !== 'development') {
-          // TODO: For now we're also reporting the error here, because in
-          // production, the "Server" environment will only get the obfuscated
-          // error (created by the Flight Client in the cache wrapper).
-          console.error(error)
-        }
+    const abortSignal = AbortSignal.any([
+      renderSignal,
+      timeoutAbortController.signal,
+    ])
 
-        if (error === timeoutError) {
-          // The timeout error already aborted the whole stream. We don't need
-          // to also push this error into the `errors` array.
-          return timeoutError.digest
-        }
+    const { prelude } = await prerender(
+      resultPromise,
+      clientReferenceManifest.clientModules,
+      {
+        environmentName: 'Cache',
+        signal: abortSignal,
+        temporaryReferences,
+        onError(error) {
+          if (renderSignal.aborted && renderSignal.reason === error) {
+            return undefined
+          }
 
-        errors.push(error)
-      },
+          return handleError(error)
+        },
+      }
+    )
+
+    clearTimeout(timer)
+
+    if (timeoutAbortController.signal.aborted && !timeoutErrorHandled) {
+      // When halting is enabled, the prerender will not call `onError` when
+      // it's aborted with the timeout abort signal, and hanging promises will
+      // also not be rejected. In this case, we're creating an erroring stream
+      // here, to ensure that the error is propagated to the server environment.
+      stream = new ReadableStream({
+        start(controller) {
+          controller.error(timeoutError)
+        },
+      })
+    } else {
+      stream = prelude
     }
-  )
+  } else {
+    stream = renderToReadableStream(
+      resultPromise,
+      clientReferenceManifest.clientModules,
+      {
+        environmentName: 'Cache',
+        temporaryReferences,
+        onError: handleError,
+      }
+    )
+  }
 
   const [returnStream, savedStream] = stream.tee()
 
@@ -419,8 +467,7 @@ async function generateCacheEntryImpl(
     outerWorkUnitStore,
     innerCacheStore,
     startTime,
-    errors,
-    timer
+    errors
   )
 
   // Return the stream as we're creating it. This means that if it ends up

filename: packages/next/types/$$compiled.internal.d.ts
patch: @@ -217,6 +217,8 @@ declare module 'react-server-dom-webpack/server.node' {
   ): Promise<ReactFormState | null>
 }
 declare module 'react-server-dom-webpack/static.edge' {
+  export type TemporaryReferenceSet = WeakMap<any, string>
+
   export function unstable_prerender(
     children: any,
     webpackMap: {
@@ -232,6 +234,7 @@ declare module 'react-server-dom-webpack/static.edge' {
       filterStackFrame?: (url: string, functionName: string) => boolean
       identifierPrefix?: string
       signal?: AbortSignal
+      temporaryReferences?: TemporaryReferenceSet
       onError?: (error: unknown) => void
       onPostpone?: (reason: string) => void
     }