상태 관리 라이브러리 만들어보기
상태 관리 라이브러리 Zustand를 직접 구현하면서 동작 원리를 이해해 보자.
Zustand의 동작 방식
Zustand는 어떻게 동작할까? 간단히 말하면 Zustand가 상태 저장소를 만들고, React는 useSyncExternalStore를 통해 이를 구독하는 구조다.
Zustand가 상태 저장소 생성React가useSyncExternalStore로 상태 저장소 구독- 상태가 변경되면
Zustand가 모든 구독자에게 알림 React가 변경을 감지하고 컴포넌트를 리렌더링
기본 상태 저장소
먼저 최소한의 기능만 가진 단순한 저장소를 만들어보자.
세 가지 인터페이스를 갖는 저장소를 만든다.
- 값을 반환할 수 있다.
- 값을 변경할 수 있다.
- 값이 변경됐음을 알릴 수 있다.
type Store<T> = {getState: () => T;setState: (value: T) => void;subscribe: (listener: () => void) => () => void;};export const createStore = <T>(initialState: T) => {let state = initialState;const listeners = new Set<() => void>();const getState: Store<T>['getState'] = () => {return state;};const setState: Store<T>['setState'] = (value) => {const nextState = value;if (Object.is(nextState, state)) return;state = nextState;listeners.forEach((listener) => listener());};const subscribe: Store<T>['subscribe'] = (listener) => {listeners.add(listener);return () => listeners.delete(listener);};return { getState, setState, subscribe };};
createStore 함수는 세 가지 인터페이스를 반환한다. 내부의 state는 노출하지 않고 클로저를 이용해 안전하게 관리하는 구조다.
이제 이 저장소를 React가 구독할 수 있도록 만들어보자. useSyncExternalStore 명세에 따라 subscribe와 getSnapshot 함수를 연결한다.
export const useStore = <T>(store: Store<T>) => {const slice = useSyncExternalStore(store.subscribe, store.getState);return slice;};
저장소는 상태가 변경되면 React에 알리고, React는 새로운 상태 스냅샷을 Object.is로 비교해 변경 시 컴포넌트를 리렌더링한다.
벌써 기본 상태 저장소가 완성됐다. 이제 간단한 카운터 컴포넌트를 만들어서 테스트해 보자.
const store = createStore(0);const Counter = () => {const value = useStore(store);const increase = () => {store.setState(store.getState() + 1);};return (<><div>Value: {value}</div><button onClick={increase}>Increase Button</button></>);};
함수형 상태 업데이트
기존 저장소는 잘 동작하지만, 이전 상태를 기반으로 다음 상태를 계산해야 하는 경우가 자주 있다. setState(state + 1) 대신 setState((prev) => prev + 1)처럼 작성할 수 있으면 더 직관적이겠지?
setState 함수를 조금 확장해 보자.
const setState: Store<T>['setState'] = (value) => {const nextState = typeof value === 'function' ? (value as (state: T) => T)(state) : value;if (Object.is(nextState, state)) return;state = nextState;listeners.forEach((listener) => listener());};
const increase = () => {store.setState((prev) => prev + 1);};
효율적인 상태 저장소
기존 저장소는 상태를 변경할 때마다 모든 구독자에게 알린다. 원시 값일 땐 문제가 없지만 객체를 다룰 때는 불필요한 렌더링이 발생한다.
객체를 저장하고 일부만 변경해 보자. age가 변경되면 name만 구독하는 컴포넌트도 리렌더링된다.
const store = createStore({name: '김도현',age: 0,});const Name = () => {const { name } = useStore(store);return <div>Name: {name}</div>;};const Age = () => {const { age } = useStore(store);const increase = () => {store.setState((prev) => ({ ...prev, age: prev.age + 1 }));};return (<><div>Age: {age}</div><button onClick={increase}>Increase Button</button></>);};
React.memo로 감싸도 동일하게 리렌더링 된다. 리렌더링 트리거가 useSyncExternalStore에서 발생하고, 전체 객체를 구독하기 때문에 객체 참조가 새로워지면 모든 컴포넌트가 리렌더링되기 때문이다.
이 문제를 해결하기 위해 selector 개념을 도입해 보자. 핵심은 useSyncExternalStore가 getSnapshot 결과를 기억한다는 점이다. 따라서 무엇을 구독할지 만 조절해 주면 불필요한 리렌더링을 막을 수 있다.
type Selector<T, U = T> = (state: T) => U;export const useStore = <T, U = T>(store: Store<T>,selector: Selector<T, U> = (state) => state as unknown as U) => {const slice = useSyncExternalStore(store.subscribe,() => selector(store.getState()),() => selector(store.getState()));return slice;};
이제 컴포넌트에서 객체의 원하는 값만 골라서 구독할 수 있고, 불필요한 리렌더링도 발생하지 않는다.
const store = createStore({name: '김도현',age: 0,});const Name = () => {const name = useStore(store, (state) => state.name);return <div>Name: {name}</div>;};const Age = () => {const age = useStore(store, (state) => state.age);const increase = () => {store.setState((prev) => ({ ...prev, age: prev.age + 1 }));};return (<><div>Age: {age}</div><button onClick={increase}>Increase Button</button></>);};
최적화된 상태 저장소
selector를 추가하면서 이제 필요한 부분만 구독할 수 있게 되었다. 하지만 여전히 문제가 하나 남아있다.
여러 값을 객체로 묶어서 반환하면 객체 참조가 매번 새로 생성돼 무한 리렌더링이 발생한다.
useStore(store, (state) => ({ name: state.name, age: state.age }));
React는 참조 동등성(Object.is)으로 상태 변화를 감지하기 때문에 새로운 객체가 만들어질 때마다 useSyncExternalStore는 "값이 달라졌다"라고 판단한다.
이를 해결하기 위해 shallow 개념을 도입해 보자. 객체의 값이 바뀌지 않았다면 동일한 참조를 전달하면 된다.
const equal = <T>(value1: T, value2: T) => {if (JSON.stringify(value1) === JSON.stringify(value2)) return true;return false;};export const useShallow = <T, U = T>(selector: Selector<T, U>) => {const prev = useRef<U>(null);return (state: T) => {const next = selector(state);return equal(prev.current, next) ? (prev.current as U) : (prev.current = next);};};
이전값을 저장한 후 새 값과 비교해 동일하면 이전값을, 다르면 새 값을 반환하는 구조다.
이제 문제없이 객체의 여러 값을 한 번에 구독할 수 있다.
const store = createStore({name: '김도현',age: 0,});const Person = () => {const { name, age } = useStore(store,useShallow((state) => ({ name: state.name, age: state.age })));const increase = () => {store.setState((prev) => ({ ...prev, age: prev.age + 1 }));};return (<><div>Name: {name}</div><div>Age: {age}</div><button onClick={increase}>Increase Button</button></>);};
더 나아가기
이제 어느 정도 구색을 갖춘 상태 저장소가 되었다. 조금 더 편리하게 쓸 수 있도록 개선해 보자.
store와 useStore를 하나로 묶기
현재는 store와 useStore를 따로 불러와야 한다.
const store = createStore({ value: 0 });const value = useStore(store, (state) => state.value);
매번 이렇게 사용하는 건 조금 번거롭다. create 함수를 만들어서 store는 내부에서 생성하고 클로저만 반환하면 깔끔해진다.
export const create = <T>(initialState: T) => {const store = createStore(initialState);return <U = T>(selector?: Selector<T, U>) => useStore(store, selector);};
const useStore = create({name: '김도현',});const Person = () => {const name = useStore((state) => state.name);return <>{name}</>;};
점점 실제로 사용하는 Zustand와 비슷해지지?
컴포넌트 외부에서 접근하기
지금은 컴포넌트 내부에서 Hook으로만 사용할 수 있다. 컴포넌트 외부에서도 사용할 수 있으면 좋겠지?
사실 상태를 관리하는 코어 로직은 Vanilla로 이미 작성되어 있다. 이 인터페이스를 노출하기만 하면 되는데, 그렇다고 애써 없앤 store를 다시 불러오는 건 피하고 싶다.
반환 중인 클로저에 Vanilla 로직을 붙여서 하나의 함수를 Hook처럼 사용하면서 getState, setState, subscribe 같은 메서드까지 함께 노출하는 구조로 만들어보자.
export const create = <T>(initialState: T) => {const store = createStore(initialState);const useBoundStore = <U = T>(selector?: Selector<T, U>) => useStore(store, selector);Object.assign(useBoundStore, store);return useBoundStore as typeof useBoundStore & typeof store;};
const useStore = create({name: '김도현',});useStore.getState().name; // '김도현'
Getter와 Setter
끝으로 지금까지는 Hook을 통해서만 상태를 가져오거나 변경했다. 하지만 컴포넌트 내부에서 Hook이 들어갈 수 없는 위치에서도 상태를 변경해야 하거나, 같은 상태의 최신 값을 읽어야 하는 경우가 생길 수 있다.
상태 저장소를 생성할 때 내부에서 직접 상태를 읽고 쓸 수 있는 메서드를 제공해 보자.
type StateCreator<T> = (set: Store<T>['setState'], get: Store<T>['getState']) => T;const createStore = <T>(createState: StateCreator<T>) => {const listeners = new Set<() => void>();const getState: Store<T>['getState'] = () => {const setState: Store<T>['setState'] = (value) => {const subscribe: Store<T>['subscribe'] = (listener) => {let state = createState(setState, getState);return { getState, setState, subscribe };};...export const create = <T>(createState: StateCreator<T>) => {const store = createStore(createState);const useBoundStore = <U = T>(selector?: Selector<T, U>) => useStore(store, selector);Object.assign(useBoundStore, store);return useBoundStore as typeof useBoundStore & typeof store;};
저장소 생성 시 getter와 setter를 매개변수로 갖는 콜백 함수를 넘긴다.
const useStore = create((set, get) => ({name: '김도현',age: 0,increase: () => set((state) => ({ ...state, age: get().age + 1 })),}));
마치며
두렵게만 느껴졌던 라이브러리 내부가 생각보다 단순하고 명확한 원리로 동작한다는 걸 알게 되었다. 특별한 마법이 아니라 익숙한 개념들이 잘 조합되어 있었다.
단계마다 어떤 문제를 해결하려 했는지, 왜 이런 구조를 선택했는지 직접 구현하면서 이해해 보는 과정이 재미있었다. 실제로 사용하면서 당연하게 여겼던 기능들이 이런 고민 끝에 만들어진 거였다.