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

@@ -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
     }

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)
+  }
 }

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