From aef60dc581f564b64cfb7667cdca0bd3e3b6e366 Mon Sep 17 00:00:00 2001
From: Jiachi Liu <inbox@huozhi.im>
Date: Fri, 12 Aug 2022 21:06:08 +0200
Subject: [PATCH] Support multiple flush effects (#39559)

Bring more flexibility for consume flush effects in separate places. Then you can move it into different client components roots
---
 packages/next/server/app-render.tsx              | 16 ++++++----------
 packages/next/shared/lib/flush-effects.tsx       |  9 +++++----
 .../rsc-basic/app/root-style-registry.client.js  | 12 ++++++------
 3 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/packages/next/server/app-render.tsx b/packages/next/server/app-render.tsx
index b3906ed2b011458..af038ae9d27ed54 100644
--- a/packages/next/server/app-render.tsx
+++ b/packages/next/server/app-render.tsx
@@ -1018,23 +1018,19 @@ export async function renderToHTMLOrFlight(
     }
   )
 
-  let flushEffectsHandler: (() => React.ReactNode) | null = null
+  const flushEffectsCallbacks: Set<() => React.ReactNode> = new Set()
   function FlushEffects({ children }: { children: JSX.Element }) {
     // Reset flushEffectsHandler on each render
-    flushEffectsHandler = null
-    const setFlushEffectsHandler = React.useCallback(
+    flushEffectsCallbacks.clear()
+    const addFlushEffects = React.useCallback(
       (handler: () => React.ReactNode) => {
-        if (flushEffectsHandler)
-          throw new Error(
-            'The `useFlushEffects` hook cannot be used more than once.'
-          )
-        flushEffectsHandler = handler
+        flushEffectsCallbacks.add(handler)
       },
       []
     )
 
     return (
-      <FlushEffectsContext.Provider value={setFlushEffectsHandler}>
+      <FlushEffectsContext.Provider value={addFlushEffects}>
         {children}
       </FlushEffectsContext.Provider>
     )
@@ -1063,7 +1059,7 @@ export async function renderToHTMLOrFlight(
 
     const flushEffectHandler = (): string => {
       const flushed = ReactDOMServer.renderToString(
-        <>{flushEffectsHandler && flushEffectsHandler()}</>
+        <>{Array.from(flushEffectsCallbacks).map((callback) => callback())}</>
       )
       return flushed
     }
diff --git a/packages/next/shared/lib/flush-effects.tsx b/packages/next/shared/lib/flush-effects.tsx
index 72c093c6cc2cd39..4aae4f1d73ff458 100644
--- a/packages/next/shared/lib/flush-effects.tsx
+++ b/packages/next/shared/lib/flush-effects.tsx
@@ -6,9 +6,10 @@ export const FlushEffectsContext = createContext<FlushEffectsHook | null>(
   null as any
 )
 
-export function useFlushEffects(callbacks: () => React.ReactNode): void {
-  const flushEffectsImpl = useContext(FlushEffectsContext)
+export function useFlushEffects(callback: () => React.ReactNode): void {
+  const addFlushEffects = useContext(FlushEffectsContext)
   // Should have no effects on client where there's no flush effects provider
-  if (!flushEffectsImpl) return
-  return flushEffectsImpl(callbacks)
+  if (addFlushEffects) {
+    addFlushEffects(callback)
+  }
 }
diff --git a/test/e2e/app-dir/rsc-basic/app/root-style-registry.client.js b/test/e2e/app-dir/rsc-basic/app/root-style-registry.client.js
index c642e14fb1a76b4..bef5c5ce392dd7b 100644
--- a/test/e2e/app-dir/rsc-basic/app/root-style-registry.client.js
+++ b/test/e2e/app-dir/rsc-basic/app/root-style-registry.client.js
@@ -19,13 +19,13 @@ export default function RootStyleRegistry({ children }) {
     return <>{styles}</>
   }
 
+  // Allow multiple useFlushEffects
   useFlushEffects(() => {
-    return (
-      <>
-        {styledJsxFlushEffect()}
-        {styledComponentsFlushEffect()}
-      </>
-    )
+    return <>{styledJsxFlushEffect()}</>
+  })
+
+  useFlushEffects(() => {
+    return <>{styledComponentsFlushEffect()}</>
   })
 
   // Only include style registry on server side for SSR