From f2fee3d6ecd9c0d370ab5eab390f95537398a5a4 Mon Sep 17 00:00:00 2001
From: Jiachi Liu <inbox@huozhi.im>
Date: Fri, 18 Mar 2022 02:55:59 +0100
Subject: [PATCH] Fix: change the timing for clearing the inlined data buffer
 (#35413)

Fix error of `Unexpected server data: missing bootstrap script` when the initial rendering is streaming rsc.

Introduced in #35344, the buffer might still keep flushing after on mount.


Co-authored-by: Shu Ding <3676859+shuding@users.noreply.github.com>
---
 packages/next/client/index.tsx | 46 +++++++++++++++++++++++++---------
 1 file changed, 34 insertions(+), 12 deletions(-)

diff --git a/packages/next/client/index.tsx b/packages/next/client/index.tsx
index c00d17f7276c..cd2f3fc315b2 100644
--- a/packages/next/client/index.tsx
+++ b/packages/next/client/index.tsx
@@ -678,9 +678,13 @@ if (process.env.__NEXT_RSC) {
   } = require('next/dist/compiled/react-server-dom-webpack')
 
   const encoder = new TextEncoder()
+
   let initialServerDataBuffer: string[] | undefined = undefined
   let initialServerDataWriter: WritableStreamDefaultWriter | undefined =
     undefined
+  let initialServerDataLoaded = false
+  let initialServerDataFlushed = false
+
   function nextServerDataCallback(seg: [number, string, string]) {
     if (seg[0] === 0) {
       initialServerDataBuffer = []
@@ -695,24 +699,45 @@ if (process.env.__NEXT_RSC) {
       }
     }
   }
+
+  // There might be race conditions between `nextServerDataRegisterWriter` and
+  // `DOMContentLoaded`. The former will be called when React starts to hydrate
+  // the root, the latter will be called when the DOM is fully loaded.
+  // For streaming, the former is called first due to partial hydration.
+  // For non-streaming, the latter can be called first.
+  // Hence, we use two variables `initialServerDataLoaded` and
+  // `initialServerDataFlushed` to make sure the writer will be closed and
+  // `initialServerDataBuffer` will be cleared in the right time.
   function nextServerDataRegisterWriter(writer: WritableStreamDefaultWriter) {
     if (initialServerDataBuffer) {
       initialServerDataBuffer.forEach((val) => {
         writer.write(encoder.encode(val))
       })
+      if (initialServerDataLoaded && !initialServerDataFlushed) {
+        writer.close()
+        initialServerDataFlushed = true
+        initialServerDataBuffer = undefined
+      }
     }
+
     initialServerDataWriter = writer
   }
+
   // When `DOMContentLoaded`, we can close all pending writers to finish hydration.
-  document.addEventListener(
-    'DOMContentLoaded',
-    function () {
-      if (initialServerDataWriter && !initialServerDataWriter.closed) {
-        initialServerDataWriter.close()
-      }
-    },
-    false
-  )
+  const DOMContentLoaded = function () {
+    if (initialServerDataWriter && !initialServerDataFlushed) {
+      initialServerDataWriter.close()
+      initialServerDataFlushed = true
+      initialServerDataBuffer = undefined
+    }
+    initialServerDataLoaded = true
+  }
+  // It's possible that the DOM is already loaded.
+  if (document.readyState === 'loading') {
+    document.addEventListener('DOMContentLoaded', DOMContentLoaded, false)
+  } else {
+    DOMContentLoaded()
+  }
 
   const nextServerDataLoadingGlobal = ((self as any).__next_s =
     (self as any).__next_s || [])
@@ -771,9 +796,6 @@ if (process.env.__NEXT_RSC) {
     React.useEffect(() => {
       rscCache.delete(cacheKey)
     })
-    React.useEffect(() => {
-      initialServerDataBuffer = undefined
-    }, [])
     const response = useServerResponse(cacheKey, serialized)
     const root = response.readRoot()
     return root