Use onPostpone to determine if segment prefetch is partial

.changeset/shaggy-pears-tell.md

@@ -0,0 +1,5 @@
+---
+"next": patch
+---
+
+Use `onPostpone` to determine if segment prefetch is partial

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

@@ -4258,27 +4258,8 @@ async function collectSegmentData(
     serverModuleMap: null,
   }
 
-  // When dynamicIO is enabled, missing data is encoded to an infinitely hanging
-  // promise, the absence of which we use to determine if a segment is fully
-  // static or partially static. However, when dynamicIO is not enabled, this
-  // trick doesn't work.
-  //
-  // So if PPR is enabled, and dynamicIO is not, we have to be conservative and
-  // assume all segments are partial.
-  //
-  // TODO: When PPR is on, we can at least optimize the case where the entire
-  // page is static. Either by passing that as an argument to this function, or
-  // by setting a header on the response like the we do for full page RSC
-  // prefetches today. The latter approach might be simpler since it requires
-  // less plumbing, and the client has to check the header regardless to see if
-  // PPR is enabled.
-  const shouldAssumePartialData =
-    renderOpts.experimental.isRoutePPREnabled === true && // PPR is enabled
-    !renderOpts.experimental.dynamicIO // dynamicIO is disabled
-
   const staleTime = prerenderStore.stale
   return await ComponentMod.collectSegmentData(
-    shouldAssumePartialData,
     fullPageDataBuffer,
     staleTime,
     clientReferenceManifest.clientModules as ManifestNode,

packages/next/src/server/app-render/collect-segment-data.tsx

@@ -74,7 +74,6 @@ function onSegmentPrerenderError(error: unknown) {
 }
 
 export async function collectSegmentData(
-  shouldAssumePartialData: boolean,
   fullPageDataBuffer: Buffer,
   staleTime: number,
   clientModules: ManifestNode,
@@ -119,7 +118,6 @@ export async function collectSegmentData(
     // inside of it, the side effects are transferred to the new stream.
     // @ts-expect-error
     <PrefetchTreeData
-      shouldAssumePartialData={shouldAssumePartialData}
       fullPageDataBuffer={fullPageDataBuffer}
       fallbackRouteParams={fallbackRouteParams}
       serverConsumerManifest={serverConsumerManifest}
@@ -150,7 +148,6 @@ export async function collectSegmentData(
 }
 
 async function PrefetchTreeData({
-  shouldAssumePartialData,
   fullPageDataBuffer,
   fallbackRouteParams,
   serverConsumerManifest,
@@ -159,7 +156,6 @@ async function PrefetchTreeData({
   segmentTasks,
   onCompletedProcessingRouteTree,
 }: {
-  shouldAssumePartialData: boolean
   fullPageDataBuffer: Buffer
   serverConsumerManifest: any
   fallbackRouteParams: FallbackRouteParams | null
@@ -199,7 +195,6 @@ async function PrefetchTreeData({
   // walk the tree, we will also spawn a task to produce a prefetch response for
   // each segment.
   const tree = collectSegmentDataImpl(
-    shouldAssumePartialData,
     flightRouterState,
     buildId,
     seedData,
@@ -211,8 +206,7 @@ async function PrefetchTreeData({
     segmentTasks
   )
 
-  const isHeadPartial =
-    shouldAssumePartialData || (await isPartialRSCData(head, clientModules))
+  const isHeadPartial = await isPartialRSCData(head, clientModules)
 
   // Notify the abort controller that we're done processing the route tree.
   // Anything async that happens after this point must be due to hanging
@@ -231,7 +225,6 @@ async function PrefetchTreeData({
 }
 
 function collectSegmentDataImpl(
-  shouldAssumePartialData: boolean,
   route: FlightRouterState,
   buildId: string,
   seedData: CacheNodeSeedData | null,
@@ -265,7 +258,6 @@ function collectSegmentDataImpl(
         : encodeSegment(childSegment)
     )
     const childTree = collectSegmentDataImpl(
-      shouldAssumePartialData,
       childRoute,
       buildId,
       childSeedData,
@@ -288,13 +280,7 @@ function collectSegmentDataImpl(
       // Since we're already in the middle of a render, wait until after the
       // current task to escape the current rendering context.
       waitAtLeastOneReactRenderTask().then(() =>
-        renderSegmentPrefetch(
-          shouldAssumePartialData,
-          buildId,
-          seedData,
-          key,
-          clientModules
-        )
+        renderSegmentPrefetch(buildId, seedData, key, clientModules)
       )
     )
   } else {
@@ -344,7 +330,6 @@ function encodeSegmentWithPossibleFallbackParam(
 }
 
 async function renderSegmentPrefetch(
-  shouldAssumePartialData: boolean,
   buildId: string,
   seedData: CacheNodeSeedData,
   key: string,
@@ -359,8 +344,7 @@ async function renderSegmentPrefetch(
     buildId,
     rsc,
     loading,
-    isPartial:
-      shouldAssumePartialData || (await isPartialRSCData(rsc, clientModules)),
+    isPartial: await isPartialRSCData(rsc, clientModules),
   }
   // Since all we're doing is decoding and re-encoding a cached prerender, if
   // it takes longer than a microtask, it must because of hanging promises
@@ -403,6 +387,11 @@ async function isPartialRSCData(
   await prerender(rsc, clientModules, {
     signal: abortController.signal,
     onError() {},
+    onPostpone() {
+      // If something postponed, i.e. when Dynamic IO is not enabled, we can
+      // infer that the RSC data is partial.
+      isPartial = true
+    },
   })
   return isPartial
 }

test/e2e/app-dir/segment-cache/client-only-opt-in/next.config.js

@@ -4,7 +4,6 @@
 const nextConfig = {
   experimental: {
     ppr: 'incremental',
-    dynamicIO: true,
     clientSegmentCache: 'client-only',
   },
 }

test/e2e/app-dir/segment-cache/export/next.config.js

@@ -4,8 +4,6 @@
 const nextConfig = {
   output: 'export',
   experimental: {
-    ppr: false,
-    dynamicIO: true,
     clientSegmentCache: true,
   },
 }

test/e2e/app-dir/segment-cache/incremental-opt-in/next.config.js

@@ -4,7 +4,6 @@
 const nextConfig = {
   experimental: {
     ppr: 'incremental',
-    dynamicIO: true,
     clientSegmentCache: true,
   },
 }

test/ppr-tests-manifest.json

@@ -97,6 +97,7 @@
       "test/e2e/app-dir/static-shell-debugging/static-shell-debugging.test.ts",
       "test/e2e/app-dir/dynamic-io-errors/dynamic-io-errors.prospective-fallback.test.ts",
       "test/e2e/app-dir/segment-cache/basic/segment-cache-basic.test.ts",
+      "test/e2e/app-dir/segment-cache/export/segment-cache-output-export.test.ts",
       "test/e2e/app-dir/segment-cache/incremental-opt-in/segment-cache-incremental-opt-in.test.ts",
       "test/e2e/app-dir/segment-cache/memory-pressure/segment-cache-memory-pressure.test.ts",
       "test/e2e/app-dir/segment-cache/prefetch-scheduling/prefetch-scheduling.test.ts",