Migrate head side effects to hooks (#37526)

* rewrite head side effects component in hooks * remove mapping from element to children in head manager since they're already the children of `<Head>` When move `SideEffect` to hooks, the effects scheduling is earlier than life cycle. We're leverage layout effects and effects at the same time, always cache the latest head updating function in head manager in layout effects, and flush them in the effects. This could help get rid of the promises delaying approach in head manager. Co-authored-by: Shu Ding <3676859+shuding@users.noreply.github.com>
vercel · Jun 8, 2022 · 16fd15b · 16fd15b
1 parent 88d0440
commit 16fd15b
Show file tree

Hide file tree

Showing 7 changed files with 134 additions and 98 deletions.
diff --git a/packages/next/client/head-manager.ts b/packages/next/client/head-manager.ts
@@ -118,57 +118,48 @@ export default function initHeadManager(): {
   mountedInstances: Set<unknown>
   updateHead: (head: JSX.Element[]) => void
 } {
-  let updatePromise: Promise<void> | null = null
-
   return {
     mountedInstances: new Set(),
     updateHead: (head: JSX.Element[]) => {
-      const promise = (updatePromise = Promise.resolve().then(() => {
-        if (promise !== updatePromise) return
-
-        updatePromise = null
-        const tags: Record<string, JSX.Element[]> = {}
-
-        head.forEach((h) => {
+      const tags: Record<string, JSX.Element[]> = {}
+
+      head.forEach((h) => {
+        if (
+          // If the font tag is loaded only on client navigation
+          // it won't be inlined. In this case revert to the original behavior
+          h.type === 'link' &&
+          h.props['data-optimized-fonts']
+        ) {
           if (
-            // If the font tag is loaded only on client navigation
-            // it won't be inlined. In this case revert to the original behavior
-            h.type === 'link' &&
-            h.props['data-optimized-fonts']
+            document.querySelector(`style[data-href="${h.props['data-href']}"]`)
           ) {
-            if (
-              document.querySelector(
-                `style[data-href="${h.props['data-href']}"]`
-              )
-            ) {
-              return
-            } else {
-              h.props.href = h.props['data-href']
-              h.props['data-href'] = undefined
-            }
+            return
+          } else {
+            h.props.href = h.props['data-href']
+            h.props['data-href'] = undefined
           }
-
-          const components = tags[h.type] || []
-          components.push(h)
-          tags[h.type] = components
-        })
-
-        const titleComponent = tags.title ? tags.title[0] : null
-        let title = ''
-        if (titleComponent) {
-          const { children } = titleComponent.props
-          title =
-            typeof children === 'string'
-              ? children
-              : Array.isArray(children)
-              ? children.join('')
-              : ''
         }
-        if (title !== document.title) document.title = title
-        ;['meta', 'base', 'link', 'style', 'script'].forEach((type) => {
-          updateElements(type, tags[type] || [])
-        })
-      }))
+
+        const components = tags[h.type] || []
+        components.push(h)
+        tags[h.type] = components
+      })
+
+      const titleComponent = tags.title ? tags.title[0] : null
+      let title = ''
+      if (titleComponent) {
+        const { children } = titleComponent.props
+        title =
+          typeof children === 'string'
+            ? children
+            : Array.isArray(children)
+            ? children.join('')
+            : ''
+      }
+      if (title !== document.title) document.title = title
+      ;['meta', 'base', 'link', 'style', 'script'].forEach((type) => {
+        updateElements(type, tags[type] || [])
+      })
     },
   }
 }
diff --git a/packages/next/shared/lib/head.tsx b/packages/next/shared/lib/head.tsx
@@ -116,22 +116,13 @@ function unique() {
 
 /**
  *
- * @param headElements List of multiple <Head> instances
+ * @param headChildrenElements List of children of <Head>
  */
 function reduceComponents(
-  headElements: Array<React.ReactElement<any>>,
+  headChildrenElements: Array<React.ReactElement<any>>,
   props: WithInAmpMode
 ) {
-  return headElements
-    .reduce(
-      (list: React.ReactChild[], headElement: React.ReactElement<any>) => {
-        const headElementChildren = React.Children.toArray(
-          headElement.props.children
-        )
-        return list.concat(headElementChildren)
-      },
-      []
-    )
+  return headChildrenElements
     .reduce(onlyReactElement, [])
     .reverse()
     .concat(defaultHead(props.inAmpMode).reverse())

diff --git a/packages/next/shared/lib/side-effect.tsx b/packages/next/shared/lib/side-effect.tsx
@@ -1,6 +1,4 @@
-import React, { Component } from 'react'
-
-const isServer = typeof window === 'undefined'
+import React, { Children, useEffect, useLayoutEffect } from 'react'
 
 type State = JSX.Element[] | undefined
 
@@ -12,49 +10,65 @@ type SideEffectProps = {
   handleStateChange?: (state: State) => void
   headManager: any
   inAmpMode?: boolean
+  children: React.ReactNode
 }
 
-export default class extends Component<SideEffectProps> {
-  private _hasHeadManager: boolean
-
-  emitChange = (): void => {
-    if (this._hasHeadManager) {
-      this.props.headManager.updateHead(
-        this.props.reduceComponentsToState(
-          [...this.props.headManager.mountedInstances],
-          this.props
-        )
-      )
+const isServer = typeof window === 'undefined'
+const useClientOnlyLayoutEffect = isServer ? () => {} : useLayoutEffect
+const useClientOnlyEffect = isServer ? () => {} : useEffect
+
+export default function SideEffect(props: SideEffectProps) {
+  const { headManager, reduceComponentsToState } = props
+
+  function emitChange() {
+    if (headManager && headManager.mountedInstances) {
+      const headElements = Children.toArray(
+        headManager.mountedInstances
+      ).filter(Boolean) as React.ReactElement[]
+      headManager.updateHead(reduceComponentsToState(headElements, props))
     }
   }
 
-  constructor(props: any) {
-    super(props)
-    this._hasHeadManager =
-      this.props.headManager && this.props.headManager.mountedInstances
+  if (isServer) {
+    headManager?.mountedInstances?.add(props.children)
+    emitChange()
+  }
 
-    if (isServer && this._hasHeadManager) {
-      this.props.headManager.mountedInstances.add(this)
-      this.emitChange()
+  useClientOnlyLayoutEffect(() => {
+    headManager?.mountedInstances?.add(props.children)
+    return () => {
+      headManager?.mountedInstances?.delete(props.children)
     }
-  }
-  componentDidMount() {
-    if (this._hasHeadManager) {
-      this.props.headManager.mountedInstances.add(this)
+  })
+
+  // We need to call `updateHead` method whenever the `SideEffect` is trigger in all
+  // life-cycles: mount, update, unmount. However, if there are multiple `SideEffect`s
+  // being rendered, we only trigger the method from the last one.
+  // This is ensured by keeping the last unflushed `updateHead` in the `_pendingUpdate`
+  // singleton in the layout effect pass, and actually trigger it in the effect pass.
+  useClientOnlyLayoutEffect(() => {
+    if (headManager) {
+      headManager._pendingUpdate = emitChange
     }
-    this.emitChange()
-  }
-  componentDidUpdate() {
-    this.emitChange()
-  }
-  componentWillUnmount() {
-    if (this._hasHeadManager) {
-      this.props.headManager.mountedInstances.delete(this)
+    return () => {
+      if (headManager) {
+        headManager._pendingUpdate = emitChange
+      }
     }
-    this.emitChange()
-  }
+  })
 
-  render() {
-    return null
-  }
+  useClientOnlyEffect(() => {
+    if (headManager && headManager._pendingUpdate) {
+      headManager._pendingUpdate()
+      headManager._pendingUpdate = null
+    }
+    return () => {
+      if (headManager && headManager._pendingUpdate) {
+        headManager._pendingUpdate()
+        headManager._pendingUpdate = null
+      }
+    }
+  })
+
+  return null
 }
diff --git a/test/integration/client-navigation-a11y/next.config.js b/test/integration/client-navigation-a11y/next.config.js
diff --git a/test/integration/client-navigation-a11y/test/index.test.js b/test/integration/client-navigation-a11y/test/index.test.js
@@ -11,7 +11,6 @@ import { join } from 'path'
 
 const context = {}
 const appDir = join(__dirname, '../')
-const nodeArgs = ['-r', join(appDir, '../../lib/react-17-require-hook.js')]
 
 const navigateTo = async (browser, selector) =>
   await browser
@@ -32,7 +31,6 @@ describe('Client Navigation accessibility', () => {
     context.appPort = await findPort()
     context.server = await launchApp(appDir, context.appPort, {
       env: { __NEXT_TEST_WITH_DEVTOOL: 1 },
-      nodeArgs,
     })
 
     const prerender = [

diff --git a/test/integration/client-navigation/pages/head-dynamic.js b/test/integration/client-navigation/pages/head-dynamic.js
@@ -0,0 +1,28 @@
+import Head from 'next/head'
+import React from 'react'
+
+function Foo() {
+  const [displayed, toggle] = React.useState(true)
+
+  return (
+    <>
+      {displayed ? (
+        <Head>
+          <title>B</title>
+        </Head>
+      ) : null}
+      <button onClick={() => toggle(!displayed)}>toggle</button>
+    </>
+  )
+}
+
+export default () => {
+  return (
+    <>
+      <Head>
+        <title>A</title>
+      </Head>
+      <Foo />
+    </>
+  )
+}
diff --git a/test/integration/client-navigation/test/index.test.js b/test/integration/client-navigation/test/index.test.js
@@ -1621,6 +1621,22 @@ describe('Client Navigation', () => {
         }
       }
     })
+
+    it('should update head when unmounting component', async () => {
+      let browser
+      try {
+        browser = await webdriver(context.appPort, '/head-dynamic')
+        expect(await browser.eval('document.title')).toBe('B')
+        await browser.elementByCss('button').click()
+        expect(await browser.eval('document.title')).toBe('A')
+        await browser.elementByCss('button').click()
+        expect(await browser.eval('document.title')).toBe('B')
+      } finally {
+        if (browser) {
+          await browser.close()
+        }
+      }
+    })
   })
 
   describe('foreign history manipulation', () => {