전역 상태 관리 Deep Dive: 바닥부터 구현하는 Zustand

✅ 개요

전역 상태 관리 라이브러리를 사용하면서 React 외부에 있는 스토어가 어떻게 React와 함께 활용할 수 있는지 궁금증이 있었습니다.

그래서 이를 가장 작은 규모부터 차근차근 구현하면서 어떻게 이루어져있는지 만들어보면서 어떠한 히스토리가 담겨져있는지 살펴보려고 합니다.

🔨 최소 모듈 스토어 구현하기

먼저 가장 단순하게 React 외부 스토어를 구현해보겠습니다.

1. 기본 구현

//✅ 상태 값과 상태 관리를 구독해주는 Set 함수
let count = 0;
const setStateFunctions = new Set<(count: number) => void>();

const Component1 = () => {
  const [state, setState] = useState(count);
  
  useEffect(() => {
    //✅ 해당 컴포넌트가 전역 상태를 구독해주는 코드
    setStateFunctions.add(setState);
    return () => {
      setStateFunctions.delete(setState);
    };
  }, []);

  const inc = () => {
    //✅ 외부 상태를 업데이트한 후, 구독 함수에 전파하였음
    count += 1;
    setStateFunctions.forEach((fn) => fn(count));
  };

  return (
    <div>
      {state} <button onClick={inc}>+1</button>
    </div>
  );
};

이 구현에서는 외부 모듈에서 count 값을 선언하고, setStateFunctions를 통해 상태 갱신 함수들을 관리합니다.
상태가 업데이트되면 모듈의 상태를 변경하고, 저장된 모든 setState 함수들에게 새로운 상태를 전달합니다.

2. 모듈 분리
실제 스토어는 컴포넌트와 분리해서 작성하므로 다음과 같이 분리해보았습니다.

// custom-store.ts
export let count = 0;
export const setStateFunctions = new Set<(count: number) => void>();

// CounterWithModule.tsx
import { count, setStateFunctions } from "./custom-store";

export function CounterWithModule() {
  const [state, setState] = useState(count);
  
  useEffect(() => {
    setStateFunctions.add(setState);
    return () => {
      setStateFunctions.delete(setState);
    };
  }, []);

  const inc = () => {
    //❌ 가져오기이므로 'count'에 할당할 수 없습니다.
    count += 1;
    setStateFunctions.forEach((fn) => fn(count));
  };

  return (
    <div>
      {state} <button onClick={inc}>+2</button>
    </div>
  );
}

하지만 import된 변수는 readonly이므로 문제가 발생했습니다.

3. 상태 변경 함수 추가하기

// custom-store.ts
export let count = 0;
export const setStateFunctions = new Set<(count: number) => void>();

//✅ 업데이트 함수 추가
export const increment = (amount: number) => {
  count += amount;
  return count;
};

// CounterWithModule.tsx
const inc = () => {
  //✅ 업데이트 함수로 변경하여 readonly 이슈 해결
  const count = increment(2);
  //❌ 구독한 컴포넌트에게 알림을 매번 선언해주어야 함.
  setStateFunctions.forEach((fn) => fn(count));
};

이 구현에서는 increment 함수로 readonly 문제를 해결하였습니다.

아쉬운 점은 상태 업데이트와 구독자 알림이 분리되어 있어 불편했습니다.
구독자들에게 상태 변경을 알리기 위해 반환값을 사용해야 했고, 매번 두 단계의 작업이 필요했습니다.

4. 상태 변경과 구독 로직 통합하기
이를 통합하여 개선해보았습니다.

// custom-store.ts
export let count = 0;
export const setStateFunctions = new Set<(count: number) => void>();
export const increment = (amount: number) => {
  count += amount;
  //✅ increment 함수로 이동
  setStateFunctions.forEach((fn) => fn(count));
};

// CounterWithModule.tsx
const inc = () => {
  //✅ 상태 업데이트와 구독자 알림이 하나의 함수로 처리됨
  increment(2);
};

상태 변경과 구독자 알림이 하나의 함수에서 처리하도록 개선하였습니다.
이제 컴포넌트에서는 단순히 상태 변경 함수만 호출하였고, 구독 알림을 추상화하였습니다.

아쉬운 점이 하나 더 있습니다.
매번 스토어를 만들 때마다 직접 입력해줘야합니다.
이를 개선하기위해 함수를 활용하여 필요한 객체를 생성할 수 있도록 만들어보겠습니다.

🏭 팩토리 함수로 개선하기

먼저 기존 모듈의 구조를 살펴보면:

// 실제 상태값
export let count = 0; 

// 구독 함수의 집합
export const setStateFunctions = new Set<(count: number) => void>(); 

// 상태 변경 함수
export const increment = (amount: number) => { 
  count += amount;
  setStateFunctions.forEach((fn) => fn(count));
};

여기서 각 요소의 역할은:

• count: 실제 상태값을 저장
• increment: 상태를 업데이트하는 setter 함수
• setStateFunctions: 관례적으로 'subscribe'라고 불리며, 상태 업데이트 시 등록된 함수들에게 알림

이러한 패턴을 재사용하기 쉽도록 팩토리 패턴으로 만들어보겠습니다:

const createStore = <T extends unknown>(initialState: T): Store<T> => {
  //✅ count를 대신하여 상태를 저장하는 변수
  let state = initialState;  
  
  //✅ setStateFunctions를 대신하여 구독 함수를 관리해주는 변수
  const callbacks = new Set<() => void>();  
  
  //✅ setter 함수
  const setState = (nextState: T) => {  
    state = nextState;
    callbacks.forEach((callback) => callback());
  };
  
  //✅ 구독 함수를 추가하고, 언마운트 시 해제해주는 함수
  const subscribe = (callback: () => void) => { 
    callbacks.add(callback);
    return () => {
      callbacks.delete(callback);
    };
  };
  
  return { state, setState, subscribe };
};

이렇게 하면 스토어가 필요할 때마다 쉽게 생성하여 사용할 수 있습니다.

하지만 이 구현에는 몇 가지 문제점이 있습니다.
차근차근 문제점들을 개선해보겠습니다.

🔒 state 캡슐화의 필요성

스토어를 생성할 때 state를 직접 노출하면 예상치 못한 문제가 발생할 수 있습니다.
가장 큰 문제는 상태 관리의 일관성이 깨질 수 있다는 점입니다.

//❌ state가 직접 노출하여 문제점 발생
const store = createStore({ count: 0 });
store.state = { count: 1 }; 

이렇게 직접 상태를 수정하면 상태 변경을 추적할 수 없습니다.

  //setter 함수
  const setState = (nextState: T) => {  
    state = nextState;
    callbacks.forEach((callback) => callback());
  };

setter 함수를 통해 구독 함수를 업데이트하는데 직접 수정하면 구독한 컴포넌트는 이를 알 수 없습니다.

이러한 문제를 해결하기 위해 상태를 캡슐화하고, 상태 접근을 위한 안전한 인터페이스를 제공해야 합니다.

👊 getState 도입하기

상태를 조회하는 함수인 getState를 추가해보겠습니다:

const createStore = <T extends unknown>(initialState: T): Store<T> => {
  let state = initialState;
  const callbacks = new Set<() => void>();
  
  //✅ 상태 읽기 전용 함수
  const getState = () => state;  
  
  const setState = (nextState: T) => {
    state = nextState;
  };
  
  const subscribe = (callback: () => void) => {
    callbacks.add(callback);
    return () => {
      callbacks.delete(callback);
    };
  };

  return { getState, setState, subscribe };
};

이렇게 개선한 스토어는

• 상태값을 직접 수정할 수 없습니다.
• getState()를 통해 안전하게 상태를 읽을 수 있습니다.
• 모든 상태 변경은 setState를 통해서만 가능합니다.

이를 통해 상태 관리의 안정성이 크게 향상됩니다.

이제는 다음 문제점인 setter 함수의 문제점을 살펴보겠습니다.

⚡ setter 함수 개선하기

먼저 기존의 setter 구현을 살펴보겠습니다:

// custom-store.ts
export const increment = (amount: number) => {
  count += amount;
  setStateFunctions.forEach((fn) => {
    fn(count);
  });
};

// CounterWithModule.tsx
const inc = () => {
  const count = increment(2);
};

이러한 단순한 값 교체 방식은 두 가지 중요한 문제가 발생할 수 있습니다.

1. 동시성 문제
여러 상태 업데이트가 연속해서 발생할 때 의도치 않은 결과가 발생할 수 있습니다:

setState(count + 1);  // count가 0일 때
setState(count + 1);  // 여전히 이전 count 값(0) 참조
setState(count + 1);  // 여전히 이전 count 값(0) 참조
//❌ 결과: count = 1 (예상과 다르게 한 번만 증가)

2. 클로저 문제
비동기 컨텍스트에서 오래된 상태값을 참조합니다.

setTimeout(() => {
  setState(count + 1);  //❌ 오래된 count 값 참조
}, 1000);

생성시점에 count 값을 참조하여 1초 뒤에 count 값과 일치하지 않아 문제가 발생합니다.

setter 개선해보기

React의 useState setter 패턴을 참고하여 이 문제들을 해결할 수 있습니다:

const setState = (nextState: T | ((prev: T) => T)) => {
  state = typeof nextState === "function"
    ? (nextState as (prev: T) => T)(state)
    : nextState;
  callbacks.forEach((callback) => callback());
};

이 개선된 구현의 장점:

• 값을 직접 전달하는 것이 가능하고, 이전 값을 활용한 함수를 사용할 수 있습니다.
• 함수를 실행하는 시점에 값을 평가하므로 최신 상태값을 참조할 수 있습니다.

문제 해결 예시

1. 동시성 문제 해결:

// 업데이트 함수를 사용하여 순차적 증가 보장
setState(prev => prev + 1);  // 0 -> 1
setState(prev => prev + 1);  // 1 -> 2
setState(prev => prev + 1);  // 2 -> 3
//✅ 결과: count = 3 (의도한 대로 세 번 증가)

2. 클로저 문제 해결:

setTimeout(() => {
  setState(prev => prev + 1);  //✅ 항상 최신 상태값 기준으로 업데이트
}, 1000);

이렇게 개선된 setter 함수로 동시성과 클로저 관련 문제들을 효과적으로 해결할 수 있습니다.

✨ createStore 완성

위 문제점들을 해결하여 최종적으로 createStore가 만들어졌습니다.
이를 통해 안전하게 모듈 기반의 전역상태관리 스토어를 만들 수 있습니다.

const createStore = <T extends unknown>(initialState: T): Store<T> => {
  let state = initialState;
  const callbacks = new Set<() => void>();
  
  const getState = () => state;
  
  const setState = (nextState: T | ((prev: T) => T)) => {
    state = typeof nextState === "function"
      ? (nextState as (prev: T) => T)(state)
      : nextState;
    callbacks.forEach((callback) => callback());
  };
  
  const subscribe = (callback: () => void) => {
    callbacks.add(callback);
    return () => callbacks.delete(callback);
  };
  
  return { getState, setState, subscribe };
};

// 카운터 스토어
const counterStore = createStore<CounterState>({ count: 0 });

// 사용자 정보 스토어
const userStore = createStore<UserState>({
  name: "",
  age: 0,
  isLoggedIn: false
});

이제 스토어를 만들었으니 React에서 이를 구독하는 로직을 훅으로 만들어서 보일러플레이트를 줄여보겠습니다.

🎣 useStore 커스텀 훅

React에서 스토어를 구독하기 위해 아래와 같은 로직이 필요합니다.

// 스토어 상태 로직
const [state, setState] = useState(store.getState());

// 스토어 구독 로직
useEffect(() => {
  const unsubscribe = store.subscribe(() => {
    setState(store.getState());
  });
  setState(store.getState());
  return unsubscribe;
}, [store]);

이 로직을 재사용 가능한 훅으로 만들어 보겠습니다.

const useStore = <T extends unknown>(store: Store<T>) => {
  const [state, setState] = useState(store.getState());
  
  useEffect(() => {
    const unsubscribe = store.subscribe(() => {
      setState(store.getState());
    });
    setState(store.getState());
    return unsubscribe;
  }, [store]);
  
  return [state, store.setState] as const;
};

스토어를 주입 받아서, 각 스토어에 맞추어 사용할 수 있도록 만들었습니다.

이제 컴포넌트에서 적은 코드로 전역 상태를 사용할 수 있습니다:

//counter-store.ts
const counterStore = createStore<CounterState>({ count: 0 });

//counter.tsx
const Component = () => {
  
  const [state, setState] = useStore(counterStore);
  
  const increment = () => {
    setState(prev => ({
      ...prev,
      count: prev.count + 1
    }));
  };
  
  return (
    <div>
      <p>Count: {state.count}</p>
      <button onClick={increment}>증가</button>
    </div>
  );
};

이렇게 만든 커스텀 훅은

• 스토어 상태 구독을 자동으로 처리
• 컴포넌트 언마운트 시 구독 해제
• useState와 유사한 친숙한 API 제공

로 더 편리한 상태 관리가 가능해졌습니다.

🎯 Top Down Store의 selector 기능

모듈 상태 관리에서는 하나의 큰 객체로 저장합니다.
그래서 전역 상태가 업데이트가 되었을 때, 참조가 변경되어 구독하는 모든 컴포넌트가 렌더링됩니다.
문제는 실제 사용하는 값이 아님에도 렌더링하는 경우가 생깁니다.

이를 해결하기 위해 selector 기능을 도입하려고 합니다.

객체를 커다란 햄버거에 비유해볼 수 있습니다.
각 컴포넌트가 항상 전체 햄버거(상태)를 필요로 하지는 않습니다.

어떤 컴포넌트는 빵만, 또 어떤 컴포넌트는 치즈만 필요할 수 있죠.
이처럼 컴포넌트의 필요에 따라 특정 상태만 구독하고 싶을 때 selector를 활용할 수 있습니다.

🔄 useStoreSelector 구현하기

useStore에 selector 기능을 추가하여 useStoreSelector를 구현해보겠습니다.

getState 값에 selector를 감싸서 필요한 값만 읽어보겠습니다.

const useStoreSelector = <T, S>(store: Store<T>, selector: (state: T) => S) => {
  //✅ 초기 값을 특정 값만 구독하여 사용함
  const [state, setState] = useState(() => selector(store.getState()));
  useEffect(() => {
    const unsubscribe = store.subscribe(() => {
      setState(selector(store.getState()))
    });
    setState(selector(store.getState()));
    return unsubscribe;
  }, [store, selector]);
  return state;
};

이렇게 구현하면 selector를 주입받아 원하는 상태만 구독하고 사용할 수 있습니다.
각 스토어마다 주입 받은 callback을 저장하기 때문에

//특정 값 사용
setState(selector(store.getState()))

//전체 값 사용
setState(store.getState())

각 컴포넌트는 자신만의 콜백을 등록하여 사용합니다.
이를 통해 컴포넌트가 사용하는 값만 사용할 수 있습니다.

useStoreSelctor 실제 사용 예시

다음과 같이 selector를 정의하고 사용할 수 있습니다:

//custom-store.ts
const store = createStore({ count1: 0, count2: 0 });

//counter2.tsx
const selectCount2 = (state: ReturnType<typeof store.getState>) => state.count2;

const Component2 = () => {
  const state = useStoreSelector(store, selectCount2);
  const inc = () => {
    store.setState((prev) => ({
      ...prev,
      count2: prev.count2 + 1,
    }));
  };
  
  return (
    <div>
      count2: {state} <button onClick={inc}>+1</button>
    </div>
  );
};

count1과 count2를 key로 갖지만, selector를 통해 count2 값만 구독할 수 있습니다.
이를 통해 복잡한 중첩 객체에서도 필요한 값만 구독할 수 있습니다.

다음으로는 useSyncExternalStore라는 React 훅을 사용하여 구독하는 로직을 줄여보는 작업으로 마무리해보겠습니다.

🔗 useSyncExternalStore 활용하기

React 18.2에서 도입된 useSyncExternalStore를 활용하면 구독 로직을 더욱 간단하게 구현할 수 있습니다:

// 이전 구현
const useStoreSelector = <T, S>(store: Store<T>, selector: (state: T) => S) => {
  const [state, setState] = useState(() => selector(store.getState()));
  useEffect(() => {
    const unsubscribe = store.subscribe(() => {
      setState(selector(store.getState()));
    });
    setState(selector(store.getState()));
    return unsubscribe;
  }, [store, selector]);
  return state;
};

// useSyncExternalStore 활용
const useStoreSelector = <T, S>(store: Store<T>, selector: (state: T) => S) => {
  return useSyncExternalStore(
    store.subscribe,
    useCallback(() => selector(store.getState()), [store, selector])
  );
};

이를 통해 보일러플레이트 코드를 줄이고, React와 외부 스토어를 더욱 간편하게 연동할 수 있습니다.

🔗 zustand 로직 살펴보기

실제 zustand 로직을 살펴보면 유사한 점을 살펴볼 수 있습니다.

const createStoreImpl: CreateStoreImpl = (createState) => {
  //✅ 기존 state와 callbacks과 동일
  let state: TState
  const listeners: Set<Listener> = new Set()
  
  //✅ setState 구현
  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
      // replace가 true이거나 nextState가 객체가 아닌 경우 전체 교체
      // 그렇지 않으면 기존 상태와 병합
      state = (replace ?? (typeof nextState !== 'object' || nextState === null))
        ? (nextState as TState)
        : Object.assign({}, state, nextState)
      
      // 리스너들에게 상태 변경 알림
      listeners.forEach((listener) => listener(state, previousState))
    }
  }
  
  //✅ subscribe 동일한 구조
    const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
    listeners.add(listener)
    // Unsubscribe
    return () => listeners.delete(listener)
  }

  
  // ... 나머지 API 구현
}

setState 업데이트를 더 효율적으로하는 것 외에 구조는 동일했습니다.
zustand는 위에 만든 스토어보다 더 많은 기능을 제공하지만 기본 구조는 동일합니다.

🔚 끝으로

이렇게 간단한 모듈 스토어를 직접 구현해보았습니다.

평소 React에서 외부 스토어를 어떻게 구독하고 활용하는지 궁금했었는데, 직접 구현해보면서 이러한 궁금증을 해소할 수 있었습니다.

처음에는 zustand 로직을 보았을 때는 잘 읽히지가 않았는데, 직접 구현하고 나서 구도가 보이기 시작해서 잘읽을 수 있었습니다.

직접 하나하나 구현하면서 단순히 전역 상태 관리의 기능을 아는 것을 넘어서, 그 뒤에 담긴 히스토리와 설계 의도를 더 깊이 이해할 수 있는 좋은 경험이었습니다.

앞으로 제가 라이브러리를 개발하게 된다면, 이번에 고민했던 setter 함수와 getState 구현 경험을 토대로 더 나은 라이브러리를 만들 수 있을 것 같습니다.
긴 글 읽어주셔서 감사합니다.

참고문헌:

• 서적: 「리액트 훅을 이용한 마이크로프론트엔드」
• 소스코드: Zustand 라이브러리 구현 (vanilla.ts)
  https://github.com/pmndrs/zustand/blob/main/src/vanilla.ts#L59