wip

Author: lubieowoce

packages/next/errors.json

@@ -852,5 +852,6 @@
   "851": "Pass either `webpack` or `turbopack`, not both.",
   "852": "Only custom servers can pass `webpack`, `turbo`, or `turbopack`.",
   "853": "Turbopack build failed",
-  "854": "Expected a %s request header."
+  "854": "Expected a %s request header.",
+  "855": "`pipelineInSequentialTasks` should not be called in edge runtime."
 }

packages/next/src/server/app-render/app-render-render-utils.ts

@@ -29,3 +29,40 @@ export function scheduleInSequentialTasks<R>(
     })
   }
 }
+
+/**
+ * This is a utility function to make scheduling sequential tasks that run back to back easier.
+ * We schedule on the same queue (setTimeout) at the same time to ensure no other events can sneak in between.
+ * The function that runs in the second task gets access to the first tasks's result.
+ */
+export function pipelineInSequentialTasks<A, B>(
+  render: () => A,
+  followup: (a: A) => B | Promise<B>
+): Promise<B> {
+  if (process.env.NEXT_RUNTIME === 'edge') {
+    throw new InvariantError(
+      '`pipelineInSequentialTasks` should not be called in edge runtime.'
+    )
+  } else {
+    return new Promise((resolve, reject) => {
+      let renderResult: A | undefined = undefined
+      setTimeout(() => {
+        try {
+          renderResult = render()
+        } catch (err) {
+          clearTimeout(followupId)
+          reject(err)
+        }
+      }, 0)
+      const followupId = setTimeout(() => {
+        // if `render` threw, then the `followup` timeout would've been cleared,
+        // so if we got here, we're guaranteed to have a `renderResult`.
+        try {
+          resolve(followup(renderResult!))
+        } catch (err) {
+          reject(err)
+        }
+      }, 0)
+    })
+  }
+}

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

@@ -168,7 +168,10 @@ import {
   prerenderAndAbortInSequentialTasks,
 } from './app-render-prerender-utils'
 import { printDebugThrownValueForProspectiveRender } from './prospective-render-utils'
-import { scheduleInSequentialTasks } from './app-render-render-utils'
+import {
+  pipelineInSequentialTasks,
+  scheduleInSequentialTasks,
+} from './app-render-render-utils'
 import { waitAtLeastOneReactRenderTask } from '../../lib/scheduler'
 import {
   workUnitAsyncStorage,
@@ -197,6 +200,7 @@ import {
   trackPendingChunkLoad,
   trackPendingImport,
   trackPendingModules,
+  trackPendingModulesInRender,
 } from './module-loading/track-module-loading.external'
 import { isReactLargeShellError } from './react-large-shell-error'
 import type { GlobalErrorComponent } from '../../client/components/builtin/global-error'
@@ -2122,57 +2126,198 @@ 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 = async (): Promise<RSCPayloadWithValidation> => {
+        const payload: RSCPayloadWithValidation =
+          await workUnitAsyncStorage.run(
+            requestStore,
+            getRSCPayload,
+            tree,
+            ctx,
+            res.statusCode === 404
+          )
+        // Placing the validation outlet in the payload is safe
+        // even if we end up discarding a render and restarting,
+        // because we're not going to wait for the stream to complete,
+        // so leaving the validation unresolved is fine.
+        payload._validation = validationOutlet
+        return payload
+      }
 
-      if (debugChannel) {
+      const setDebugChannelForClientRender = (
+        debugChannel: DebugChannelPair
+      ) => {
         const [readableSsr, readableBrowser] =
           debugChannel.clientSide.readable.tee()
 
         reactDebugStream = readableSsr
 
-        setReactDebugChannel(
+        setReactDebugChannel!(
           { readable: readableBrowser },
           htmlRequestId,
           requestId
         )
       }
 
-      const reactServerStream = await workUnitAsyncStorage.run(
+      const environmentName = () =>
+        requestStore.prerenderPhase === true ? 'Prerender' : 'Server'
+
+      // 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.
+
+      // 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()
+
+      // If we encounter async modules that delay rendering, we'll also need to restart.
+      // TODO(restart-on-cache-miss): technically, we only need to wait for pending *server* modules here,
+      // but `trackPendingModules` doesn't distinguish between client and server.
+      trackPendingModulesInRender(cacheSignal)
+
+      const prerenderResumeDataCache = createPrerenderResumeDataCache()
+
+      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 intialRenderDebugChannel =
+        setReactDebugChannel && createDebugChannel()
+
+      const initialRscPayload = await getPayload()
+      const maybeInitialServerStream = 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,
+        () =>
+          pipelineInSequentialTasks(
+            () => {
+              // Static stage
+              requestStore.prerenderPhase = true
+              return ComponentMod.renderToReadableStream(
+                initialRscPayload,
+                clientReferenceManifest.clientModules,
+                {
+                  onError: serverComponentsErrorHandler,
+                  environmentName,
+                  filterStackFrame,
+                  debugChannel: intialRenderDebugChannel?.serverSide,
+                  signal: initialRenderReactController.signal,
+                }
+              )
+            },
+            async (stream) => {
+              // Dynamic stage
+              // Note: if we had cache misses, things that would've happened statically otherwise
+              // may be marked as dynamic instead.
+              requestStore.prerenderPhase = false
+
+              // If all cache reads initiated in the static stage have completed,
+              // then all of the necessary caches have to be warm (or there's no caches on the page).
+              // On the other hand, if we still have pending cache reads, then we had a cache miss,
+              // and the static stage didn't render all the content that it normally would have.
+              const hadCacheMiss = cacheSignal.hasPendingReads()
+              if (!hadCacheMiss) {
+                // No cache misses. We can use the stream as is.
+                return stream
+              } else {
+                // Cache miss. We'll discard this stream, and render again.
+                return null
+              }
             }
           )
-        },
-        () => {
-          requestStore.prerenderPhase = false
-        }
       )
 
+      if (maybeInitialServerStream !== null) {
+        // No cache misses. We can use the stream as is.
+
+        // We're using this render, so we should pass its debug channel to the client render.
+        if (intialRenderDebugChannel) {
+          setDebugChannelForClientRender(intialRenderDebugChannel)
+        }
+
+        reactServerResult = new ReactServerResult(maybeInitialServerStream)
+      } 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()
+        initialRenderReactController.abort()
+
+        //===============================================
+
+        // The initial render acted as a prospective render to warm the caches.
+        // Now, we need to do another render.
+
+        // TODO(restart-on-cache-miss): we should use a separate request store for this instead
+
+        // We've filled the caches, so now we can render as usual.
+        requestStore.prerenderResumeDataCache = null
+        requestStore.renderResumeDataCache = createRenderResumeDataCache(
+          prerenderResumeDataCache
+        )
+        requestStore.cacheSignal = null
+
+        // Reset mutable fields.
+        requestStore.prerenderPhase = undefined
+        requestStore.usedDynamic = undefined
+
+        // The initial render already wrote to its debug channel. We're not using it,
+        // so we need to create a new one.
+        const finalRenderDebugChannel =
+          setReactDebugChannel && createDebugChannel()
+        // We know that we won't discard this render, so we can set the debug channel up immediately.
+        if (finalRenderDebugChannel) {
+          setDebugChannelForClientRender(finalRenderDebugChannel)
+        }
+
+        const finalRscPayload = await getPayload()
+        const finalServerStream = await workUnitAsyncStorage.run(
+          requestStore,
+          scheduleInSequentialTasks,
+          () => {
+            // Static stage
+            requestStore.prerenderPhase = true
+            return ComponentMod.renderToReadableStream(
+              finalRscPayload,
+              clientReferenceManifest.clientModules,
+              {
+                onError: serverComponentsErrorHandler,
+                environmentName,
+                filterStackFrame,
+                debugChannel: finalRenderDebugChannel?.serverSide,
+              }
+            )
+          },
+          () => {
+            // Dynamic stage
+            requestStore.prerenderPhase = false
+          }
+        )
+
+        reactServerResult = new ReactServerResult(finalServerStream)
+      }
+
+      // 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.
       consoleAsyncStorage.run(
         { dim: true },
         spawnDynamicValidationInDev,
@@ -2184,8 +2329,6 @@ async function renderToStream(
         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(
@@ -2537,12 +2680,21 @@ async function renderToStream(
   }
 }
 
-function createDebugChannel():
-  | {
-      serverSide: { readable?: ReadableStream; writable: WritableStream }
-      clientSide: { readable: ReadableStream; writable?: WritableStream }
-    }
-  | undefined {
+type DebugChannelPair = {
+  serverSide: DebugChannelServer
+  clientSide: DebugChannelClient
+}
+
+type DebugChannelServer = {
+  readable?: ReadableStream<Uint8Array>
+  writable: WritableStream<Uint8Array>
+}
+type DebugChannelClient = {
+  readable: ReadableStream<Uint8Array>
+  writable?: WritableStream<Uint8Array>
+}
+
+function createDebugChannel(): DebugChannelPair | undefined {
   if (process.env.NODE_ENV === 'production') {
     return undefined
   }

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.hasPendingReads())
+  }
+
   beginRead() {
     this.count++
 
@@ -114,6 +132,10 @@ export class CacheSignal {
     }
   }
 
+  hasPendingReads(): boolean {
+    return this.count > 0
+  }
+
   trackRead<T>(promise: Promise<T>) {
     this.beginRead()
     // `promise.finally()` still rejects, so don't use it here to avoid unhandled rejections

packages/next/src/server/app-render/module-loading/track-module-loading.external.ts

@@ -6,6 +6,12 @@ import {
   trackPendingChunkLoad,
   trackPendingImport,
   trackPendingModules,
+  trackPendingModulesInRender,
 } from './track-module-loading.instance' with { 'turbopack-transition': 'next-shared' }
 
-export { trackPendingChunkLoad, trackPendingImport, trackPendingModules }
+export {
+  trackPendingChunkLoad,
+  trackPendingImport,
+  trackPendingModules,
+  trackPendingModulesInRender,
+}