Zustand 잘 사용하고 계신가요? - 리팩터링 여정

2024-05-2910 min read

Zustand 정복하기

처음 프로젝트 코드를 파악하고 분석하는 과정에서 가장 많은 시간을 소요한 부분은 Zustand 관련 로직이었다.

해당 로직은 공식 문서를 기반으로 코드가 작성되지 않았고 임의로 코드가 구성되어 있었다.

타입 또한 전혀 고려되지 않아 코드의 안정성은 낮았으며 매번 공식 문서가 아닌 프로젝트 내 코드와 사람에 의존해서 코드가 구성되는 상황이었다.

이와 더불어 여태까지 Zustand를 사용해본 적이 없었기에 미지의 영역에서 코드를 작성한다는 두려움이 존재했다.

여기서 어떻게 판단해야 할까?

현재 상황에 순응하며 반복될 것으로 예상되는 어려움을 감수해야 할까?

코드를 빠르게 이해하면서 앞으로의 작업을 효율적으로 진행하기 위해서는 Zustand 관련 리팩터링 작업은 반드시 진행되어야 하는 작업이라는 생각이 들었다.

더욱이 해당 로직은 프로젝트 대부분의 영역에 쓰이는 코어 부분이다.

모든 작업자에게 프로젝트의 복잡성을 높이게 만드는 현재의 악순환을 끊고 싶었다.

준비 과정

우선 나부터가 Zustand에 대해 전혀 알지 못하는 문제를 해결해야 했다.

이를 해결하기 위해 팀스터디를 주도했다.

마침 Zustand의 메인테이너가 작성한 책이 존재해서 팀원분에게 스터디를 제안을 했다.

매주 정해진 분량을 읽어오며 토론하는 방식으로 약 한 달간의 스터디를 진행하며 라이브러리와 관련된 이해도를 높였다.

상태 관리 스터디 학습 기록

리액트 훅을 활용한 마이크로 상태 관리

해당 글에서는 스터디를 진행하며 개인적으로 핵심적이다고 느낀 Zustand의 내부 코드만 공유하고자 한다.

아마 책을 읽으신 분들은 코드들이 낯설지 않을 것이다.

모듈 상태 관리

핵심 로직 - Store 생성

외부 스토어 useSyncExternalStore

스터디를 진행하고 Zustand에 대한 두려움을 없앨 수 있었다.

이제 본격적인 문제 상황을 해결해 나가자.

문제 상황

문제 상황에 작성된 코드들은 리팩터링 전 상황을 보여주고자 임의로 재구성한 코드입니다.

참고하시길 바랍니다.

프로젝트에서 Zustand가 쓰이는 부분들만 분리해보자.

rootStore.tsx

import { createContext, useContext, useLayoutEffect } from 'react'
import { create } from 'zustand'
import { useStoreTest1 } from './store.test1'

type StoreState = { test1: number }

type StoreSelector<T> = (state: StoreState) => T

const initialState = {}

export const zustandContext = createContext(null)
export const StoreProvider = zustandContext.Provider

export const useStoreSSR = <T, _>(selector: StoreSelector<T>) => {
  const useStore: any = useContext(zustandContext)

  if (typeof useStore === 'function') {
    // eslint-disable-next-line react-hooks/rules-of-hooks
    return useStore(selector)
  }
}

export const initializeStore = (preloadedState?: any) => {
  const _create = (set: any, get: any) => ({
    ...initialState,
    ...preloadedState,

    ...useStoreTest1(set),
    // ...useStoreTest2(set),
    // ...useStoreTest3(set),
    // 기타 Store 정의들...
  })

  return create(_create)
}

export let store: any

export function useCreateStore(initialState: any) {
  if (typeof window === 'undefined') {
    return initializeStore(initialState)
  }

  store = store ?? initializeStore(initialState)

  // eslint-disable-next-line react-hooks/rules-of-hooks
  useLayoutEffect(() => {
    if (initialState && store) {
      store.setState({ ...initialState, ...store.getState() })
    }
  }, [initialState])

  return store
}

초기 상태가 정의되는 부분으로 크게 useStoreSSR, initializeStore, useCreateStore로 볼 수 있다.

useStoreSSR

역할

관심있는 상태만을 가져올 때 활용하고자 만든 함수이다.

인자로 selector가 넘어오는데 예시에서는 test1이 넘어온다고 생각하면 된다.

추후에 useTest1 함수에서 확인할 수 있다.

문제점

해당 함수에서는 상태 관리 로직인데 SSR이라는 네이밍도 적합하지 않다는 생각이 들며,

조건부로 훅을 사용하려다보니 훅 규칙을 위반했다.

initializeStore

역할

하나의 스토어지만 역할을 구분하기 위해 기능별로 스토어 역할 함수를 분리해서 정의한다.

예시에서는 store.test1 파일에서의 useStoreTest1 스토어를 구조 분해해서 _create 내부에 정의한다.

문제점

set, get은 Zustand에서 쓰이는 것과 동일한 역할을 하지만 타입이 전혀 정의되어 있지 않다.

즉 작업자가 미리 "여기는 set이 올거야. get이 올거야."라고 예측하고 코드를 작성하도록 한다.

useCreateStore

역할

서버와 클라이언트 상태의 싱크를 맞춰주는 역할을 한다.

문제점

2가지 코드 스멜이 의심된다.

store가 따로 정의될 필요가 있는가?
정말 useLayoutEffect이 필요한 것인가?

_app.tsx

import { StoreProvider, initializeStore, useCreateStore } from '@/store/rootStore'
import type { AppProps } from 'next/app'

export default function App({ Component, pageProps }: AppProps) {
  const store = useCreateStore(pageProps.initialZustandState)

  return (
    <StoreProvider value={store}>
      <Component {...pageProps} />
    </StoreProvider>
  )
}

App.getInitialProps = async () => {
  const zustandStore = initializeStore()

  zustandStore.getState().test1Prefetch()
  // zustandStore.getState().test2Prefetch();
  // zustandStore.getState().test3Prefetch();
  // 기타 Store 호출...

  return {
    pageProps: {
      initialZustandState: zustandStore.getState(),
    },
  }
}

_app.tsx 내 getInitialProps 부분에서 생성한 스토어들을 호출한다. 이는 서버에서 호출하는 것으로 클라이언트 이전에 미리 상태를 정의하는 역할을 한다.

서버에서 호출해서 변경된 상태를 pageProps 내 initialZustandState로 정의해서 전달한다. 전달된 값은 클라이언트에서 사용하도록 Context로 정의된 StoreProvider에 할당한다.

store.test1.ts

export const useStoreTest1 = (set: any) => ({
  test1Prefetch: () => {
    set((state: any) => {
      return { ...state, test1: 'hello test1' }
    })
  },
})

스토어 내부 상태를 정의하는 함수로 set 함수 내부에서 상태를 설정한다.

예시에서는 test1Prefetch 메서드를 통해 test에 hello test1을 정의한다.

이는 _app.tsx 내 getInitialProps에서 호출된다.

useTest1.tsx

import { useStoreSSR } from './rootStore'

const useTest1 = () => useStoreSSR((state) => state.test1)

export default useTest1

특정 상태만을 가져오고자 구성된 커스텀훅 부분이다.

예시에서는 test1만 가져오기 위해 useTest1으로 정의한다.

사용 예시

import useTest1 from '@/store/useTest1'

export default function Home() {
  const test1 = useTest1()

  return <div>{test1}</div>
}

useTest1 훅을 통해 test1을 불러오면 hello test1을 볼 수 있다.

정리하면 다음과 같은 문제점이 있다고 판단했다.

처음 코드를 접하는 작업자는 공식 문서가 아닌 프로젝트 내 코드만을 따라가면서 학습해야 하는 상황이다.
더욱이 타입이 any로 구성되어 있어서 코드의 추적을 어렵게 만든다.
모든 상태 기반 코드를 기억해야 하는데 이는 사실상 불가능하다.
결국 왜 사용하는지, 어떠한 흐름으로 동작하는지도 모르고 불안감을 가지고 코드를 작성해야 한다.

타입이 잡히지 않는 코드들

상황 분석

프로젝트 내에서 Zustand는 어떻게 쓰고 있는가?

현재 프로젝트에서 Zustand는 다음과 같이 사용되고 있다.

서버에서 사용되는 스토어내 상태들을 정의하기 위해 서버에서 호출한다.

클라어언트로 넘어오면 컨텍스트를 활용해서 서버에서 정의된 스토어를 다시 호출한다.

사용할 때는 커스텀훅을 활용해서 사용한다.

initializeStore에서 상태 및 스토어 구성
_app.tsx에서 초기 상태 정의
useCreateStore로 서버와 클라이언트 상태 동기화
useStoreSSR을 통해 커스텀훅을 구성해서 사용

리팩터링하기

큰 흐름은 파악이 되었다.

공식 문서를 기반으로 리팩터링을 진행해보자.

Slice 패턴 적용

initializeStore를 보면 Redux에서 많이 본 Slice 패턴을 적용할 수 있겠다는 생각이 들었다.

코드를 수정하면 다음과 같이 수정할 수 있다.

rootStore.tsx

import { create } from 'zustand'
import { createTest1Slice } from './store.test1.ts'

export type BaseStore = Test1Slice
// & Test2Slice
// & Test3Slice
// 기타 Slice 타입들...

export const useBoundStore = create<BaseStore>()((...a) => ({
  ...createTest1Slice(...a),
  // createTest2Slice(...a)
  // createTest3Slice(...a)
  // 기타 createSlice 정의들...
}))

store.test1.ts

import { StateCreator } from 'zustand'

export type Test1Slice = {
  test1: string
  test1Prefetch: () => void
}

export const createTest1Slice: StateCreator<Test1Slice, [], [], Test1Slice> = (set) => ({
  test1: '',
  test1Prefetch: () => {
    set(() => ({ test1: 'hello test1' }))
  },
})

타입을 위해 Slice에 사용되는 BaseStore 타입을 정의해주고 initializeStore는 useBoundStore로 대체된다.

서버와 클라이언트 상태 동기화

변경된 Slice 패턴에 맞추어 Store 및 Context 관련 코드도 타입 적용 및 로직 개선을 진행해보자.

rootStore.tsx

import { createContext, useContext, useRef } from 'react'
import { StoreApi, create, useStore } from 'zustand'

// Slice 적용 코드 생략

export const getStore = () => {
  return useBoundStore.getState()
}

export const InitStoreContext = createContext<StoreApi<BaseStore> | null>(null)

export const useInitStore = <T, _>(selector: (store: BaseStore) => T) => {
  const initStoreContext = useContext(InitStoreContext)

  if (!initStoreContext) {
    throw new Error('useInitStore must be use within InitStoreContextProvider')
  }

  return useStore(initStoreContext, selector)
}

export const initStore = (initialState: BaseStore) => {
  useBoundStore.setState(() => initialState)

  return useBoundStore
}

export const InitStoreProvider = ({
  children,
  initialState,
}: {
  children: React.ReactNode
  initialState: BaseStore
}) => {
  const storeRef = useRef<StoreApi<BaseStore>>()

  if (!storeRef.current) {
    storeRef.current = initStore(initialState)
  }

  return (
    <InitStoreContext.Provider value={storeRef.current}>
      {children}
    </InitStoreContext.Provider>
  )
}

useStoreSSR는 해당 상태 설정이 처음 한 번에만 이루어진다고 생각되어 useInitStore으로 변경했다.

Context 부분도 기존의 store 및 useLayoutEffect을 없애고 useRef를 활용한 방식으로 수정했다.

ref.current가 존재하지 않을 때, 즉 처음 클라이언트 호출할 때는 initStore 함수를 사용해 스토어의 초기 상태를 정의한다. 여기서 서버에서 정의되었던 스토어와의 싱크를 다시 맞추는 역할을 한다.

이를 통해 서버에서 설정한 상태 초기값을 클라이언트 스토어에 동기화시켜서 사용할 수 있다.

_app.tsx

import { InitStoreProvider, getStore } from '@/store/rootStore'
import type { AppProps } from 'next/app'

export default function App({ Component, pageProps }: AppProps) {
  return (
    <InitStoreProvider initialState={pageProps.initialZustandState}>
      <Component {...pageProps} />
    </InitStoreProvider>
  )
}

App.getInitialProps = async () => {
  const { test1Prefetch } = getStore()

  test1Prefetch()

  return {
    pageProps: {
      initialZustandState: getStore(),
    },
  }
}

getStore 함수는 rootStore.tsx에서 정의한 함수이다.

서버에서 스토어에 접근할 때 주의할 점은 스토어 액션없이 접근이 가능해야 한다.

Zustand는 외부 스토어와 상태간에 동기화를 위해 usesyncexternalstore를 사용한다.

이때 스토어 액션을 통해 접근을 시도하면 에러가 발생한다.

이를 피하기 위해 getState 메서드를 활용해서 스토어에 접근한다.

useTest1.ts

import { useInitStore } from './rootStore'

const useTest1 = () => useInitStore((state) => state.test1)

export default useTest1

useStoreSSR가 아닌 useInitStore로 커스텀훅을 구성한다.

사용 예시

import useTest1 from '@/store/useTest1'

export default function Home() {
  const test1 = useTest1()

  return <div>{test1}</div>
}

동일하게 사용하면 된다.

타입이 잡히는 코드

번외 - Zustand에서 Context를 사용하는 이유는 뭘까?

리팩터링 과정에서 들었던 의문은 Zustand에서 Context를 사용하는 부분이었다.

Context에서 발생하는 렌더링 문제를 해결하기 위해 Zustand와 같은 전역 상태 라이브러리들이 등장한 것으로 이해했기 때문에 괴리감을 느끼게 되었다.

다행히 이러한 의문점을 해소해준 좋은 아티클을 읽게 되었다.

Zustand and React Context

props로 초기화할 수 없는 문제
테스트
재사용성

여기서 우리가 Context로 넘기는 것은 Store라는 점이 핵심이다.

최종 코드

최종 코드는 다음과 같다.

rootStore.tsx

import { createContext, useContext, useRef } from 'react'
import { StoreApi, create, useStore } from 'zustand'
import { Test1Slice, createTest1Slice } from './store.test1'

export type BaseStore = Test1Slice
// & Test2Slice
// & Test3Slice
// 기타 Slice 타입들...

export const useBoundStore = create<BaseStore>()((...a) => ({
  ...createTest1Slice(...a),
  // createTest2Slice(...a)
  // createTest3Slice(...a)
  // 기타 createSlice 정의들...
}))

export const getStore = () => {
  return useBoundStore.getState()
}

export const InitStoreContext = createContext<StoreApi<BaseStore> | null>(null)

export const useInitStore = <T, _>(selector: (store: BaseStore) => T) => {
  const initStoreContext = useContext(InitStoreContext)

  if (!initStoreContext) {
    throw new Error('useInitStore must be use within InitStoreContextProvider')
  }

  return useStore(initStoreContext, selector)
}

export const initStore = (initialState: BaseStore) => {
  useBoundStore.setState(() => initialState)

  return useBoundStore
}

export const InitStoreProvider = ({
  children,
  initialState,
}: {
  children: React.ReactNode
  initialState: BaseStore
}) => {
  const storeRef = useRef<StoreApi<BaseStore>>()

  if (!storeRef.current) {
    storeRef.current = initStore(initialState)
  }

  return (
    <InitStoreContext.Provider value={storeRef.current}>
      {children}
    </InitStoreContext.Provider>
  )
}

store.test1.ts

import { StateCreator } from 'zustand'

export type Test1Slice = {
  test1: string
  test1Prefetch: () => void
}

export const createTest1Slice: StateCreator<Test1Slice, [], [], Test1Slice> = (set) => ({
  test1: '',
  test1Prefetch: () => {
    set(() => ({ test1: 'hello test1' }))
  },
})

_app.tsx

import { InitStoreProvider, getStore } from '@/store/rootStore'
import type { AppProps } from 'next/app'

export default function App({ Component, pageProps }: AppProps) {
  return (
    <InitStoreProvider initialState={pageProps.initialZustandState}>
      <Component {...pageProps} />
    </InitStoreProvider>
  )
}

App.getInitialProps = async () => {
  const { test1Prefetch } = getStore()

  test1Prefetch()

  return {
    pageProps: {
      initialZustandState: getStore(),
    },
  }
}

useTest1.tsx

import { useInitStore } from './rootStore'

const useTest1 = () => useInitStore((state) => state.test1)

export default useTest1

마치며

일하면서 성장하기

이번 Zustand 리팩터링 작업을 진행하면서 좋았던 점은 업무와 함께 나의 성장도 이끌어낼 수 있었다.

무지의 영역을 팀원과 함께 스터디하면서 잘 극복한 것 같다.

앞으로도 프로젝트내 기술을 기반으로 학습을 진행하고자 한다.

나의 성장과 코드 품질 개선을 함께 이끌어내는 선순환을 만들어내고 싶다.

Zustand를 의존성 주입 도구로만 활용하는 것은 아닌가?

현재 프로젝트에서 쓰고 있는 형태는 전역 상태 관리가 아닌 의존성 주입을 편리하게 도와주는 도구라는 생각이 든다.

물론 의존성 주입의 형태라고 Zustand를 쓰지 말라는 이유는 없어서 굳이 Zustand를 없애야 할 필요는 없는 것 같다.

어떠한 흐름으로 쓰이고 있는지 이해하고 상황에 맞게 적절하게 사용하면 괜찮지 않나 싶다.

테스트 코드의 필요성을 느끼며

기존 코드를 수정하다보니 미처 고려하지 못한 사이드이펙트가 생기는 것은 아닐까 라는 고민을 하게 되었다.

그러면서 자연스럽게 테스트 코드에 관심이 생기게 되는 것 같다.

그래서 다음 프로젝트 개선 작업으로 테스트코드 도입을 하고 싶다는 생각이 들었다.

지금처럼 팀과 함께 학습하면서 프로젝트에 적용하는 단계를 거칠 것 같다.

처음 프로젝트를 접하는 누구라도 불안감없이 편하게 코드를 작성하는 환경이 됐으면 한다.

Zustand 잘 사용하고 계신가요? - 리팩터링 여정

.css-199glky{-webkit-transition:0.3;transition:0.3;font-weight:700;}.css-199glky:not(.anchor){border-bottom:2px solid #c471f5;}.css-199glky:hover{color:#c471f5;}Zustand 정복하기

준비 과정

문제 상황

상황 분석

프로젝트 내에서 Zustand는 어떻게 쓰고 있는가?

리팩터링하기

Slice 패턴 적용

서버와 클라이언트 상태 동기화

번외 - Zustand에서 Context를 사용하는 이유는 뭘까?

최종 코드

마치며

일하면서 성장하기

Zustand를 의존성 주입 도구로만 활용하는 것은 아닌가?

테스트 코드의 필요성을 느끼며

참고 문서

Zustand 정복하기