Zustand 렌더링 최적화

최근 프로젝트 내의 ContextAPI를 모두 Zustand와 CustomHook 조합으로 갈아치우는 작업을 하였는데요. Zustand를 처음 도입해보느라, 엉성하게 코드를 작성한 것이 결국 작은 화를 불러일으켰습니다.

그럼 먼저 엉성한 코드의 슬픈 최후를 보고 가시겠습니다

보이시나요..? 이 아름다운 렌더링의 향연...!

AS-IS

기존에 headerContext로 관리하던 코드를 headerStore로 변경해, header에 관련된 상태와 setter함수를 모두 store에 보관하고, useHeaderStore를 통해 가져오도록 코드를 작성하였습니다.

export const useHeaderStore = create<HeaderStoreType>()(
  devtools(
    set => ({
      isTopBannerVisible: undefined,
      isScrollNavSticky: false,
      setIsTopBannerVisible: (isTopBannerVisible: boolean | undefined) =>
        set({
          isTopBannerVisible
        }),
      setScrollNavSticky: (isScrollNavSticky: boolean) =>
        set({
          isScrollNavSticky
        })
    }),
    { name: 'header-storage' }
  )
);

문제는 바로 이 zustand(정확히는 제가 쓴 Zustand 코드...)에 있었습니다.

const { isScrollNavSticky, setScrollNavSticky } = useHeaderStore(state => state);

저는 위와 같은 형태로 상태를 각 컴포넌트에서 구독하고 있었습니다!

TO-BE

위 코드의 문제는 두 가지가 있는데요. 아래와 같습니다.

1) useHeaderStore 내부에 콜백함수 전달 시 전체 state를 리턴시키고 있다는 점
2) 객체 자체를 strict하게 비교하고 있었다는 점

1 복수의 스테이트를 구독해야 할 때!

스토어로 부터 단일한 상태를 구독할 때는 아래와 같이 간단하게 작성할 수 있었습니다.

  const setIsTopBannerVisible = useHeaderStore(state => state.setIsTopBannerVisible);

두 개 이상의 상태를 구독해야 하다보니, state 전체를 반환하고, 구조분해 할당을 통해 상태를 추출해왔었는데요! 공식문서를 살펴보면, Zustand는 기본적으로 이전 상태와 새 상태를 엄격하게 비교하므로, atomic하게 상태를 추출할 것을 제안합니다!

It detects changes with strict-equality (old === new) by default, this is efficient for atomic state picks.

const nuts = useBearStore((state) => state.nuts)
const honey = useBearStore((state) => state.honey)

더 나아가 하나의 객체로 복수 상태를 뽑아낼 수도 있는데, 그럴 경우 아래와 같은 방법을 사용하라고 제안합니다.

2 얕은 비교하기!

useCustomStore의 옵션으로 우리는 Equality Function을 전달할 수 있습니다. 그리고 친절하게도 Zustand는 엄격한 비교가 아닌, 얕은 비교를 할 수 있는 Shallow라는 비교 함수를 제공해줍니다. 이 함수를 두번째 인자로 전달하면, 더 이상 변화된 상태를 비교할 때 객체의 주소값을 비교하는 것이 아니라, 내부의 원시값들 자체를 비교하게 되는 것이죠.

If you want to construct a single object with multiple state-picks inside, similar to redux's mapStateToProps, you can tell zustand that you want the object to be diffed shallowly by passing the shallow equality function.

import { shallow } from 'zustand/shallow'

// Object pick, re-renders the component when either state.nuts or state.honey change
const { nuts, honey } = useBearStore(
  (state) => ({ nuts: state.nuts, honey: state.honey }),
  shallow
)

// Array pick, re-renders the component when either state.nuts or state.honey change
const [nuts, honey] = useBearStore(
  (state) => [state.nuts, state.honey],
  shallow
)

// Mapped picks, re-renders the component when state.treats changes in order, count or keys
const treats = useBearStore((state) => Object.keys(state.treats), shallow)

그래서 저도 아래와 같은 코드를 완성하였습니다.

import { shallow } from 'zustand/shallow';

  const { isScrollNavSticky, setScrollNavSticky } = useHeaderStore(
    state => ({
      isScrollNavSticky: state.isScrollNavSticky,
      setScrollNavSticky: state.setScrollNavSticky
    }),
    shallow
  );

위와 같이 코드를 변경해주고 나니, 렌더링이 최소화된 것을 살펴볼 수 있었습니다.

참고

useStore에 두번째 인자로 전달하는 Equality Function에는 커스텀 Function도 얼마든지 적용할 수 있다고 합니다.

const treats = useBearStore(
  (state) => state.treats,
  (oldTreats, newTreats) => compare(oldTreats, newTreats)
)