Handle dev fouc for layout styling

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,34 +147,29 @@ 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)) cssLoadingPromises.push(loadCss(cssJson))
     }
   }
 
@@ -193,29 +189,55 @@ function useInitialServerResponse(cacheKey: string) {
           controller.enqueue(new TextEncoder().encode(line))
         }
       }
+
       if (buffer && !buffer.startsWith('CSS:')) {
         controller.enqueue(new TextEncoder().encode(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 +257,18 @@ 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 async 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/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.
@@ -326,8 +325,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__ || []
 
@@ -351,18 +353,15 @@ function getCssFlightData(ComponentMod: any, serverComponentManifest: any) {
   const cssFiles = importedServerCSSFiles.map(
     (css) => serverComponentManifest[css].default
   )
-  if (process.env.NODE_ENV === 'development') {
-    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 +382,7 @@ export async function renderToHTML(
     runtime,
     ComponentMod,
   } = renderOpts
+  const dev = !!renderOpts.dev
 
   const isFlight = query.__flight__ !== undefined
 
@@ -807,9 +807,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 +867,8 @@ export async function renderToHTML(
       transformStream: serverComponentsInlinedTransformStream,
       serverComponentManifest,
       serverContexts,
-    }
+    },
+    dev
   )
 
   const jsxStyleRegistry = createStyleRegistry()
@@ -916,6 +918,7 @@ export async function renderToHTML(
     }
 
     return await continueFromInitialStream(renderStream, {
+      dev,
       dataStream: serverComponentsInlinedTransformStream?.readable,
       generateStaticHTML: generateStaticHTML || !hasConcurrentFeatures,
       flushEffectHandler,

packages/next/server/node-web-streams-helper.ts

@@ -136,6 +136,33 @@ export function createFlushEffectStream(
   })
 }
 
+export function createDevScriptTransformStream(): TransformStream<
+  Uint8Array,
+  Uint8Array
+> {
+  let injected = false
+  const foucTags = `<style data-next-hide-fouc>body{display:none}</style>
+  <noscript data-next-hide-fouc>
+    <style>body{display:block}</style>
+  </noscript>`
+  return new TransformStream({
+    transform(chunk, controller) {
+      const content = decodeText(chunk)
+      let index
+      if (!injected && (index = content.indexOf('</head')) !== -1) {
+        injected = true
+        // head content + fouc tags + </head
+        const injectedContent =
+          content.slice(0, index) + foucTags + content.slice(index)
+
+        controller.enqueue(encodeText(injectedContent))
+      } else {
+        controller.enqueue(chunk)
+      }
+    },
+  })
+}
+
 export function renderToInitialStream({
   ReactDOMServer,
   element,
@@ -151,11 +178,13 @@ export function renderToInitialStream({
 export async function continueFromInitialStream(
   renderStream: ReactReadableStream,
   {
+    dev,
     suffix,
     dataStream,
     generateStaticHTML,
     flushEffectHandler,
   }: {
+    dev?: boolean
     suffix?: string
     dataStream?: ReadableStream<Uint8Array>
     generateStaticHTML: boolean
@@ -172,9 +201,10 @@ export async function continueFromInitialStream(
   const transforms: Array<TransformStream<Uint8Array, Uint8Array>> = [
     createBufferedTransformStream(),
     flushEffectHandler ? createFlushEffectStream(flushEffectHandler) : null,
-    suffixUnclosed != null ? createDeferredPrefixStream(suffixUnclosed) : null,
+    suffixUnclosed != null ? createDeferredSuffixStream(suffixUnclosed) : null,
     dataStream ? createInlineDataStream(dataStream) : null,
     suffixUnclosed != null ? createSuffixStream(closeTag) : null,
+    dev ? createDevScriptTransformStream() : null,
   ].filter(nonNullable)
 
   return transforms.reduce(
@@ -198,44 +228,47 @@ export function createSuffixStream(
 export function createPrefixStream(
   prefix: string
 ): TransformStream<Uint8Array, Uint8Array> {
-  let prefixEnqueued = false
+  let prefixFlushed = false
   return new TransformStream({
     transform(chunk, controller) {
-      if (!prefixEnqueued) {
+      if (!prefixFlushed && prefix) {
+        prefixFlushed = true
         controller.enqueue(encodeText(prefix))
-        prefixEnqueued = true
       }
       controller.enqueue(chunk)
     },
   })
 }
 
-export function createDeferredPrefixStream(
-  prefix: string
+// Suffix after main body content - scripts before </body>,
+// but wait for the major chunks to be enqueued.
+export function createDeferredSuffixStream(
+  suffix: string
 ): TransformStream<Uint8Array, Uint8Array> {
-  let prefixFlushed = false
-  let prefixPrefixFlushFinished: Promise<void> | null = null
+  let suffixFlushed = false
+  let suffixFlushTask: Promise<void> | null = null
+
   return new TransformStream({
     transform(chunk, controller) {
       controller.enqueue(chunk)
-      if (!prefixFlushed && prefix) {
-        prefixFlushed = true
-        prefixPrefixFlushFinished = new Promise((res) => {
+      if (!suffixFlushed && suffix) {
+        suffixFlushed = true
+        suffixFlushTask = new Promise((res) => {
           // NOTE: streaming flush
-          // Enqueue prefix part before the major chunks are enqueued so that
-          // prefix won't be flushed too early to interrupt the data stream
+          // Enqueue suffix part before the major chunks are enqueued so that
+          // suffix won't be flushed too early to interrupt the data stream
           setTimeout(() => {
-            controller.enqueue(encodeText(prefix))
+            controller.enqueue(encodeText(suffix))
             res()
           })
         })
       }
     },
     flush(controller) {
-      if (prefixPrefixFlushFinished) return prefixPrefixFlushFinished
-      if (!prefixFlushed && prefix) {
-        prefixFlushed = true
-        controller.enqueue(encodeText(prefix))
+      if (suffixFlushTask) return suffixFlushTask
+      if (!suffixFlushed && suffix) {
+        suffixFlushed = true
+        controller.enqueue(encodeText(suffix))
       }
     },
   })