zustand 라이브러리 코드 분석하기
Zustand는 React 상태 관리 라이브러리입니다. 그동안 Redux와 Recoil을 써왔는데, Redux는 스토어 하나를 만드는 데도 보일러플레이트가 많았고, Recoil은 더 이상 업데이트되지 않고 있었습니다. 그 대안을 찾다가 zustand로 넘어왔습니다.
zustand는 스토어를 만들고 훅으로 꺼내 쓰는 구조가 전부라 API 표면이 좁습니다. 미들웨어로 기능을 덧붙일 수 있고, 상태가 바뀌면 그 상태를 구독한 컴포넌트만 다시 렌더링합니다. 다만 여기까지는 사용하는 쪽에서 본 설명이고, 이 동작이 내부에서 어떻게 이뤄지는지는 써 보는 것만으로는 알기 어려웠습니다.
그래서 이 글에서는 zustand의 소스 코드를 직접 따라가며, 스토어가 상태를 저장하고 구독자에게 알리는 과정, 그리고 그 스토어를 React 훅으로 잇는 과정을 코드 레벨에서 정리해 보려고 합니다.
zustand로 상태 관리하기
우선 zustand 공식 문서를 보면 다음과 같은 예제 코드를 볼 수 있습니다.
import { create } from 'zustand';
const useStore = create((set) => ({
count: 1,
inc: () => set((state) => ({ count: state.count + 1 })),
}));
function Counter() {
const { count, inc } = useStore();
return (
<div>
<span>{count}</span>
<button onClick={inc}>one up</button>
</div>
);
}
zustand는 상태를 스토어에 저장합니다. create 함수에 저장할 상태와 그 상태를 다루는 액션을 함께 정의하면, 스토어를 구독하는 훅 useStore가 반환됩니다. 컴포넌트에서는 이 훅을 호출해 상태와 액션을 꺼내 씁니다.
위 예제에서 Counter는 useStore를 호출해 count와 inc를 구조 분해로 받아, 화면에 그리고 버튼에 연결합니다. 사용하는 쪽 코드는 이게 전부입니다. 이 좁은 표면 뒤에서 스토어가 어떻게 상태를 관리하는지 다음 절에서 코드로 확인해 보겠습니다.
zustand 핵심 개념 코드 분석
zustand는 ‘구독’이라는 개념으로 상태를 관리합니다. 컴포넌트가 스토어의 상태 변경을 구독해 두면, 상태가 바뀔 때 스토어가 구독자에게 변경을 알립니다.
코드를 한 줄씩 뜯어보기 전에, 스토어가 어떤 함수들로 이뤄지는지 먼저 확인해 보겠습니다.
createStore: 스토어 생성 상태와 상태 관리 API를 정의합니다.getState: 상태 읽기 현재 상태를 읽습니다.setState: 상태 변경 상태를 변경하며, 구독자에게 알립니다.subscribe: 구독 및 알림 상태 변경에 반응하며, 필요 시 구독을 해제합니다.destroy: 스토어 종료 모든 구독자를 제거하고 스토어를 초기화합니다.
여기서 가장 중요한 함수는 createStore와 subscribe입니다.
createStore는 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) => {
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;
};
차근차근 위에서부터 하나씩 살펴보겠습니다.
상태 저장
상태의 저장과 구독 관리를 위해 state와 listeners를 정의해줍니다.
let state: TState;
const listeners: Set<Listener> = new Set();
상태 업데이트
const setState: StoreApi<TState>['setState'] = (partial, replace) => {
const nextState =
typeof partial === 'function'
? (partial as (state: TState) => TState)(state)
: partial;
setState 함수는 인자로 partial과 replace 값을 받습니다.
- partial: 상태 또는 액션(함수)
- replace: 상태를 특정 상태(nextState)로 대체할 것인지의 여부
partial이 만약 함수라면 현재 상태를 입력으로 받아, 새로운 상태(nextState)를 반환합니다.
partial이 함수가 아닌 값일 경우 그대로 새로운 상태(nextState)로 사용됩니다.
현재 상태와 새로운 상태 비교
if (!Object.is(nextState, state)) {
const previousState = state;
Object.is로 현재 상태와 새로운 상태가 같은 값인지 판단합니다. 다른 값일 때만 블록 안으로 들어가, 갱신 전 상태를 previousState에 보관합니다. 같은 값이면 아무 일도 하지 않으니, 불필요한 구독자 호출을 여기서 걸러 냅니다.
현재 상태 업데이트
state =
(replace ?? (typeof nextState !== 'object' || nextState === null))
? (nextState as TState)
: Object.assign({}, state, nextState);
만약 replace 값이 true라면 현재 상태를 새로운 상태로 대체시킵니다.
false라면 Object.assign 연산자를 통해 현재 상태와 새로운 상태를 합칩니다.
등록된 모든 구독자 호출
listeners.forEach((listener) => listener(state, previousState));
}
이후 listeners를 순회하며, 모든 구독자에게 새로운 상태와 이전 상태를 전달합니다.
여기서의 구독자는 이후에 자세히 설명하겠지만, 리액트에서 상태 변경을 감지하는 콜백 함수를 등록합니다.
상태 반환
const getState: StoreApi<TState>['getState'] = () => state;
const getInitialState: StoreApi<TState>['getInitialState'] = () => initialState;
getState는 현재 상태를 반환하는 역할을 하고, getInitialState는 초기값을 반환하는 역할을 합니다.
구독자 등록 및 해제
const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
listeners.add(listener);
// Unsubscribe
return () => listeners.delete(listener);
};
subscribe함수는 호출되면 인자로 들어온 함수를 listeners에 추가합니다.
구독을 해제할 수 있는 함수를 반환해, 이를 호출하면 listeners에서 해당 함수를 제거시킵니다.
클로저를 통한 상태 유지
const api = { setState, getState, getInitialState, subscribe }
const initialState = (state = createState(setState, getState, api))
return api as any
}
createStoreImpl 함수가 반환하는 api는 setState, getState, subscribe 같은 함수들입니다. 이 함수들은 모두 createStoreImpl 안에 선언된 state와 listeners를 참조합니다. 클로저 덕분에 함수 실행이 끝난 뒤에도 이 두 변수는 사라지지 않고, 반환된 api 함수들이 계속 붙들어 둡니다. 스토어의 상태와 구독자 목록이 컴포넌트 바깥에서 유지되는 근거가 여기에 있습니다.
Closure(클로저): 어떤 함수가 다른 함수 내부에서 선언되었을 때, 그 함수가 외부 함수의 변수와 환경에 접근할 수 있는 기능
여기까지가 상태를 저장하고 업데이트하는 부분입니다. 이제 이 스토어를 React에서 쓰려면 훅으로 감싸 주는 코드가 필요합니다.
zustand 리액트 훅 코드 분석
전체 코드
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);
return slice;
}
const createImpl = <T>(createState: StateCreator<T, [], []>) => {
const api = createStore(createState);
const useBoundStore: any = (selector?: any) => useStore(api, selector);
Object.assign(useBoundStore, api);
return useBoundStore;
};
상태 저장 훅
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);
return slice;
}
React 18 이전에는 이 부분이 더 긴 코드로 작성되어 있었지만, React 18이 useSyncExternalStore를 지원하면서 짧아졌습니다.
React.useSyncExternalStore: external state의 변경사항을 관찰하고 있다가, tearing이 발생하지 않도록 상태 변경이 관찰되면 다시 렌더링을 시작합니다.
const snapshot = useSyncExternalStore(subscribe, getSnapshot, getServerSnapshot?)
useSyncExternalStore는 subscribe, getSnapshot, getServerSnapshot 세 인자를 받습니다. useStore가 넘긴 세 인자가 앞서 만든 스토어의 함수들과 그대로 맞물립니다.
api.subscribe→ 스토어의 구독 함수. React가 이 함수로 스토어에 콜백을 등록해 두고, 상태가 바뀌면 알림을 받습니다.() => selector(api.getState())→ 스냅샷 함수. 현재 상태에서selector가 고른 부분만 반환합니다.() => selector(api.getInitialState())→ 서버 렌더링용 초기 스냅샷.
여기서 selector가 성능의 핵심입니다. 컴포넌트가 상태 전체가 아니라 selector로 고른 일부만 구독하면, 그 부분이 바뀔 때만 다시 렌더링됩니다. 앞서 setState가 Object.is로 변경 여부를 거르고 listeners에 알린 흐름이, 여기서 컴포넌트의 리렌더링으로 이어집니다.
스토어와 훅 잇기
useStore는 api를 인자로 받는데, 정작 위 사용 예제에서는 useStore(api, selector)가 아니라 useStore()만 호출했습니다. 이 간극을 메우는 것이 create의 실제 구현인 createImpl입니다.
const createImpl = <T>(createState: StateCreator<T, [], []>) => {
const api = createStore(createState);
const useBoundStore: any = (selector?: any) => useStore(api, selector);
Object.assign(useBoundStore, api);
return useBoundStore;
};
createImpl은 먼저 createStore로 스토어의 api를 만듭니다. 그다음 이 api를 고정으로 물고 있는 훅 useBoundStore를 만드는데, 이 훅을 호출하면 내부에서 useStore(api, selector)가 실행됩니다. 덕분에 컴포넌트는 api를 넘기지 않고 useStore()만 호출하면 됩니다.
마지막 Object.assign(useBoundStore, api)는 훅 함수 자체에 api의 메서드를 붙입니다. 그래서 컴포넌트 밖에서도 useStore.getState()나 useStore.setState()처럼 훅을 거치지 않고 스토어를 직접 다룰 수 있습니다.
마무리
zustand의 코드는 크게 두 부분으로 나뉩니다. createStore는 클로저로 state와 listeners를 컴포넌트 바깥에 붙들어 두고, setState가 상태를 바꿀 때 Object.is로 실제 변경만 걸러 구독자에게 알립니다. 여기까지는 React와 무관한 순수한 스토어입니다.
그 스토어를 React에 잇는 것이 useStore이고, 다리 역할은 useSyncExternalStore가 합니다. 스토어의 subscribe와 getState를 그대로 넘겨 외부 상태를 구독하고, selector로 고른 부분이 바뀔 때만 컴포넌트를 다시 렌더링합니다.
라이브러리를 쓸 때는 create와 useStore 두 함수가 전부처럼 보이지만, 그 뒤에는 클로저로 상태를 붙들고 구독으로 렌더링을 잇는 구조가 있었습니다. 이렇게 원리를 정확히 알고 나니 올바른 활용법도 자연스럽게 따라왔습니다. selector로 필요한 상태만 구독해 리렌더링을 줄이는 것도, 그 동작을 코드로 확인한 뒤에는 근거를 가지고 쓰게 되었습니다.
참고
https://github.com/pmndrs/zustand