Ramu Narasinga

Posted on Sep 19, 2024 • Edited on Apr 4 • Originally published at thinkthroo.com

useLayoutEffect in Zustand test-case explained.

#javascript #typescript #react #zustand

React’s useLayoutEffect is one of the lesser-used but crucial hooks that offers precise control over side effects that need to be executed before the browser paints the UI. In this article, we'll break down how useLayoutEffect works and demonstrate its usage with a test case from Zustand's source code. The test case provides a perfect example of how useLayoutEffect can help manage state and updates effectively when timing is critical.

This test case validates a bug fix related to an issue reported where the subscribed listener gets overwritten. Please note this issue was reported on Zustand’s version — 2.2.1. Link to codesandbox: https://codesandbox.io/s/quirky-taussig-ng90m

The latest stable Zustand’s version is 4.5.5 but it’s good to learn what this test case is about, especially when it uses useLayoutEffect

What is `useLayoutEffect`?

useLayoutEffect is similar to useEffect, but it fires synchronously after all DOM mutations and before the browser repaints the screen. It ensures that updates inside this hook are reflected on the page immediately, without the user experiencing any visual inconsistency.

In contrast, useEffect runs after the browser repaints the screen, meaning the user might see the DOM in an interim state before the effect takes place.

You can read more about useLayoutEffect in React docs.

When Should You Use `useLayoutEffect`?

You should use useLayoutEffect in situations where you need to ensure that the DOM has been updated before the browser paints the UI. Typical use cases include:

Measuring the DOM layout (e.g., for animations or measurements).
Synchronizing state or making changes that must be reflected in the DOM immediately.
Updating state that depends on DOM changes, such as adjusting styles or positions based on layout changes.

For less critical side effects, such as data fetching or logging, useEffect is generally preferred because it's non-blocking and doesn’t delay browser painting.

Zustand Test Case

Here’s a test case from Zustand’s source code that demonstrates useLayoutEffect in action:

// https://github.com/pmndrs/zustand/issues/84
it('ensures the correct subscriber is removed on unmount', async () => {
  const useBoundStore = create(() => ({ count: 0 }))
  const api = useBoundStore

  function increment() {
    api.setState(({ count }) => ({ count: count + 1 }))
  }

  function Count() {
    const c = useBoundStore((s) => s.count)
    return <div>count: {c}</div>
  }

  function CountWithInitialIncrement() {
    useLayoutEffect(increment, [])
    return <Count />
  }

  function Component() {
    const [Counter, setCounter] = useState(() => CountWithInitialIncrement)
    useLayoutEffect(() => {
      setCounter(() => Count)
    }, [])
    return (
      <>
        <Counter />
        <Count />
      </>
    )
  }

  const { findAllByText } = render(
    <>
      <Component />
    </>,
  )

  expect((await findAllByText('count: 1')).length).toBe(2)

  act(increment)

  expect((await findAllByText('count: 2')).length).toBe(2)
})

Breaking Down the Code

This test validates that Zustand correctly removes subscribers when components unmount and re-renders state changes appropriately. The core element we will focus on is how useLayoutEffect plays a critical role in controlling when state changes occur relative to the component lifecycle.

1. Creating the Store and Increment Function

The store is initialized with an object containing a single state property count set to 0.

const useBoundStore = create(() => ({ count: 0 }))
const api = useBoundStore

The increment function updates the count in the store:

function increment() {
    api.setState(({ count }) => ({ count: count + 1 }))
}

2. The `Count` Component

The Count component subscribes to the store and renders the current count:

function Count() {
    const c = useBoundStore((s) => s.count)
    return <div>count: {c}</div>
}

This component simply pulls the count from Zustand’s store and displays it in the DOM.

3. `CountWithInitialIncrement` and `useLayoutEffect`

This is where useLayoutEffect comes into play. In the CountWithInitialIncrement component, we use useLayoutEffect to trigger the increment function immediately after the component mounts but before the browser paints:

function CountWithInitialIncrement() {
    useLayoutEffect(increment, [])
    return <Count />
}

This ensures that the count is incremented before the component's UI is rendered. If we used useEffect here, the UI would first render count: 0, then update to count: 1 after the effect runs. However, with useLayoutEffect, the UI skips the initial count: 0 and directly renders count: 1.

4. Switching Components Dynamically

The Component component demonstrates a more advanced scenario where useLayoutEffect is used to switch between two different components on mount:

function Component() {
    const [Counter, setCounter] = useState(() => CountWithInitialIncrement)
    useLayoutEffect(() => {
      setCounter(() => Count)
    }, [])
    return (
      <>
        <Counter />
        <Count />
      </>
    )
}

Initially, Counter is set to CountWithInitialIncrement, which triggers the increment function when it mounts.
Then, useLayoutEffect runs synchronously after the DOM is updated, changing Counter to the Count component. This switch ensures that after the layout is painted, the next time the component renders, it doesn’t include the initial increment logic.

By the time Counter switches to Count, the count is already incremented, and the correct values are displayed.

5. The Test Expectations

Finally, the test verifies the following:

Initially, both the Counter and Count components display count: 1.
After triggering another increment, both components should update to count: 2.

expect((await findAllByText('count: 1')).length).toBe(2)

act(increment)

expect((await findAllByText('count: 2')).length).toBe(2)

The test passes because useLayoutEffect ensures the state update happens before the browser renders the UI, avoiding any intermediate render where count would still be 0.

Why Use `useLayoutEffect` Here?

this test case closely resembles the issue repro provided in the codesandbox.

The aim of this test case is to validate that subscribers do not get overwritten when the components unmount. Since useLayoutEffect renders the state before the browser repaints the UI, this is to ensure the listeners all work as expected. The issue states that one of listeners simply does not get updated, which is weird.

Conclusion

React’s useLayoutEffect provides control over state updates and DOM changes that need to happen before the user sees the page. In the Zustand test case we reviewed, useLayoutEffect ensures that the increment function is executed synchronously after the DOM updates, making sure that state changes are reflected immediately.

useLayoutEffect can hurt performance. Prefer useEffect when possible. — React Docs

About me:

Hey, my name is Ramu Narasinga. I study large open-source projects and create content about their codebase architecture and best practices, sharing it through articles, videos.

I am open to work on interesting projects. Send me an email at ramu.narasinga@gmail.com

My Github — https://github.com/ramu-narasinga

My website — https://ramunarasinga.com

My Youtube channel — https://www.youtube.com/@ramu-narasinga

Learning platform — https://thinkthroo.com

Codebase Architecture — https://app.thinkthroo.com/architecture

Best practices — https://app.thinkthroo.com/best-practices

Production-grade projects — https://app.thinkthroo.com/production-grade-projects