Basic implementation of cache-wrapper with RSC serialization (vercel#…

…70435)

Co-authored-by: JJ Kasper <[email protected]>

packages/next/src/build/webpack/loaders/next-flight-loader/cache-wrapper.ts

@@ -1,3 +1 @@
-export function cache(_kind: string, _id: string, fn: any) {
-  return fn
-}
+export { cache } from '../../../../server/use-cache/use-cache-wrapper'

packages/next/src/client/components/async-local-storage.ts

@@ -40,3 +40,15 @@ export function createAsyncLocalStorage<
   }
   return new FakeAsyncLocalStorage()
 }
+
+export function createSnapshot(): <R, TArgs extends any[]>(
+  fn: (...args: TArgs) => R,
+  ...args: TArgs
+) => R {
+  if (maybeGlobalAsyncLocalStorage) {
+    return maybeGlobalAsyncLocalStorage.snapshot()
+  }
+  return function (fn: any, ...args: any[]) {
+    return fn(...args)
+  }
+}

packages/next/src/client/components/static-generation-async-storage.external.ts

@@ -47,6 +47,7 @@ export interface StaticGenerationStore {
   forceStatic?: boolean
   dynamicShouldError?: boolean
   pendingRevalidates?: Record<string, Promise<any>>
+  pendingRevalidateWrites?: Array<Promise<void>> // This is like pendingRevalidates but isn't used for deduping.
 
   dynamicUsageDescription?: string
   dynamicUsageStack?: string

packages/next/src/server/app-render/action-handler.ts

@@ -122,6 +122,7 @@ async function addRevalidationHeader(
       staticGenerationStore.revalidatedTags || []
     ),
     ...Object.values(staticGenerationStore.pendingRevalidates || {}),
+    ...(staticGenerationStore.pendingRevalidateWrites || []),
   ])
 
   // If a tag was revalidated, the client router needs to invalidate all the
@@ -533,6 +534,7 @@ export async function handleAction({
             staticGenerationStore.revalidatedTags || []
           ),
           ...Object.values(staticGenerationStore.pendingRevalidates || {}),
+          ...(staticGenerationStore.pendingRevalidateWrites || []),
         ])
 
         const promise = Promise.reject(error)
@@ -924,6 +926,7 @@ export async function handleAction({
           staticGenerationStore.revalidatedTags || []
         ),
         ...Object.values(staticGenerationStore.pendingRevalidates || {}),
+        ...(staticGenerationStore.pendingRevalidateWrites || []),
       ])
       const promise = Promise.reject(err)
       try {

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

@@ -227,8 +227,8 @@ function parseRequestHeaders(
 
   const flightRouterState = shouldProvideFlightRouterState
     ? parseAndValidateFlightRouterState(
-        headers[NEXT_ROUTER_STATE_TREE_HEADER.toLowerCase()]
-      )
+      headers[NEXT_ROUTER_STATE_TREE_HEADER.toLowerCase()]
+    )
     : undefined
 
   const csp =
@@ -930,7 +930,7 @@ async function renderToHTMLOrFlightImpl(
             })
             .end(
               metrics.clientComponentLoadStart +
-                metrics.clientComponentLoadTimes
+              metrics.clientComponentLoadTimes
             )
         }
       }
@@ -1086,13 +1086,15 @@ async function renderToHTMLOrFlightImpl(
     // If we have pending revalidates, wait until they are all resolved.
     if (
       staticGenerationStore.pendingRevalidates ||
-      staticGenerationStore.revalidatedTags
+      staticGenerationStore.revalidatedTags ||
+      staticGenerationStore.pendingRevalidateWrites
     ) {
       options.waitUntil = Promise.all([
         staticGenerationStore.incrementalCache?.revalidateTag(
           staticGenerationStore.revalidatedTags || []
         ),
         ...Object.values(staticGenerationStore.pendingRevalidates || {}),
+        ...(staticGenerationStore.pendingRevalidateWrites || []),
       ])
     }
 
@@ -1194,13 +1196,15 @@ async function renderToHTMLOrFlightImpl(
     // If we have pending revalidates, wait until they are all resolved.
     if (
       staticGenerationStore.pendingRevalidates ||
-      staticGenerationStore.revalidatedTags
+      staticGenerationStore.revalidatedTags ||
+      staticGenerationStore.pendingRevalidateWrites
     ) {
       options.waitUntil = Promise.all([
         staticGenerationStore.incrementalCache?.revalidateTag(
           staticGenerationStore.revalidatedTags || []
         ),
         ...Object.values(staticGenerationStore.pendingRevalidates || {}),
+        ...(staticGenerationStore.pendingRevalidateWrites || []),
       ])
     }
 
@@ -1849,8 +1853,8 @@ async function prerenderToStream(
           ctx,
           res.statusCode === 404
         )
-        function voidOnError() {}
-        ;(
+        function voidOnError() { }
+        ; (
           prerenderAsyncStorage.run(
             // The store to scope
             prospectiveRenderPrerenderStore,
@@ -1869,7 +1873,7 @@ async function prerenderToStream(
               signal: flightController.signal,
             }
           ) as Promise<ReactServerPrerenderResolveToType>
-        ).catch(() => {})
+        ).catch(() => { })
 
         // When this resolves the cache has no inflight reads and we can ascertain the dynamic outcome
         await cacheSignal.cacheReady()
@@ -2064,7 +2068,7 @@ async function prerenderToStream(
             const resumeStream = await resume(
               <App
                 reactServerStream={foreverStream}
-                preinitScripts={() => {}}
+                preinitScripts={() => { }}
                 clientReferenceManifest={clientReferenceManifest}
                 ServerInsertedHTMLProvider={ServerInsertedHTMLProvider}
                 nonce={ctx.nonce}
@@ -2157,24 +2161,24 @@ async function prerenderToStream(
           ctx,
           res.statusCode === 404
         )
-        // We're not going to use the result of this render because the only time it could be used
-        // is if it completes in a microtask and that's likely very rare for any non-trivial app
-        ;(
-          prerenderAsyncStorage.run(
-            // The store to scope
-            prospectiveRenderPrerenderStore,
-            // The function to run
-            ComponentMod.prerender,
-            // ... the arguments for the function to run
-            firstAttemptRSCPayload,
-            clientReferenceManifest.clientModules,
-            {
-              nonce: ctx.nonce,
-              onError,
-              signal: flightController.signal,
-            }
-          ) as Promise<ReactServerPrerenderResolveToType>
-        ).catch(() => {})
+          // We're not going to use the result of this render because the only time it could be used
+          // is if it completes in a microtask and that's likely very rare for any non-trivial app
+          ; (
+            prerenderAsyncStorage.run(
+              // The store to scope
+              prospectiveRenderPrerenderStore,
+              // The function to run
+              ComponentMod.prerender,
+              // ... the arguments for the function to run
+              firstAttemptRSCPayload,
+              clientReferenceManifest.clientModules,
+              {
+                nonce: ctx.nonce,
+                onError,
+                signal: flightController.signal,
+              }
+            ) as Promise<ReactServerPrerenderResolveToType>
+          ).catch(() => { })
 
         // When this resolves the cache has no inflight reads and we can ascertain the dynamic outcome
         await cacheSignal.cacheReady()
@@ -2513,7 +2517,7 @@ async function prerenderToStream(
           const resumeStream = await resume(
             <App
               reactServerStream={foreverStream}
-              preinitScripts={() => {}}
+              preinitScripts={() => { }}
               clientReferenceManifest={clientReferenceManifest}
               ServerInsertedHTMLProvider={ServerInsertedHTMLProvider}
               nonce={ctx.nonce}