jotai 내부 구조 훑어보기

회사에서 복잡한 react 컴포넌트(테이블)를 만들 일이 생겨서

c++의 namespace 기능처럼 provider 내부에서만 전역이고 밖에서는 사용 불가능한 context 기반의 라이브러리를 찾고 있었다.

 

그래서 찾은것중 하나가 jotai였고, 용량이 작고 꾸준히 업데이트 되는 점이 마음에 들어서 jotai를 사내 라이브러리 용으로 

쓰다보니 다음과 같은 여러가지 의문이 생겨서 코드를 훑어보기로 했다.

 

바로, 

 

1. context api는 상태가 하나 바뀌면 provider 하위 컴포넌트들이 전부 rerendering 되는데 jotai는 해당 문제를 해결하기 위해서 어떻게 구현했는지?

2. react에서 리렌더링이 되지 않고 현 상태값을 어떻게 알고 있는지??

 

2번 의문이 뭔지 구체적으로 설명하기 위해서 우선 jotai에 대해 설명해보자면,

jotai 상태를 재정의할때 쓰는 write 전용 useSetAtom이라는 hook 혹은 wrtieOnlyAtom이라는 atom이 있다.

const [state, setState] = useState의 setState의 jotai 버젼이라고 보면 된다. 

 

App.js
import { atom, useAtom } from "jotai";

const dotsAtom = atom([]);
const drawingAtom = atom(false);

const handleMouseDownAtom = atom(
  null,
  (get, set) => { // get(drawingAtom) 을 쓰면 현 상태 값을 불러올 수 있다! 
    set(drawingAtom, true);
  }
);

const handleMouseUpAtom = atom(null, (get, set) => {
  set(drawingAtom, false);
});

(https://tutorial.jotai.org/quick-start/write-only-atoms에서 발췌)

 

이 기능들을 쓰기 위해선 "현 상태"를 알아야할 때도 있을텐데 (주석 부분 참고)

writeOnlyAtom만을 쓰면 그 atom을 쓰는 컴포넌트는 상태가 바뀌어도 리렌더링이 되지 않는다.

react에서 리렌더링이 되지 않고 현 상태값을 어떻게 알고 있는지??

 

예를 들어 숫자를 +1, -1 하는 counter 컴포넌트를 만든다고 했을때

+1을 하기 위해 지금 상태 값을 알고, 그 값에 +1 혹은 -1를 해줘야할텐데

react에서 리렌더링을 유발하지 않고 어떻게 지금 상태값을 알고 있는건지???

 

두가지 의문점이 생겨서 내부 코드를 뜯어보기로 했다. (특히 2번)

깊게 보지는 않고, 위 2가지 의문점을 해결하는 것 중심으로 살펴볼 것이다.

 

 

우선 파일구조부터 보자.

https://github.com/pmndrs/jotai/tree/main/src

 

구조를 보니 vanila 폴더에 기본 로직이 있고 react 폴더에 react와 vanila를 이어주는 부분이 있을 것 같다.

 우선 react 부분부터 보자..

 

https://github.com/pmndrs/jotai/tree/main/src/react

 

생각보다 심플한 구조다. 

우선 provider를 봐보자 

 

provider.ts

export const useStore = (options?: Options): Store => {
  const store = useContext(StoreContext)
  return options?.store || store || getDefaultStore()
}

export const Provider = ({
  children,
  store,
}: {
  children?: ReactNode
  store?: Store
}): FunctionComponentElement<{ value: Store | undefined }> => {
  const storeRef = useRef<Store>()
  if (!store && !storeRef.current) {
    storeRef.current = createStore()
  }
  return createElement(
    StoreContext.Provider,
    {
      value: store || storeRef.current,
    },
    children,
  )
}

https://github.com/pmndrs/jotai/blob/main/src/react/Provider.ts

 

react 폴더 안에 provider.ts란 파일이 있다. 

 

useRef 안에 아까 본 "vanila 로직이 든 객체"를 넣고, 그 ref를 provider에 넣어주고 있다.

이럼 ref값이 바뀌어도, 상태가 아니니 react에서 리렌더링이 발생하지 않게 되는데

어디선가 리렌더링을 강제로 발생시키는 로직을 쓰고 있을것 같다.

 

추가적으로 useStore에서 getDefaultStore라는 값을 쓰고 있는데

provider가 명시가 안되면 전역에 한 createStore 객체를 만들고,

그 객체를 계속 참조하는것 같다.

 

useState의 jotai판인 useAtom을 봐보자.

 

useAtom.ts

export function useAtom<Value, Args extends any[], Result>(
  atom: Atom<Value> | WritableAtom<Value, Args, Result>,
  options?: Options,
) {
  return [
    useAtomValue(atom, options),
    // We do wrong type assertion here, which results in throwing an error.
    useSetAtom(atom as WritableAtom<Value, Args, Result>, options),
  ]
}

https://github.com/pmndrs/jotai/blob/main/src/react/useAtom.ts

 

useAtom이 뭔지 간단히 설명하자면, 말 그대로 useState인데 jotai atom을 참조해서 사용하는

useState의 전역 버젼이다. 

const [state, setState] = useAtom(jotaiAtom)

 

이런식으로, react를 알고 있다면 러닝커브가 낮아 쉽게 사용 가능해서 jotai를 고른 이유도 있다.

useAtom은 useAtomValue와 useSetAtom을 반환하는데 각각, read/write 기능이다. 

useAtomValue에 우리가 궁금했던 1번 로직이 담겨있을 것 같다.

 

보고 싶은 부분만 추려보자..

 

useAtomValue.ts

export function useAtomValue<Value>(atom: Atom<Value>, options?: Options) {
  const store = useStore(options)

  const [[valueFromReducer, storeFromReducer, atomFromReducer], rerender] =
    useReducer<
      ReducerWithoutAction<readonly [Value, Store, typeof atom]>,
      undefined
    >(
      (prev) => {
        const nextValue = store.get(atom)
        if (
          Object.is(prev[0], nextValue) &&
          prev[1] === store &&
          prev[2] === atom
        ) {
          return prev
        }
        return [nextValue, store, atom]
      },
      undefined,
      () => [store.get(atom), store, atom],
    )

  let value = valueFromReducer
  if (storeFromReducer !== store || atomFromReducer !== atom) {
    rerender()
    value = store.get(atom)
  }

  const delay = options?.delay
  useEffect(() => {
    const unsub = store.sub(atom, () => {
      if (typeof delay === 'number') {
        // delay rerendering to wait a promise possibly to resolve
        setTimeout(rerender, delay)
        return
      }
      rerender()
    })
    rerender()
    return unsub
  }, [store, atom, delay])

  useDebugValue(value)
  // TS doesn't allow using `use` always.
  // The use of isPromiseLike is to be consistent with `use` type.
  // `instanceof Promise` actually works fine in this case.
  return isPromiseLike(value) ? use(value) : (value as Awaited<Value>)
}

 

 

핵심 부분은 useEffect 부분의, 

  useEffect(() => {
    const unsub = store.sub(atom, () => {
      if (typeof delay === 'number') {
        // delay rerendering to wait a promise possibly to resolve
        setTimeout(rerender, delay)
        return
      }
      rerender()
    })
    rerender()
    return unsub
  }, [store, atom, delay])

 

sub/pub 부분인것 같다. 이 코드를 보고 나니 아까 품었던 의문 1,2가 풀린것 같다. 

jotai는 부분 렌더링을 지원하기 위해서 Publish-subscribe pattern(https://en.wikipedia.org/wiki/Publish%E2%80%93subscribe_pattern)을 사용하고 있었다.

 

어떤 "atom"을 component에서 사용한다면 해당 컴포넌트를 subscribe를 해주고 

atom의 값이 바뀌면 이벤트를 발행해서 해당 컴포넌트를 모두 rerendering 시키고 있는 간단한 형식이었다.

 

그리고, atom의 값은 완전한 상태가 아니고 store라는 객체가 가지고 있는 값이여서 아까 품었던 2번 의문(writeOnlyAtom에서 리렌더링 하지 않고도 최신 값을 가져올 수 있었던 방법)도 해결 되었다.

 

useAtom, 정확히는 useAtomValue를 쓰지 않는 컴포넌트는 리렌더링 되지 않았던 것..

 

생각보다 코드가 단순했던것 같다.

react에서 외부 라이브러리와 리엑트 상태를 동기화시키고 tearing을 방지 시키기 위해서 usesyncexternalstore? 를 쓰는 경우도 있다고 들었는데 jotai는 useEffect로 동기화 시키고 있었다. 

 

그외에 바닐라 폴더에서 비동기/동기 처리를 위해

여러 흥미로운 로직을 쓰고 있었는데 이건 다음 시간에 살펴봐야겠다.