abonglog logoabonglog

zustand는 어떻게 마법같이 동작할까? 의 썸네일

zustand는 어떻게 마법같이 동작할까?

웹 브라우저 지식

프로필 이미지
yonghyeun

해당 동작을 공부하게 된 계기는 다음과 같다.

선택적 구독을 통해 특정 값만 구독하려 한 경우
interface Chracters {
  ...
  id : number
}[]
 
const CharacterInputItem: React.FC<CharacterInputItemProps> = ({ id }) => {
  const characterIds = useNovelSettingFormStore((state) =>
    state.characters.map(({ id }) => id)
  );

zustand 스토어에 저장된 어떤 배열에서 원하는 값만 꺼내와 구독을 해두려고 했을 때 이런 에러가 발생했다.

해당 에러의 콜스택해당 에러의 콜스택

console.error : The result of getSnapshot should be cached to avoid an infinite loop

그래서 에러가 발생한 콜스택을 따라가다보니 에러가 발생한 부분이 zustand 가 아닌 리액트의 useSyncExternalState 에서 발생했더라

말이 나온김에 zustand 의 코드 내부를 살펴보며 왜 저런 에러가 발생했는지 살펴보도록 하자

Zustand의 내부구현 살펴보기

zustand 의 가장 큰 장점이라 한다면 아마 Context 없이 사용 가능하단 점과 선택적 구독을 통해 렌더링 최적화가 가능하다는 점일 것이다.

zustand 의 깃허브 코드를 들어가보면 코드가 매우 단순하게 작성되어있어 이해하는데 오랜 시간이 걸리지 않았다.

전체적인 흐름은 src/vanila.ts 에 생성되어있는 100줄 남짓한 코드와 해당 코드의 메소드를 활용하여 구현된 src/react.ts 의 60줄 남짓한 코드만 보면 됐기 때문이다.

타입선언들을 제외하면 두 코드의 합은 100줄도 되지 않는다!

이번 글에선 복잡한 타입 선언들을 제외하고 적도록 한다. 타입시스템이 정말 정말 고수스러운 코드여서 아직 내가 소화하긴 어렵더라 🫠

src/vanila.ts

스토어를 생성하는 createStoreImpl
const createStoreImpl: CreateStoreImpl = (createState) => {
  type TState = ReturnType<typeof createState>
  type Listener = (state: TState, prevState: TState) => void
  let state: TState
  const listeners: Set<Listener> = new Set()
 
  const setState: StoreApi<TState>['setState'] = (partial, replace) => {
    // TODO: Remove type assertion once https://github.com/microsoft/TypeScript/issues/37663 is resolved
    // https://github.com/microsoft/TypeScript/issues/37663#issuecomment-759728342
    const nextState =
      typeof partial === 'function'
        ? (partial as (state: TState) => TState)(state)
        : partial
    if (!Object.is(nextState, state)) {
      const previousState = state
      state =
        (replace ?? (typeof nextState !== 'object' || nextState === null))
          ? (nextState as TState)
          : Object.assign({}, state, nextState)
 
 
      listeners.forEach((listener) => listener(state, previousState))
    }
  }
 
  const getState: StoreApi<TState>['getState'] = () => state
 
  const getInitialState: StoreApi<TState>['getInitialState'] = () =>
    initialState
 
  const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
    listeners.add(listener)
    // Unsubscribe
    return () => listeners.delete(listener)
  }
 
  const api = { setState, getState, getInitialState, subscribe }
  const initialState = (state = createState(setState, getState, api))
  return api as any
}
 
export const createStore = ((createState) =>
  createState ? createStoreImpl(createState) : createStoreImpl) as CreateStore

해당 메소드는 상태를 저장하는 뮤테이블한 객체인 state 와 리스너들을 담는 immutable 한 Set 자료구조인 lisnters 객체가 존재한다.

이후 해당 객체들을 바라보는 메소드들을 담은 객체 api 를 반환한다.

이게 끝이다.

api 메소드들을 보면 알 수 있듯이 getState , getInitialState , subscribe 모두 우리가 create 메소드를 통해 생성한 스토어에서 호출 할 수 있는 static method 인걸 확인 할 수 있다.

차례 차례 로직들을 살펴보자

  1. subscribe 메소드가 호출되면 현재의 상태와 이전의 상태를 인수로 받는 listener 메소드가 listeners 자료구조에 담긴다.
  2. setState 메소드가 호출되면 state 객체가 새롭게 변경되고 listeners 자료구조에 존재하는 모든 listener 메소드들을 호출한다.

이게 끝이다.

src/react.ts

우리가 사용하는 흐름에 따라 차례차례 살펴보자

우리가 zustand 를 사용 할 때 이렇게 사용한다.

일반적으로 사용하는 방식
/* create 메소드로 스토어 생성 */
const useSomethingStore = create<T>(({set , get})=>{...})
/* selector를 이용해 선택적 구독 */
const something = useSomethingStore(state => state.something)

우리가 create 메소드를 호출하면 이런 일이 발생한다.

create 가 호출됐을 때
const createImpl = <T>(createState: StateCreator<T, [], []>) => {
  const api = createStore(createState) // src/vanila 로 구현된 createStore 메소드
 
  const useBoundStore: any = (selector?: any) => useStore(api, selector)
 
  Object.assign(useBoundStore, api)
 
  return useBoundStore
}
 
export const create = (<T>(createState: StateCreator<T, [], []> | undefined) =>
  createState ? createImpl(createState) : createImpl) as Create
  1. createState 를 이용하여 state , listeners 를 바라보는 메소드들을 담은 객체 api 를 생성한다. ({getState , setState ...})

createState 는 다음과 같이 생겼다. 우리가 create(({set , get})=>{...}) 내부 인수들의 타입을 의미한다.

StateCreator
export type StateCreator<
 T,
 Mis extends [StoreMutatorIdentifier, unknown][] = [],
 Mos extends [StoreMutatorIdentifier, unknown][] = [],
 U = T,
> = ((
 setState: Get<Mutate<StoreApi<T>, Mis>, 'setState', never>,
 getState: Get<Mutate<StoreApi<T>, Mis>, 'getState', never>,
 store: Mutate<StoreApi<T>, Mis>,
) => U) & { $$storeMutators?: Mos }
  1. useBoundStore 훅을 생성한다. 해당 훅은 selector 를 이용하여 useStore 를 호출한다.

  2. useBoundStore 훅에 api 들을 Objest.assign 로 적용해 api 내부 메소드들을 정적으로 호출 할 수 있게 한다.

  3. useBoundStore 반환

와우 이렇게 간단 할 수가

이후 create 로 생성한 스토어를 호출 할 때 useSomethingStore(selector) 을 사용하게 되면 이는 사실 아직 말하지 않은 useStore(api ,selector) 를 호출하게 되는 것이다.

useStore

useStore
export function useStore<TState, StateSlice>(
  api: ReadonlyStoreApi<TState>,
  selector: (state: TState) => StateSlice = identity as any,
) {
  const slice = React.useSyncExternalStore(
    api.subscribe,
    () => selector(api.getState()),
    () => selector(api.getInitialState()),
  )
  React.useDebugValue(slice) // react devtools를 이용하기 위한 훅
  return slice
}

useStore 는 더 간단하다. 단순히 useSyncExternalStore 를 호출하는 것 외에 없다.

useSyncExternalStore 의 첫 번째 인수인 api.subscribe 가 호출 될 때 마다 리액트는 리렌더링 해야 될지 말지를 판단한다.

두 번째 인수인 getSnapshot 파트에서 api.subscribe 함수가 호출된 시점의 selector(api.getState())와 이전에 반환했던 상태인 selector(api.getState()) 의 비교를 통해 변경이 일어났다면 리렌더링을 유발한다.

useSyncExternalState 를 통해 받는 상태는 두 번째 인수의 실행 결과이다.

오 마이 갓, useSyncExternalState 를 이용했으니 Context 없이도 전역 상태 관리가 가능했고 선택적 구독이 가능했구나

오 마이 갓

그럼 이전 내 코드는 왜 에러가 발생했을까?

이전 내 코드에선 외부에서 생성된 상태를 변형하여 새로운 배열을 반환하는 map 메소드를 이용했었다.

리액트는 getSnapshot 의 이전 호출 결과와 현재 호출 결과가 다를 경우 리렌더링이 일어나는데 매번 렌더링 시점마다 새로운 배열이 생성되었기 때문이다.

회고

물론 zustand 의 동작 방식이 처음부터 이렇게 미니멀했던 것은 아니더라

실제 src/react 파트의 커밋 히스토리에서 가장 첫 커밋을 보면 다양한 커스텀 훅과 useRef 로 범벅이 된 코드를 볼 수 있다.

단순 zustand 는 리액트 api인 useSyncExternalState 를 활용한 라이브러리였구나싶다.

물론 이 안에서 상당히 고수스러운 타입 선언도 있는데 그건 이해하다가 포기했다.

고수스러운 타입 선언
type SetStateInternal<T> = {
  _(
    partial: T | Partial<T> | { _(state: T): T | Partial<T> }['_'],
    replace?: false,
  ): void
  _(state: T | { _(state: T): T }['_'], replace: true): void
}['_']

이 타입은 setState 의 타입 선언 방식인데 도저히 이해가 안간다. 🫠