refactor types for dev-only RequestStore properties

Author: lubieowoce

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

@@ -16,6 +16,7 @@ import {
   type WorkStore,
 } from '../app-render/work-async-storage.external'
 import type {
+  DevStoreModernPartial,
   PrerenderStoreModernRuntime,
   RequestStore,
 } from '../app-render/work-unit-async-storage.external'
@@ -2826,18 +2827,20 @@ async function renderWithRestartOnCacheMissInDev(
     initialHangingPromiseController.signal
   )
 
-  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.stagedRendering = initialStageController
-  requestStore.asyncApiPromises = createAsyncApiPromisesInDev(
-    initialStageController,
-    requestStore.cookies,
-    requestStore.mutableCookies,
-    requestStore.headers
-  )
-  requestStore.cacheSignal = cacheSignal
+  Object.assign(requestStore, {
+    stagedRendering: initialStageController,
+    asyncApiPromises: createAsyncApiPromisesInDev(
+      initialStageController,
+      requestStore.cookies,
+      requestStore.mutableCookies,
+      requestStore.headers
+    ),
+    prerenderResumeDataCache,
+    cacheSignal,
+  } satisfies DevStoreModernPartial)
 
   let debugChannel = setReactDebugChannel && createDebugChannel()
 
@@ -2938,18 +2941,20 @@ async function renderWithRestartOnCacheMissInDev(
 
   // We've filled the caches, so now we can render as usual,
   // without any cache-filling mechanics.
-  requestStore.prerenderResumeDataCache = null
   requestStore.renderResumeDataCache = createRenderResumeDataCache(
     prerenderResumeDataCache
   )
-  requestStore.stagedRendering = finalStageController
-  requestStore.cacheSignal = null
-  requestStore.asyncApiPromises = createAsyncApiPromisesInDev(
-    finalStageController,
-    requestStore.cookies,
-    requestStore.mutableCookies,
-    requestStore.headers
-  )
+  Object.assign(requestStore, {
+    stagedRendering: finalStageController,
+    asyncApiPromises: createAsyncApiPromisesInDev(
+      finalStageController,
+      requestStore.cookies,
+      requestStore.mutableCookies,
+      requestStore.headers
+    ),
+    prerenderResumeDataCache: null,
+    cacheSignal: null,
+  } satisfies DevStoreModernPartial)
 
   // The initial render already wrote to its debug channel.
   // We're not using it, so we need to create a new one.

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

@@ -29,7 +29,7 @@ export interface CommonWorkUnitStore {
   readonly implicitTags: ImplicitTags
 }
 
-export interface RequestStore extends CommonWorkUnitStore {
+interface BaseRequestStore extends CommonWorkUnitStore {
   readonly type: 'request'
 
   /**
@@ -65,14 +65,46 @@ export interface RequestStore extends CommonWorkUnitStore {
    * The resume data cache for this request. This will be a immutable cache.
    */
   renderResumeDataCache: RenderResumeDataCache | null
+}
+
+export type RequestStore = ProdRequestStore | DevRequestStore
+
+export type ProdRequestStore = BaseRequestStore & AllMissing<DevStore>
+export type DevRequestStore = BaseRequestStore & DevStore
 
-  // DEV-only
+// If `cacheComponents` is enabled, we add multiple extra properties on the store.
+// We either want all of them to be present, or all of them to be undefined.
+// Note that we don't want to remove the properties altogether, as in `{ a: A } | { a: A, b: B }`,
+// because then typescript wouldn prevent us from doing `if (requestStore.someProp) { ... }` --
+// we'd need to check `if ('someProp' in requestStore && requestStore.someProp) { ... }` instead,
+// which is annoying to write everywhere.
+type DevStore = DevStoreLegacy | DevStoreModern
+type DevStoreLegacy = DevStoreCommon & AllMissing<DevStoreModernPartial>
+type DevStoreModern = DevStoreCommon & DevStoreModernPartial
+
+type DevStoreCommon = {
   usedDynamic?: boolean
   devFallbackParams?: OpaqueFallbackRouteParams | null
-  stagedRendering?: StagedRenderingController | null
-  asyncApiPromises?: DevAsyncApiPromises
-  cacheSignal?: CacheSignal | null
-  prerenderResumeDataCache?: PrerenderResumeDataCache | null
+}
+
+export type DevStoreModernPartial = {
+  readonly stagedRendering: StagedRenderingController
+  readonly asyncApiPromises: DevAsyncApiPromises
+} & (
+  | {
+      // In the initial render, we track and fill caches
+      readonly cacheSignal: CacheSignal
+      readonly prerenderResumeDataCache: PrerenderResumeDataCache
+    }
+  | {
+      // In the final (restarted) render, we do not track or fill caches
+      readonly cacheSignal: null
+      readonly prerenderResumeDataCache: null
+    }
+)
+
+type AllMissing<TObj extends Record<string, any>> = {
+  [key in keyof TObj]?: undefined
 }
 
 type DevAsyncApiPromises = {

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

@@ -17,6 +17,7 @@ import { prerender } from 'react-server-dom-webpack/static'
 import type { WorkStore } from '../app-render/work-async-storage.external'
 import { workAsyncStorage } from '../app-render/work-async-storage.external'
 import type {
+  DevRequestStore,
   PrerenderStoreModernClient,
   PrerenderStoreModernRuntime,
   PrivateUseCacheStore,
@@ -1848,7 +1849,7 @@ function isRecentlyRevalidatedTag(tag: string, workStore: WorkStore): boolean {
 
 async function delayBeforeCacheReadStartInDev(
   stage: NonStaticRenderStage,
-  requestStore: RequestStore
+  requestStore: DevRequestStore
 ): Promise<void> {
   const { stagedRendering } = requestStore
   if (stagedRendering && stagedRendering.currentStage < stage) {
@@ -1859,7 +1860,7 @@ async function delayBeforeCacheReadStartInDev(
 /** Note: Only call this after a `cacheSignal.beginRead()`. */
 async function delayOrHangStartedCacheReadInDev(
   stage: NonStaticRenderStage,
-  requestStore: RequestStore,
+  requestStore: DevRequestStore,
   cacheSignal: CacheSignal | null,
   route: string,
   expression: string