Handle dev fouc for layout styling (#38557)

* Handle dev fouc for layout styling

* refactor

* fix renderOpts.dev

* dedupe css loading

* keep css module id in dev

packages/next/client/app-index.tsx

@@ -75,6 +75,7 @@ const getCacheKey = () => {
 }
 
 const encoder = new TextEncoder()
+const loadedCss: Set<string> = new Set()
 
 let initialServerDataBuffer: string[] | undefined = undefined
 let initialServerDataWriter: ReadableStreamDefaultController | undefined =
@@ -146,76 +147,103 @@ function createResponseCache() {
 }
 const rscCache = createResponseCache()
 
-function useInitialServerResponse(cacheKey: string) {
-  const response = rscCache.get(cacheKey)
-  if (response) return response
-
-  const readable = new ReadableStream({
-    start(controller) {
-      nextServerDataRegisterWriter(controller)
-    },
-  })
-
-  async function loadCss(cssChunkInfoJson: string) {
-    const data = JSON.parse(cssChunkInfoJson)
-    await Promise.all(
-      data.chunks.map((chunkId: string) => {
-        // load css related chunks
-        return (self as any).__next_chunk_load__(chunkId)
-      })
-    )
-    // In development mode, import css in dev when it's wrapped by style loader.
-    // In production mode, css are standalone chunk that doesn't need to be imported.
-    if (data.id) {
-      ;(self as any).__next_require__(data.id)
-    }
+async function loadCss(cssChunkInfoJson: string) {
+  const data = JSON.parse(cssChunkInfoJson)
+  await Promise.all(
+    data.chunks.map((chunkId: string) => {
+      // load css related chunks
+      return (self as any).__next_chunk_load__(chunkId)
+    })
+  )
+  // In development mode, import css in dev when it's wrapped by style loader.
+  // In production mode, css are standalone chunk that doesn't need to be imported.
+  if (data.id) {
+    return (self as any).__next_require__(data.id)
   }
 
+  return Promise.resolve()
+}
+
+function createLoadFlightCssStream(callback?: () => Promise<void>) {
+  const cssLoadingPromises: Promise<any>[] = []
   const loadCssFromStreamData = (data: string) => {
-    if (data.startsWith('CSS:')) {
-      loadCss(data.slice(4).trim())
+    if (data.startsWith('CSS')) {
+      const cssJson = data.slice(4).trim()
+      if (!loadedCss.has(cssJson)) {
+        loadedCss.add(cssJson)
+        cssLoadingPromises.push(loadCss(cssJson))
+      }
     }
   }
 
   // TODO-APP: Refine the buffering code here to make it more correct.
   let buffer = ''
   const loadCssFromFlight = new TransformStream({
     transform(chunk, controller) {
+      const process = (buf: string) => {
+        if (buf) {
+          if (buf.startsWith('CSS:')) {
+            loadCssFromStreamData(buf)
+          } else {
+            controller.enqueue(new TextEncoder().encode(buf))
+          }
+        }
+      }
+
       const data = new TextDecoder().decode(chunk)
       buffer += data
       let index
       while ((index = buffer.indexOf('\n')) !== -1) {
         const line = buffer.slice(0, index + 1)
         buffer = buffer.slice(index + 1)
-        if (line.startsWith('CSS:')) {
-          loadCssFromStreamData(line)
-        } else {
-          controller.enqueue(new TextEncoder().encode(line))
-        }
-      }
-      if (buffer && !buffer.startsWith('CSS:')) {
-        controller.enqueue(new TextEncoder().encode(buffer))
-        buffer = ''
+        process(line)
       }
+      process(buffer)
+      buffer = ''
     },
-    flush() {
-      loadCssFromStreamData(buffer)
+  })
+
+  if (process.env.NODE_ENV === 'development') {
+    Promise.all(cssLoadingPromises).then(() => {
+      // TODO: find better timing for css injection
+      setTimeout(() => {
+        callback?.()
+      })
+    })
+  }
+
+  return loadCssFromFlight
+}
+
+function useInitialServerResponse(cacheKey: string, onFlightCssLoaded: any) {
+  const response = rscCache.get(cacheKey)
+  if (response) return response
+
+  const readable = new ReadableStream({
+    start(controller) {
+      nextServerDataRegisterWriter(controller)
     },
   })
 
   const newResponse = createFromReadableStream(
-    readable.pipeThrough(loadCssFromFlight)
+    readable.pipeThrough(createLoadFlightCssStream(onFlightCssLoaded))
   )
 
   rscCache.set(cacheKey, newResponse)
   return newResponse
 }
 
-function ServerRoot({ cacheKey }: { cacheKey: string }) {
+function ServerRoot({
+  cacheKey,
+  onFlightCssLoaded,
+}: {
+  cacheKey: string
+  onFlightCssLoaded: any
+}) {
   React.useEffect(() => {
     rscCache.delete(cacheKey)
   })
-  const response = useInitialServerResponse(cacheKey)
+  const response = useInitialServerResponse(cacheKey, onFlightCssLoaded)
   const root = response.readRoot()
   return root
 }
@@ -235,16 +263,16 @@ function Root({ children }: React.PropsWithChildren<{}>): React.ReactElement {
   return children as React.ReactElement
 }
 
-function RSCComponent() {
+function RSCComponent(props: any) {
   const cacheKey = getCacheKey()
-  return <ServerRoot cacheKey={cacheKey} />
+  return <ServerRoot {...props} cacheKey={cacheKey} />
 }
 
-export function hydrate() {
+export function hydrate(opts?: { onFlightCssLoaded?: () => Promise<void> }) {
   renderReactElement(appElement!, () => (
     <React.StrictMode>
       <Root>
-        <RSCComponent />
+        <RSCComponent onFlightCssLoaded={opts?.onFlightCssLoaded} />
       </Root>
     </React.StrictMode>
   ))

packages/next/client/app-next-dev.js

@@ -1,4 +1,5 @@
 import { hydrate, version } from './app-index'
+import { displayContent } from './dev/fouc'
 
 // TODO-APP: implement FOUC guard
 
@@ -9,6 +10,6 @@ window.next = {
   appDir: true,
 }
 
-hydrate()
+hydrate({ onFlightCssLoaded: displayContent })
 
 // TODO-APP: build indicator

packages/next/client/components/app-router.client.tsx

@@ -49,26 +49,26 @@ function fetchFlight(url: URL, flightRouterStateData: string): ReadableStream {
   let buffer = ''
   const loadCssFromFlight = new TransformStream({
     transform(chunk, controller) {
+      const process = (buf: string) => {
+        if (buf) {
+          if (buf.startsWith('CSS:')) {
+            loadCssFromStreamData(buf)
+          } else {
+            controller.enqueue(new TextEncoder().encode(buf))
+          }
+        }
+      }
+
       const data = new TextDecoder().decode(chunk)
       buffer += data
       let index
       while ((index = buffer.indexOf('\n')) !== -1) {
         const line = buffer.slice(0, index + 1)
         buffer = buffer.slice(index + 1)
-
-        if (line.startsWith('CSS:')) {
-          loadCssFromStreamData(line)
-        } else {
-          controller.enqueue(new TextEncoder().encode(line))
-        }
-      }
-      if (buffer && !buffer.startsWith('CSS:')) {
-        controller.enqueue(new TextEncoder().encode(buffer))
-        buffer = ''
+        process(line)
       }
-    },
-    flush() {
-      loadCssFromStreamData(buffer)
+      process(buffer)
+      buffer = ''
     },
   })
 

packages/next/server/app-render.tsx

@@ -148,13 +148,11 @@ function useFlightResponse(
           writer.close()
         } else {
           const responsePartial = decodeText(value)
-          writer.write(
-            encodeText(
-              `<script>(self.__next_s=self.__next_s||[]).push(${htmlEscapeJsonString(
-                JSON.stringify([1, id, responsePartial])
-              )})</script>`
-            )
-          )
+          const scripts = `<script>(self.__next_s=self.__next_s||[]).push(${htmlEscapeJsonString(
+            JSON.stringify([1, id, responsePartial])
+          )})</script>`
+
+          writer.write(encodeText(scripts))
           process()
         }
       })
@@ -178,7 +176,8 @@ function createServerComponentRenderer(
     transformStream: TransformStream<Uint8Array, Uint8Array>
     serverComponentManifest: NonNullable<RenderOpts['serverComponentManifest']>
     serverContexts: Array<[ServerContextName: string, JSONValue: any]>
-  }
+  },
+  dev: boolean
 ) {
   // We need to expose the `__webpack_require__` API globally for
   // react-server-dom-webpack. This is a hack until we find a better way.
@@ -192,7 +191,11 @@ function createServerComponentRenderer(
     globalThis.__next_chunk_load__ = () => Promise.resolve()
   }
 
-  const cssFlightData = getCssFlightData(ComponentMod, serverComponentManifest)
+  const cssFlightData = getCssFlightData(
+    ComponentMod,
+    serverComponentManifest,
+    dev
+  )
 
   let RSCStream: ReadableStream<Uint8Array>
   const createRSCStream = () => {
@@ -326,8 +329,11 @@ function getSegmentParam(segment: string): {
 
 function getCSSInlinedLinkTags(
   ComponentMod: any,
-  serverComponentManifest: any
+  serverComponentManifest: any,
+  dev: boolean
 ) {
+  if (dev) return []
+
   const importedServerCSSFiles: string[] =
     ComponentMod.__client__?.__next_rsc_css__ || []
 
@@ -344,25 +350,31 @@ function getCSSInlinedLinkTags(
   )
 }
 
-function getCssFlightData(ComponentMod: any, serverComponentManifest: any) {
+function getCssFlightData(
+  ComponentMod: any,
+  serverComponentManifest: any,
+  dev: boolean
+) {
   const importedServerCSSFiles: string[] =
     ComponentMod.__client__?.__next_rsc_css__ || []
 
   const cssFiles = importedServerCSSFiles.map(
     (css) => serverComponentManifest[css].default
   )
-  if (process.env.NODE_ENV === 'development') {
+
+  if (dev) {
+    // Keep `id` in dev mode css flight to require the css module
     return cssFiles.map((css) => `CSS:${JSON.stringify(css)}`).join('\n') + '\n'
   }
 
   // Multiple css chunks could be merged into one by mini-css-extract-plugin,
   // we use a set here to dedupe the css chunks in production.
-  const cssSet = cssFiles.reduce((res, css) => {
+  const cssSet: Set<string> = cssFiles.reduce((res, css) => {
     res.add(...css.chunks)
     return res
   }, new Set())
 
-  return `CSS:${JSON.stringify({ chunks: [...cssSet] })}\n`
+  return cssSet.size ? `CSS:${JSON.stringify({ chunks: [...cssSet] })}\n` : ''
 }
 
 export async function renderToHTML(
@@ -383,6 +395,7 @@ export async function renderToHTML(
     runtime,
     ComponentMod,
   } = renderOpts
+  const dev = !!renderOpts.dev
 
   const isFlight = query.__flight__ !== undefined
 
@@ -781,7 +794,8 @@ export async function renderToHTML(
 
     const cssFlightData = getCssFlightData(
       ComponentMod,
-      serverComponentManifest
+      serverComponentManifest,
+      dev
     )
     const flightData: FlightData = [
       // TODO-APP: change walk to output without ''
@@ -807,9 +821,10 @@ export async function renderToHTML(
   // /blog/[slug] /blog/hello-world -> ['children', 'blog', 'children', ['slug', 'hello-world']]
   const initialTree = createFlightRouterStateFromLoaderTree(tree)
 
-  const initialStylesheets = getCSSInlinedLinkTags(
+  const initialStylesheets: string[] = getCSSInlinedLinkTags(
     ComponentMod,
-    serverComponentManifest
+    serverComponentManifest,
+    dev
   )
 
   const { Component: ComponentTree } = createComponentTree({
@@ -866,7 +881,8 @@ export async function renderToHTML(
       transformStream: serverComponentsInlinedTransformStream,
       serverComponentManifest,
       serverContexts,
-    }
+    },
+    dev
   )
 
   const jsxStyleRegistry = createStyleRegistry()
@@ -916,6 +932,7 @@ export async function renderToHTML(
     }
 
     return await continueFromInitialStream(renderStream, {
+      dev,
       dataStream: serverComponentsInlinedTransformStream?.readable,
       generateStaticHTML: generateStaticHTML || !hasConcurrentFeatures,
       flushEffectHandler,