clean up CacheSignal changes

Author: lubieowoce

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

@@ -200,7 +200,6 @@ 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'
@@ -733,7 +732,7 @@ async function prospectiveRuntimeServerPrerender(
 
   // The cacheSignal helps us track whether caches are still filling or we are ready
   // to cut the render off.
-  const cacheSignal = new CacheSignal()
+  const cacheSignal = new CacheSignal('prerender')
 
   const initialServerPrerenderStore: PrerenderStoreModernRuntime = {
     type: 'prerender-runtime',
@@ -2184,12 +2183,12 @@ async function renderToStream(
 
       // 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 cacheSignal = new CacheSignal('render')
 
       // 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)
+      trackPendingModules(cacheSignal)
 
       const prerenderResumeDataCache = createPrerenderResumeDataCache()
 
@@ -2265,11 +2264,7 @@ async function renderToStream(
         // 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()
+        await cacheSignal.cacheReady()
         initialRenderReactController.abort()
 
         //===============================================
@@ -2805,7 +2800,7 @@ async function spawnDynamicValidationInDev(
 
   // The cacheSignal helps us track whether caches are still filling or we are
   // ready to cut the render off.
-  const cacheSignal = new CacheSignal()
+  const cacheSignal = new CacheSignal('prerender')
 
   const captureOwnerStackClient = React.captureOwnerStack
   const captureOwnerStackServer = ComponentMod.captureOwnerStack
@@ -3545,7 +3540,7 @@ async function prerenderToStream(
 
       // The cacheSignal helps us track whether caches are still filling or we are ready
       // to cut the render off.
-      const cacheSignal = new CacheSignal()
+      const cacheSignal = new CacheSignal('prerender')
 
       let resumeDataCache: RenderResumeDataCache | PrerenderResumeDataCache
       let renderResumeDataCache: RenderResumeDataCache | null = null

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

@@ -8,6 +8,8 @@
 import { waitAtLeastOneReactRenderTask } from '../../lib/scheduler'
 import { InvariantError } from '../../shared/lib/invariant-error'
 
+type CacheSignalMode = 'prerender' | 'render'
+
 export class CacheSignal {
   private count = 0
   private earlyListeners: Array<() => void> = []
@@ -17,7 +19,7 @@ export class CacheSignal {
 
   private subscribedSignals: Set<CacheSignal> | null = null
 
-  constructor() {
+  constructor(private mode: CacheSignalMode) {
     if (process.env.NEXT_RUNTIME === 'edge') {
       // we rely on `process.nextTick`, which is not supported in edge
       throw new InvariantError(
@@ -71,7 +73,23 @@ export class CacheSignal {
    * if there are no inflight cache reads then we wait at least one task to allow initial
    * cache reads to be initiated.
    */
-  cacheReady() {
+  async cacheReady(): Promise<void> {
+    if (this.mode === 'prerender') {
+      return await this.cacheReadyImpl()
+    } else {
+      // 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.cacheReadyImpl()
+        await waitAtLeastOneReactRenderTask()
+      } while (this.hasPendingReads())
+    }
+  }
+
+  private cacheReadyImpl() {
     return new Promise<void>((resolve) => {
       this.listeners.push(resolve)
       if (this.count === 0) {
@@ -80,23 +98,6 @@ 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++
 

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

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

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

@@ -8,7 +8,9 @@ import { isThenable } from '../../../shared/lib/is-thenable'
 let _moduleLoadingSignal: CacheSignal | null
 function getModuleLoadingSignal() {
   if (!_moduleLoadingSignal) {
-    _moduleLoadingSignal = new CacheSignal()
+    // The mode doesn't really matter here, because we never wait for `cacheReady()` on this signal,
+    // only other signals that subscribe to it via `trackPendingModules`.
+    _moduleLoadingSignal = new CacheSignal('prerender')
   }
   return _moduleLoadingSignal
 }
@@ -52,21 +54,3 @@ export function trackPendingModules(cacheSignal: CacheSignal): void {
   // we can unsubscribe it.
   cacheSignal.cacheReady().then(unsubscribe)
 }
-
-/**
- * Like `trackPendingModules`, but for use in a render, not a prerender.
- * See `CacheSignal#cacheReadyInRender` for an explanation.
- */
-export function trackPendingModulesInRender(cacheSignal: CacheSignal): void {
-  const moduleLoadingSignal = getModuleLoadingSignal()
-
-  // We can't just use `cacheSignal.trackRead(moduleLoadingSignal.cacheReady())`,
-  // because we might start and finish multiple batches of module loads while waiting for caches,
-  // and `moduleLoadingSignal.cacheReady()` would resolve after the first batch.
-  // Instead, we'll keep notifying `cacheSignal` of each import/chunk-load.
-  const unsubscribe = moduleLoadingSignal.subscribeToReads(cacheSignal)
-
-  // Later, when `cacheSignal` is no longer waiting for any caches (or imports that we've notified it of),
-  // we can unsubscribe it.
-  cacheSignal.cacheReadyInRender().then(unsubscribe)
-}

packages/next/src/server/route-modules/app-route/module.ts

@@ -376,7 +376,7 @@ export class AppRouteRouteModule extends RouteModule<
            */
           const prospectiveController = new AbortController()
           let prospectiveRenderIsDynamic = false
-          const cacheSignal = new CacheSignal()
+          const cacheSignal = new CacheSignal('prerender')
           let dynamicTracking = createDynamicTrackingState(undefined)
 
           // TODO: Route handlers are never resumed, so it's counter-intuitive