perf(ThemeProvider): Reduce unnecessary renders and effect cascades

.changeset/improve-theme-provider-perf.md

@@ -0,0 +1,11 @@
+---
+"@primer/react": patch
+"@primer/styled-react": patch
+---
+
+perf(ThemeProvider): Reduce unnecessary renders and effect cascades
+
+- Replace `useState` + `useEffect` SSR hydration handoff with `useSyncExternalStore` — eliminates post-hydration re-render
+- Replace `useState` + `useEffect` in `useSystemColorMode` with `useSyncExternalStore` — eliminates effect gap and stale-then-update flicker
+- Cache `getServerHandoff` DOM read + JSON.parse per ID (runs once, not on every call)
+- Memoize context value object to prevent unnecessary re-renders of all consumers

packages/react/src/ThemeProvider.tsx

@@ -1,5 +1,4 @@
 import React from 'react'
-import ReactDOM from 'react-dom'
 import defaultTheme from './theme'
 import deepmerge from 'deepmerge'
 import {useId} from './hooks'
@@ -39,16 +38,32 @@ const ThemeContext = React.createContext<{
 })
 
 // inspired from __NEXT_DATA__, we use application/json to avoid CSRF policy with inline scripts
+const serverHandoffCache = new Map<string, Record<string, unknown>>()
+const emptyHandoff: Record<string, unknown> = {}
 const getServerHandoff = (id: string) => {
+  if (typeof document === 'undefined') return emptyHandoff
+
+  const cached = serverHandoffCache.get(id)
+  if (cached !== undefined) return cached
+
   try {
     const serverData = document.getElementById(`__PRIMER_DATA_${id}__`)?.textContent
-    if (serverData) return JSON.parse(serverData)
+    if (serverData) {
+      const parsed = JSON.parse(serverData)
+      serverHandoffCache.set(id, parsed)
+      return parsed
+    }
   } catch (_error) {
     // if document/element does not exist or JSON is invalid, suppress error
   }
-  return {}
+
+  const empty = {}
+  serverHandoffCache.set(id, empty)
+  return empty
 }
 
+const emptySubscribe = () => () => {}
+
 export const ThemeProvider: React.FC<React.PropsWithChildren<ThemeProviderProps>> = ({children, ...props}) => {
   // Get fallback values from parent ThemeProvider (if exists)
   const {
@@ -62,63 +77,54 @@ export const ThemeProvider: React.FC<React.PropsWithChildren<ThemeProviderProps>
   const theme = fallbackTheme ?? defaultTheme
 
   const uniqueDataId = useId()
-  const {resolvedServerColorMode} = getServerHandoff(uniqueDataId)
-  const resolvedColorModePassthrough = React.useRef(resolvedServerColorMode)
 
   const [colorMode, setColorMode] = useSyncedState(props.colorMode ?? fallbackColorMode ?? defaultColorMode)
   const [dayScheme, setDayScheme] = useSyncedState(props.dayScheme ?? fallbackDayScheme ?? defaultDayScheme)
   const [nightScheme, setNightScheme] = useSyncedState(props.nightScheme ?? fallbackNightScheme ?? defaultNightScheme)
   const systemColorMode = useSystemColorMode()
-  // eslint-disable-next-line react-hooks/refs
-  const resolvedColorMode = resolvedColorModePassthrough.current || resolveColorMode(colorMode, systemColorMode)
+  const clientColorMode = resolveColorMode(colorMode, systemColorMode)
+  // During SSR/hydration, use the server-rendered color mode from the handoff script tag
+  // to avoid mismatches. After hydration, resolve from client state.
+  const resolvedColorMode = React.useSyncExternalStore(
+    emptySubscribe,
+    () => clientColorMode,
+    () => getServerHandoff(uniqueDataId).resolvedServerColorMode ?? clientColorMode,
+  )
   const colorScheme = chooseColorScheme(resolvedColorMode, dayScheme, nightScheme)
   const {resolvedTheme, resolvedColorScheme} = React.useMemo(
     () => applyColorScheme(theme, colorScheme),
     [theme, colorScheme],
   )
 
-  // this effect will only run on client
-  React.useEffect(
-    function updateColorModeAfterServerPassthrough() {
-      const resolvedColorModeOnClient = resolveColorMode(colorMode, systemColorMode)
-
-      if (resolvedColorModePassthrough.current) {
-        // if the resolved color mode passed on from the server is not the resolved color mode on client, change it!
-        if (resolvedColorModePassthrough.current !== resolvedColorModeOnClient) {
-          window.setTimeout(() => {
-            // use ReactDOM.flushSync to prevent automatic batching of state updates since React 18
-            // ref: https://github.com/reactwg/react-18/discussions/21
-            ReactDOM.flushSync(() => {
-              // override colorMode to whatever is resolved on the client to get a re-render
-              setColorMode(resolvedColorModeOnClient)
-            })
-
-            // immediately after that, set the colorMode to what the user passed to respond to system color mode changes
-            setColorMode(colorMode)
-          })
-        }
-
-        resolvedColorModePassthrough.current = null
-      }
-    },
-    [colorMode, systemColorMode, setColorMode],
+  const contextValue = React.useMemo(
+    () => ({
+      theme: resolvedTheme,
+      colorScheme,
+      colorMode,
+      resolvedColorMode,
+      resolvedColorScheme,
+      dayScheme,
+      nightScheme,
+      setColorMode,
+      setDayScheme,
+      setNightScheme,
+    }),
+    [
+      resolvedTheme,
+      colorScheme,
+      colorMode,
+      resolvedColorMode,
+      resolvedColorScheme,
+      dayScheme,
+      nightScheme,
+      setColorMode,
+      setDayScheme,
+      setNightScheme,
+    ],
   )
 
   return (
-    <ThemeContext.Provider
-      value={{
-        theme: resolvedTheme,
-        colorScheme,
-        colorMode,
-        resolvedColorMode,
-        resolvedColorScheme,
-        dayScheme,
-        nightScheme,
-        setColorMode,
-        setDayScheme,
-        setNightScheme,
-      }}
-    >
+    <ThemeContext.Provider value={contextValue}>
       <div
         data-color-mode={colorMode === 'auto' ? 'auto' : colorScheme.includes('dark') ? 'dark' : 'light'}
         data-light-theme={dayScheme}
@@ -146,54 +152,22 @@ export function useColorSchemeVar(values: Partial<Record<string, string>>, fallb
   return values[colorScheme] ?? fallback
 }
 
-function useSystemColorMode() {
-  const [systemColorMode, setSystemColorMode] = React.useState(getSystemColorMode)
-
-  React.useEffect(() => {
-    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
-    const media = window?.matchMedia?.('(prefers-color-scheme: dark)')
-
-    function matchesMediaToColorMode(matches: boolean) {
-      return matches ? 'night' : 'day'
-    }
-
-    function handleChange(event: MediaQueryListEvent) {
-      const isNight = event.matches
-      setSystemColorMode(matchesMediaToColorMode(isNight))
-    }
-
-    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
-    if (media) {
-      // just in case the preference changed before the event listener was attached
-      const isNight = media.matches
-      setSystemColorMode(matchesMediaToColorMode(isNight))
-      // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
-      if (media.addEventListener !== undefined) {
-        media.addEventListener('change', handleChange)
-        return function cleanup() {
-          media.removeEventListener('change', handleChange)
-        }
-      }
-      // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
-      else if (media.addListener !== undefined) {
-        media.addListener(handleChange)
-        return function cleanup() {
-          media.removeListener(handleChange)
-        }
-      }
-    }
-  }, [])
+function subscribeToSystemColorMode(callback: () => void) {
+  // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+  const media = window?.matchMedia?.('(prefers-color-scheme: dark)')
+  // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+  media?.addEventListener('change', callback)
+  // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+  return () => media?.removeEventListener('change', callback)
+}
 
-  return systemColorMode
+function useSystemColorMode() {
+  return React.useSyncExternalStore<ColorMode>(subscribeToSystemColorMode, getSystemColorMode, () => 'day')
 }
 
 function getSystemColorMode(): ColorMode {
   // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
-  if (typeof window !== 'undefined' && window.matchMedia?.('(prefers-color-scheme: dark)')?.matches) {
-    return 'night'
-  }
-
-  return 'day'
+  return window?.matchMedia?.('(prefers-color-scheme: dark)')?.matches ? 'night' : 'day'
 }
 
 function resolveColorMode(colorMode: ColorModeWithAuto, systemColorMode: ColorMode) {