wip

Author: lubieowoce

packages/next/src/build/templates/app-page.ts

@@ -409,7 +409,8 @@ export async function handler(
         if (
           nextConfig.experimental.cacheComponents &&
           !isPrefetchRSCRequest &&
-          !context.renderOpts.isPossibleServerAction
+          !context.renderOpts.isPossibleServerAction &&
+          process.env.NEXT_RESTART_ON_CACHE_MISS === '0'
         ) {
           const warmup = await routeModule.warmup(nextReq, nextRes, context)
 

packages/next/src/server/app-render/app-render.tsx

@@ -2258,70 +2258,217 @@ async function renderToStream(
       // We only have a Prerender environment for projects opted into cacheComponents
       experimental.cacheComponents
     ) {
-      // This is a dynamic render. We don't do dynamic tracking because we're not prerendering
-      const RSCPayload: InitialRSCPayload & {
+      type RSCPayloadWithValidation = InitialRSCPayload & {
         /** Only available during cacheComponents development builds. Used for logging errors. */
         _validation?: Promise<React.ReactNode>
-      } = await workUnitAsyncStorage.run(
-        requestStore,
-        getRSCPayload,
-        tree,
-        ctx,
-        res.statusCode === 404
-      )
-      const [resolveValidation, validationOutlet] = createValidationOutlet()
-      RSCPayload._validation = validationOutlet
+      }
 
-      const debugChannel = setReactDebugChannel && createDebugChannel()
+      const getPayload = (): Promise<RSCPayloadWithValidation> =>
+        workUnitAsyncStorage.run(
+          requestStore,
+          getRSCPayload,
+          tree,
+          ctx,
+          res.statusCode === 404
+        )
 
-      if (debugChannel) {
-        const [readableSsr, readableBrowser] =
-          debugChannel.clientSide.readable.tee()
+      const environmentName = () =>
+        requestStore.prerenderPhase === true ? 'Prerender' : 'Server'
 
-        reactDebugStream = readableSsr
+      if (process.env.NEXT_RESTART_ON_CACHE_MISS !== '0') {
+        // Try to render the page and see if there's any cache misses.
+        // If there are, wait for caches to finish and restart the render.
 
-        setReactDebugChannel(
-          { readable: readableBrowser },
-          htmlRequestId,
-          requestId
-        )
-      }
+        const [resolveValidation, validationOutlet] = createValidationOutlet()
 
-      const reactServerStream = await workUnitAsyncStorage.run(
-        requestStore,
-        scheduleInSequentialTasks,
-        () => {
-          requestStore.prerenderPhase = true
-          return ComponentMod.renderToReadableStream(
-            RSCPayload,
-            clientReferenceManifest.clientModules,
-            {
-              onError: serverComponentsErrorHandler,
-              environmentName: () =>
-                requestStore.prerenderPhase === true ? 'Prerender' : 'Server',
-              filterStackFrame,
-              debugChannel: debugChannel?.serverSide,
+        const renderRestartable = async (
+          signal: AbortSignal | undefined,
+          onPrerenderStageEnd: (() => void) | undefined
+        ) => {
+          const rscPayload = await getPayload()
+
+          // Placing the validation outlet in the payload is safe
+          // even if we end up discarding this render and restarting,
+          // because it's just an output produced independently.
+          rscPayload._validation = validationOutlet
+
+          return workUnitAsyncStorage.run(
+            requestStore,
+            scheduleInSequentialTasks,
+            () => {
+              // Static stage
+              requestStore.prerenderPhase = true
+              return ComponentMod.renderToReadableStream(
+                rscPayload,
+                clientReferenceManifest.clientModules,
+                {
+                  onError: serverComponentsErrorHandler,
+                  environmentName,
+                  filterStackFrame,
+                  // TODO(restart-on-cache-miss): implement `debugChannel`
+                  // debugChannel: debugChannel?.serverSide,
+                  signal,
+                }
+              )
+            },
+            () => {
+              // Dynamic stage
+              requestStore.prerenderPhase = false
+              onPrerenderStageEnd?.()
             }
           )
-        },
-        () => {
-          requestStore.prerenderPhase = false
         }
-      )
 
-      devLogsAsyncStorage.run(
-        { dim: true },
-        spawnDynamicValidationInDev,
-        resolveValidation,
-        tree,
-        ctx,
-        res.statusCode === 404,
-        clientReferenceManifest,
-        requestStore,
-        devValidatingFallbackParams
-      )
+        // This render might end up being used as a prospective render (if there's cache misses),
+        // so we need to set it up for filling caches.
+        const cacheSignal = new CacheSignal()
+        const prerenderResumeDataCache = createPrerenderResumeDataCache()
 
-      reactServerResult = new ReactServerResult(reactServerStream)
+        requestStore.prerenderResumeDataCache = prerenderResumeDataCache
+        // `getRenderResumeDataCache` will fall back to using `prerenderResumeDataCache` as `renderResumeDataCache`,
+        // so not having a resume data cache won't break any expectations in case we don't need to restart.
+        requestStore.renderResumeDataCache = null
+        requestStore.cacheSignal = cacheSignal
+
+        const initialRenderReactController = new AbortController()
+        const hadCacheMissInStaticStagePromise =
+          createPromiseWithResolvers<boolean>()
+
+        console.debug(`renderToStream (1) :: attempting render`)
+
+        const reactServerStreamPromise = renderRestartable(
+          initialRenderReactController.signal,
+          () => {
+            console.debug(
+              `renderToStream (1) :: static task finished with ${cacheSignal['count']} caches pending`
+            )
+            // If all cache reads initiated in the static stage have completed,
+            // then either we don't need to fill any caches, or all of them are warm.
+            // On the other hand, if we have pending cache reads, then we had a cache miss.
+            hadCacheMissInStaticStagePromise.resolve(
+              cacheSignal.hasPendingReads()
+            )
+          }
+        )
+        reactServerStreamPromise.catch((err) =>
+          hadCacheMissInStaticStagePromise.reject(err)
+        )
+
+        const hasCacheMissInStaticStage =
+          await hadCacheMissInStaticStagePromise.promise
+
+        if (!hasCacheMissInStaticStage) {
+          // No cache misses. Use the stream as is.
+          reactServerResult = new ReactServerResult(
+            await reactServerStreamPromise
+          )
+        } else {
+          // Cache miss. We will use the initial render to fill caches, and discard its result.
+          // Then, we can render again with warm caches.
+
+          // TODO(restart-on-cache-miss):
+          // This might end up waiting for more caches than strictly necessary,
+          // because we can't abort the render yet, and we'll let runtime/dynamic APIs resolve.
+          // Ideally we'd only wait for caches that are needed in the static stage.
+          // This will be optimized in the future by not allowing runtime/dynamic APIs to resolve.
+
+          // During a render, React pings pending tasks using `setImmediate`,
+          // and only waiting for a single `cacheReady` can make us stop filling caches too soon.
+          // To avoid this, we await `cacheReady` repeatedly with an extra delay to let React try render new content
+          // (and potentially discover more caches).
+          await cacheSignal.cacheReadyInRender()
+          console.debug(`renderToStream (1) :: cacheReady`)
+          initialRenderReactController.abort()
+
+          console.debug(
+            `renderToStream :: restarting render (cache entries: ${prerenderResumeDataCache.cache.size})`
+          )
+          // The initial render acted as a prospective render.
+          // Now, we need to clear the state we've set up for it and do a regular render.
+          requestStore.prerenderResumeDataCache = null
+          requestStore.renderResumeDataCache = createRenderResumeDataCache(
+            prerenderResumeDataCache
+          )
+          requestStore.cacheSignal = null
+
+          reactServerResult = new ReactServerResult(
+            await renderRestartable(undefined, () => {
+              console.debug(
+                `renderToStream (2) :: end of static stage after restart. ${cacheSignal['count']} caches pending`
+              )
+            })
+          )
+        }
+        // TODO(restart-on-cache-miss):
+        // This can probably be optimized to do less work,
+        // because we've already made sure that we have warm caches.
+        devLogsAsyncStorage.run(
+          { dim: true },
+          spawnDynamicValidationInDev,
+          resolveValidation,
+          tree,
+          ctx,
+          res.statusCode === 404,
+          clientReferenceManifest,
+          requestStore,
+          devValidatingFallbackParams
+        )
+      } else {
+        const rscPayload = await getPayload()
+
+        const [resolveValidation, validationOutlet] = createValidationOutlet()
+        rscPayload._validation = validationOutlet
+
+        const debugChannel = setReactDebugChannel && createDebugChannel()
+
+        if (debugChannel) {
+          const [readableSsr, readableBrowser] =
+            debugChannel.clientSide.readable.tee()
+
+          reactDebugStream = readableSsr
+
+          setReactDebugChannel(
+            { readable: readableBrowser },
+            htmlRequestId,
+            requestId
+          )
+        }
+
+        const reactServerStream = await workUnitAsyncStorage.run(
+          requestStore,
+          scheduleInSequentialTasks,
+          () => {
+            requestStore.prerenderPhase = true
+            return ComponentMod.renderToReadableStream(
+              rscPayload,
+              clientReferenceManifest.clientModules,
+              {
+                onError: serverComponentsErrorHandler,
+                environmentName,
+                filterStackFrame,
+                debugChannel: debugChannel?.serverSide,
+              }
+            )
+          },
+          () => {
+            requestStore.prerenderPhase = false
+          }
+        )
+
+        devLogsAsyncStorage.run(
+          { dim: true },
+          spawnDynamicValidationInDev,
+          resolveValidation,
+          tree,
+          ctx,
+          res.statusCode === 404,
+          clientReferenceManifest,
+          requestStore,
+          devValidatingFallbackParams
+        )
+
+        reactServerResult = new ReactServerResult(reactServerStream)
+      }
     } else {
       // This is a dynamic render. We don't do dynamic tracking because we're not prerendering
       const RSCPayload = await workUnitAsyncStorage.run(

packages/next/src/server/app-render/cache-signal.ts

@@ -5,6 +5,7 @@
  * and should only be used in codepaths gated with this feature.
  */
 
+import { waitAtLeastOneReactRenderTask } from '../../lib/scheduler'
 import { InvariantError } from '../../shared/lib/invariant-error'
 
 export class CacheSignal {
@@ -79,6 +80,23 @@ export class CacheSignal {
     })
   }
 
+  /**
+   * Like `cacheReady`, but for use when rendering (not prerendering).
+   * React schedules work differently between these two, which affects the timing
+   * of waiting for all caches to be discovered.
+   **/
+  async cacheReadyInRender() {
+    // During a render, React pings pending tasks (that are waiting for something async to resolve) using `setImmediate`.
+    // This is unlike a prerender, where they are pinged in a microtask.
+    // This means that, if we're warming caches via a render (not a prerender),
+    // we need to give React more time to continue rendering after a cache has resolved
+    // in order to make sure we've discovered all the caches needed for the current render.
+    do {
+      await this.cacheReady()
+      await waitAtLeastOneReactRenderTask()
+    } while (this.count > 0)
+  }
+
   beginRead() {
     this.count++
 

packages/next/src/server/app-render/work-unit-async-storage.external.ts

@@ -69,6 +69,8 @@ export interface RequestStore extends CommonWorkUnitStore {
   usedDynamic?: boolean
   prerenderPhase?: boolean
   devFallbackParams?: OpaqueFallbackRouteParams | null
+  cacheSignal?: CacheSignal | null
+  prerenderResumeDataCache?: PrerenderResumeDataCache | null
 }
 
 /**
@@ -351,8 +353,14 @@ export function getPrerenderResumeDataCache(
       // TODO eliminate fetch caching in client scope and stop exposing this data
       // cache during SSR.
       return workUnitStore.prerenderResumeDataCache
+    case 'request': {
+      // In dev, we might fill caches even during a dynamic request.
+      if (workUnitStore.prerenderResumeDataCache) {
+        return workUnitStore.prerenderResumeDataCache
+      }
+      // fallthrough
+    }
     case 'prerender-legacy':
-    case 'request':
     case 'cache':
     case 'private-cache':
     case 'unstable-cache':
@@ -367,7 +375,6 @@ export function getRenderResumeDataCache(
 ): RenderResumeDataCache | null {
   switch (workUnitStore.type) {
     case 'request':
-      return workUnitStore.renderResumeDataCache
     case 'prerender':
     case 'prerender-runtime':
     case 'prerender-client':
@@ -380,7 +387,7 @@ export function getRenderResumeDataCache(
     case 'prerender-ppr':
       // Otherwise we return the mutable resume data cache here as an immutable
       // version of the cache as it can also be used for reading.
-      return workUnitStore.prerenderResumeDataCache
+      return workUnitStore.prerenderResumeDataCache ?? null
     case 'cache':
     case 'private-cache':
     case 'unstable-cache':
@@ -503,9 +510,15 @@ export function getCacheSignal(
     case 'prerender-client':
     case 'prerender-runtime':
       return workUnitStore.cacheSignal
+    case 'request': {
+      // In dev, we might fill caches even during a dynamic request.
+      if (workUnitStore.cacheSignal) {
+        return workUnitStore.cacheSignal
+      }
+      // fallthrough
+    }
     case 'prerender-ppr':
     case 'prerender-legacy':
-    case 'request':
     case 'cache':
     case 'private-cache':
     case 'unstable-cache':

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

@@ -466,6 +466,8 @@ async function collectResult(
         // then it shouldn't have any effects on the prerender. We'll decide
         // whether or not this cache should have its life & tags propagated when
         // we read the entry in the final prerender from the resume data cache.
+
+        // TODO(restart-on-cache-miss): might need to do the same here?
         break
       }
       case 'request':