Deep Dive into Zustand

April 20, 2024

들어가며

그동안 실무나 개인 프로젝트에서 클라이언트 상태 관리를 위해 zustand를 자주 사용해왔다. 적은 패키지 사이즈, 불필요한 보일러 플레이트 코드, flux 아키텍쳐를 계승한 익숙한 스타일 등이 장점으로 작용했던 것 같다. 그런데 정작 zustand를 사용하며 zustand의 내부 동작원리에 대해 깊이 고민해본 적은 없었던 것 같다. zustand를 사용해보면 알겠지만, Provider나 별도의 설정없이도 알아서 상태 변화를 감지하고 컴포넌트를 리렌더링한다. 어떻게 이런 마법이 가능한 것일까? 이번 글에서는 zustand의 소스코드를 직접 살펴보며 동작원리를 파헤쳐 보는 시간을 가지려고 한다.

zustand

해당 글은 zustand v4.5.2 버전을 기준으로 작성되었다.

먼저 zustand의 간단한 사용법을 알아보자. zustand와 같은 상태 관리 라이브러리를 사용하며 가장 먼저 할 일은 아마 상태와 액션을 저장하는 스토어를 만드는 일이다. zustand에서는 create 함수를 통해 손쉽게 스토어를 만들 수 있다. create 함수는 객체를 반환하는 함수를 인자로 갖는데 이때 반환되는 객체에 우리가 관리하고자 하는 상태와 액션이 담긴다.

// store.js
import { create } from 'zustand'

const useBearStore = create((set) => ({
  bears: 0,
  increasePopulation: () => set((state) => ({ bears: state.bears + 1 })),
  removeAllBears: () => set({ bears: 0 }),
}))

스토어를 만들었으면 이제 상태를 사용하는 컴포넌트에서 필요한 데이터를 불러오면 된다. zustand는 스토어를 리액트 훅의 형태로 제공해서 컴포넌트 내에서 간편하게 상태를 가져올 수 있다. 또한 가져온 상태가 변경될 때마다 컴포넌트는 리렌더링되고 zustand 스토어와 리액트 컴포넌트의 동기화가 일어난다. 예를 들어 Controls 컴포넌트에서 increatePopulation 액션을 호출한다면, bears 값이 변경되므로 BearCount 컴포넌트도 새로운 상태를 위해 리렌더링된다.

function BearCount() {
  const bears = useBearStore((state) => state.bears)
  return <h1>{bears} around here ...</h1>
}

function Controls() {
  const increasePopulation = useBearStore((state) => state.increasePopulation)
  return <button onClick={increasePopulation}>one up</button>
}

이렇듯 사용법은 정말 간단하다. 그런데 어떻게 이렇게 간단한 방법만으로 상태를 관리할 수 있는 것일까? 어떻게 React Context와 같은 별도의 설정없이도 컴포넌트는 상태의 변경을 감지할 수 있는 것일까? zustand의 소스코드를 분석해보며 그 답을 찾아보자.

Core

이번 글에서는 주요 로직에 집중하기 위해 타입스크립트 대신 자바스크립트를 사용했다. 일부 예외 상황에 대한 부분도 생락했다. 원본 코드를 확인하고 싶으면 직접 Repository를 방문해보기를 권한다.

우리는 위 예시에서 스토어를 만들때 zustand의 create 함수를 사용했다. 그렇다면 이 함수는 어떻게 구현되어 있을까? 아래는 zustand의 react.ts 파일에 구현된 create 함수다.

// https://github.com/pmndrs/zustand/blob/main/src/react.ts
import { createStore } from './vanilla.js'

const create = (createState) => {
  const api = typeof createState === 'function' ? createStore(createState) : createState

  const useBoundStore = (selector, equalityFn) => useStore(api, selector, equalityFn)

  Object.assign(useBoundStore, api)

  return useBoundStore
}

create 함수는 createState라는 인자를 받고있다. createState도 함수인데 이 함수는 우리가 처음의 예시에서 useBearStore를 만들 때 create 함수에 넘겨준 그 함수를 뜻한다.

const useBearStore = create((set) => ({
  bears: 0,
  // (...)
}))

createState 함수에 대해서는 나중에 더 자세히 살펴보기로 하고 넘어가보자.

다음으로는 vanilla.js 파일에서 createStore라는 함수를 불러오고 그를 통해 api 객체를 만들어주고 있다. 먼저 createStore라는 함수에 대해 알아봐야 할 것 같다. vanilla.js 파일로 가보자.

createStore

// https://github.com/pmndrs/zustand/blob/main/src/vanilla.ts

// (...)

export const createStore = (createState) => {
  let state
  const listeners = new Set()

  const setState = (partial, replace) => {
    const nextState = typeof partial === 'function' ? partial(state) : partial

    if (!Object.is(nextState, state)) {
      const previousState = state
      state =
        replace ?? (typeof nextState !== 'object' || nextState === null)
          ? nextState
          : Object.assign({}, state, nextState)

      listeners.forEach((listener) => listener(state, previousState))
    }
  }

  const getState = () => state

  const getInitialState = () => initialState

  const subscribe = (listener) => {
    listeners.add(listener)
    return () => listeners.delete(listener)
  }

  const api = { setState, getState, getInitialState, subscribe }
  const initialState = (state = createState(setState, getState, api))
  return api
}

// (...)

createStore 함수는 createState 함수를 인자로 받고 api라는 객체를 반환하고 있다. 그리고 api 객체는 setState, getState, getInitialState, subscribe이라는 메서드를 담고 있다. zustand에서 가장 중요한 역할을 하는 함수이니 코드를 한줄씩 자세히 살펴보자.

상태와 리스너

// store의 모든 상태가 담기는 변수
let state
// 상태가 변경될때마다 실행되는 listener들이 담겨있는 Set 자료형
const listeners = new Set()

스토어가 관리하는 상태가 담기는 state와 상태 변경 시 변경을 전파하는 listener를 담은 listeners 변수는 함수 내부에서 선언되었다. 아마 클로저를 통해 함수 외부에서 해당 변수들에 접근할 수 있을 것으로 보인다.

`setState`

setState는 상태의 변경을 담당하는 함수이다.

const setState = (partial, replace) => {
  const nextState = typeof partial === 'function' ? partial(state) : partial

  if (!Object.is(nextState, state)) {
    const previousState = state
    state =
      replace ?? (typeof nextState !== 'object' || nextState === null)
        ? nextState
        : Object.assign({}, state, nextState)

    listeners.forEach((listener) => listener(state, previousState))
  }
}

먼저 전달받은 partial 인자를 통해 nextState을 구한다. 이때 partial은 객체일 수도 있고 상태를 업데이트하는 함수일 수도 있다. 만약 partial을 통해 구한 상태가 현재 상태와 달라졌을 경우 상태를 업데이트 한다. 이때 nextState 가 객체라면 기존 상태와 병합된 상태를 현재 상태로 할당한다. 그 외에 replace 인자가 true이거나, nextState가 객체 타입이 아닌 값이라면 병합없이 해당 값을 새로운 상태로 덮어씌운다. 갑자기 병합이라는 개념이 나와서 헷갈릴 수도 있는데 이 부분은 실제 zustand의 사용 예시를 생각해보면 이해하기 쉽다.

동물의 수를 관리하는 useAnimalStore라는 스토어가 있다고 생각해보자.

import { createStore } from 'zustand'

const useAnimalStore = createStore((set) => ({
  bears: 0,
  dogs: 0,
  incrementBear: () => set((state) => ({ bears: state.bears + 1 })),
  incrementDog: () => set((state) => ({ dogs: state.dogs + 1 })),
  deleteEverything: () => set({}, true),
}))

const { getState } = useAnimalStore
log(getState()) // {bears: 0, dogs: 0, ...}
getState().incrementBear()
log(getState()) // {bears: 1, dogs: 0, ...}
getState().incrementDog()
log(getState()) // {bears: 1, dogs: 1, ...}
getState().incrementBear()
log(getState()) // {bears: 2, dogs: 1, ...}
getState().deleteEverything()
log(getState()) // {}

incrementBear나 incrementDog 메서드는 각각 bear와 dog 하나의 상태만 업데이트하고 있다. 하지만 해당 메서드를 실행한 후 스토어의 상태를 콘솔에 찍어보면 전체 상태는 그대로 유지된 채 변경된 상태가 병합되는 형태로 업데이트가 일어나는 것을 알 수 있다. 반면, deleteEverything 메서드와 같이 set 함수의 두번째 인자로 true를 넘겨주면 새로운 상태를 기존 상태에 덮어씌우는 식으로 업데이트되어 최종적으로 빈 객체가 출력되는 것을 알 수 있다.

state 변수에 새로운 상태를 할당한 이후에는 스토어를 구독하는 리스너들을 호출한다. 이 부분을 통해 스토어를 구독하고 있는 컴포넌트가 변경을 감지하고 리렌더링될 수 있을 것으로 예상된다. 자세한 내용은 뒤에서 다뤄보겠다.

listeners.forEach((listener) => listener(state, previousState))

`getState` & `getInitialState`

const getState = () => state
const getInitialState = () => initialState

이 두 메서드는 각각 스토의 현재 상태와 초기 상태를 리턴한다. 이를 통해 createStore 함수 외부에서 항상 현재 상태를 조회할 수 있다.

`subscribe`

const subscribe = (listener) => {
  listeners.add(listener)
  return () => listeners.delete(listener)
}

상태 변경시 실행되는 리스너 함수를 추가할 수 있는 메서드다. 동시에 해당 리스너를 제거할 수 있는 cleanup 함수를 반환한다.

최종적으로 createStore 함수는 위의 메서드들을 api라는 객체에 담아 반환하고, 전달받은 createState 함수를 사용해 상태를 초기화한다.

// (...)
const api = { setState, getState, getInitialState, subscribe }
const initialState = (state = createState(setState, getState, api))
return api

위에서 한번 언급했다시피 createState 함수는 스토어를 만드는 create 함수의 인자로 제공되는 콜백함수인데, 예시에서 사용된 set 함수이외에도 getState, api 라는 2가지 인자를 더 갖는다는 것을 확인할 수 있다. 그러니까 create 함수는 아래 코드블록과 같이도 사용될 수 있다.

const useCountStore = create((set, get, api) => ({
  increment: () => {
    const { count } = get()
    set({ count: count + 1 })
    console.log(`초기값은 ${api.getInitialState()}`)
  },
}))

이제 createStore 함수를 모두 살펴보았다. 간단히 정리해보면, 스토어의 초기 상태를 결정하는 createState 함수를 인자로 받아 스토어를 만들고 상태를 읽고, 쓰고, 구독하는 메서드(api)들을 만들어 외부로 반환하는 역할을 하고 있다.

createStore 함수가 어떻게 동작하는지 파악했다. 이제 다시 create 함수로 돌아가보자.

create

// https://github.com/pmndrs/zustand/blob/main/src/react.ts
export function useStore(api, selector = identity, equalityFn) {
  const slice = useSyncExternalStoreWithSelector(
    api.subscribe,
    api.getState,
    api.getServerState || api.getInitialState,
    selector,
    equalityFn,
  )
  return slice
}

const create = (createState) => {
  const api = typeof createState === 'function' ? createStore(createState) : createState

  const useBoundStore = (selector, equalityFn) => useStore(api, selector, equalityFn)

  Object.assign(useBoundStore, api)

  return useBoundStore
}

create는 createStore 함수를 리액트에서 사용할 수 있도록 확장한 함수라고 생각하면 된다. 우리는 이제 createStore 함수가 어떤 내부 메서드를 노출하는지 알고 있으므로 아래 코드블록을 이해할 수 있다.

const api = typeof createState === 'function' ? createStore(createState) : createState

그 다음으로는 useBoundStore이라는 함수를 선언하고 내부에서 useStore이라는 함수를 다시 호출하고 있는데 useStore 함수에 대해서도 이해할 필요가 있어 보인다.

const useBoundStore = (selector, equalityFn) => useStore(api, selector, equalityFn)

export function useStore(api, selector = identity, equalityFn) {
  const slice = useSyncExternalStoreWithSelector(
    api.subscribe,
    api.getState,
    api.getServerState || api.getInitialState,
    selector,
    equalityFn,
  )
  return slice
}

useStore 는 api 객체와 selector 함수, 그리고 비교를 수행하는 것으로 예상되는 equalityFn이라는 함수를 인자로 받고 있다. 그리고 함수 내부에서는 useSyncExternalStoreWithSelector 라는 훅의 결과값을 slice에 할당하고 있는데 useSyncExternalStoreWithSelector 훅이 다소 생소하다. 이를 이해하기 위해서는 React18에 추가된 useSyncExternalStore 훅을 먼저 이해해야 한다.

`useSyncExternalStore`

useSyncExternalStore is a React Hook that lets you subscribe to an external store. - react.dev

useSyncExternalStore은 리액트 외부에 존재하는 스토어를 구독할 수 있게 도와주는 훅이다. 여기서 리액트 외부에 존재하는 스토어란 props, state, context 등과 같이 리액트 라이프 사이클 내에서 관리되는 데이터가 아닌 Redux와 같은 전역 상태 라이브러리나 window.navigator, window.location 등과 같은 Browser API를 말한다.

const snapshot = useSyncExternalStore(subscribe, getSnapshot, getServerSnapshot?)

useSyncExternalStore은 2개의 필수 인자와 1개의 옵셔널한 인자를 받는다.

subscribe: 스토어를 구독하는 로직이 담기는 함수다. subscribe 함수는 callback을 인자로 받는데 스토어에서 변경이 일어날 때 callback 함수를 호출함으로써 리액트에게 상태의 변경을 알릴 수 있다.
getSnapshot: 스토어의 스냅샷을 반환하는 함수다.
getServerSnapshot (optional): 스토어 초기 상태의 스냅샷을 반환하는 함수다. hydration을 위해 서버와 클라이언트에서의 스냅샷 결과를 일치시키기 위한 용도로 사용된다.

예시를 통해 자세히 살펴보자. 아래는 todoStore라는 외부 스토어를useSyncExternalStore를 통해 구독하고 있다. todoStore의 subscribe 메서드는 전달받은 listener를 (위에서 언급한 callback 인자) 배열에 저장해두고, addTodo 메서드가 스토어의 상태를 변경할 때 emitChange를 호출함으로써 변경을 알리는 역할을 한다. getSnapshot 메서드는 현재 스토어의 상태를 반환하고 있다.

import { useSyncExternalStore } from 'react'
import { todosStore } from './todoStore.js'

// App.js
function TodosApp() {
  const todos = useSyncExternalStore(todosStore.subscribe, todosStore.getSnapshot)
  return (
    <>
      <button onClick={() => todosStore.addTodo()}>Add todo</button>
      <hr />
      <ul>
        {todos.map((todo) => (
          <li key={todo.id}>{todo.text}</li>
        ))}
      </ul>
    </>
  )
  // ...
}

// todoStore.js
let nextId = 0
let todos = [{ id: nextId++, text: 'Todo #1' }]
let listeners = []

export const todosStore = {
  addTodo() {
    todos = [...todos, { id: nextId++, text: 'Todo #' + nextId }]
    emitChange()
  },
  subscribe(listener) {
    listeners = [...listeners, listener]
    return () => {
      listeners = listeners.filter((l) => l !== listener)
    }
  },
  getSnapshot() {
    return todos
  },
}

function emitChange() {
  for (let listener of listeners) {
    listener()
  }
}

TodosApp 컴포넌트에서 Add todo 버튼이 클릭되면 todoStore의 addTodo 메서드가 호출되고 todos 상태가 변경되는 동시에 listener 함수가 실행된다. 이를 통해 TodosApp의 useSyncExternalStore 훅이 변화를 감지하고 컴포넌트를 리렌더링한다. 이러한 방식으로 리액트는 외부 스토어의 상태 변화를 구독할 수 있게 된다.

자 이제 useSyncExternalStore의 동작방식을 어느정도 이해했으니 다시 zustand의 useStore 함수로 돌아가보자.

// https://github.com/pmndrs/zustand/blob/main/src/react.ts
import { useSyncExternalStoreWithSelector } from 'use-sync-external-store/shim/with-selector'

export function useStore(api, selector = identity, equalityFn) {
  const slice = useSyncExternalStoreWithSelector(
    api.subscribe,
    api.getState,
    api.getServerState || api.getInitialState,
    selector,
    equalityFn,
  )
  return slice
}

useSyncExternalStoreWithSelector 는 리액트 18 이전 버전에서도 useSyncExternalStore 훅을 사용할 수 있도록 리액트 팀에서 제공한 shim이다. selector 함수를 통해 스토어 상태의 특정 값을 받아올 수 있고 직접 렌더링 여부를 결정하는 equalityFn 함수를 넘길 수 있다는 점을 제외하면 useSyncExternalStore 과 동작은 같다.

useStore 훅은 createStore 함수에서 반환된 메서드들(api 객체)을 useSyncExternalStoreWithSelector의 인자로 넘겨주는 역할을 한다. createStore 내부 상태의 변경이 일어나면 api.subscribe 함수가 구독중인 listener들을 호출하고 useSyncExternalStoreWithSelector 훅은 그에 따라 컴포넌트를 리렌더링할 것임을 예상할 수 있다.

create 함수는 마지막으로 useBoundStore에 api 객체를 할당하고 반환한다. 굳이 useBoundStore에 api 객체를 추가하는 이유는 훅을 사용할 수 없는 리액트 컴포넌트 밖에서도 스토어의 api를 통해 상태를 관리할 수 있도록 지원하기 위함인 것 같다. 공식 문서에도 이러한 용례에 대한 설명이 있다.( ## Reading/writing state and reacting to changes outside of components)

const create = (createState) => {
  // ✅ createStore 함수를 통해 상태를 조회하고 변경하고 구독할 수 있는 메서드들을 반환한다.
  const api = typeof createState === 'function' ? createStore(createState) : createState

  // ✅ useSyncExternalStoreWithSelector 훅으로 스토어에서 원하는 상태를 구독한다.
  const useBoundStore = (selector, equalityFn) => useStore(api, selector, equalityFn)

  // ✅ useBoundStore에 api 객체를 할당한다.
  Object.assign(useBoundStore, api)

  return useBoundStore
}

이게 끝이다! zustand는 이렇게 간단한 로직만으로 동작한다. 물론 더 많은 기능을 미들웨어를 통해 제공하고 있지만 핵심 로직은 우리가 위헤서 살펴본 것이 전부다. 이제 다시 처음의 예시로 돌아가면 코드가 완전히 새롭게 보일 것이다.

// store.js
import { create } from 'zustand'

const useBearStore = create((set) => ({
  bears: 0,
  increasePopulation: () => set((state) => ({ bears: state.bears + 1 })),
  removeAllBears: () => set({ bears: 0 }),
}))

// bear.js
function BearCount() {
  const bears = useBearStore((state) => state.bears)
  return <h1>{bears} around here ...</h1>
}

function Controls() {
  const increasePopulation = useBearStore((state) => state.increasePopulation)
  return <button onClick={increasePopulation}>one up</button>
}

create 함수가 인자로 받고 있는 함수는 set이라는 함수를 또 인자로 받는데, 우리는 이 set이 createStore 함수 내에서 만들어진 setState 함수라는 것을 알 수 있다. 그리고 set 함수는 state를 인자로 가지는 콜백 함수를 받는데 이 함수가 setState 함수의 인자인 partial 이라는 것도 알 수 있다. set 함수가 실행되면 useSyncExternalStoreWithSelector 훅에 의해 변화가 감지되고 상태를 구독중인 컴포넌트(BearCount)가 리렌더링될 것이다.

마치며

오늘은 그동안 자주 사용하던 zustand의 소스코드를 분석하며 동작원리를 이해해보는 시간을 가졌다. 사실 라이브러리의 역할은 복잡한 구현은 감추고 간단한 API만을 노출함으로써 사용자들에게 편리함을 제공하는 것이다. 하지만 강력한 라이브러리들을 사용할수록 역설적으로 그 내부가 어떻게 구성되어 있을지 항상 궁금하곤 했다. 이번 계기로 zustand에 대해 더 깊은 이해가 생긴 것 같아서 개인적으로도 가치있는 시간이었다. zustand의 제작자 카도 다이시는 zustand 외에도 jotai, valtio 등 가볍고 강력한 상태 관리 라이브러리들을 만든 것으로 유명한데, 기회가 되면 이들에 대해서도 소스코드 레벨에서 분석해보는 시간을 가져보려고 한다.