jotai 뜯어보기

개요

• atomic한 상태관리를 위한 상태관리 라이브러리 jotai 기본 동작원리에 대해 이해한다.

• recoil과 동일한 컨셉인 만큼 리코일과 비교하면서 이해 해보자

추측해보기

• 상태가 변하면 구독중인 곳을 모두 업데이트 해줘야 하므로, 옵저버 패턴을 쓸 거 같다.

• atom을 선언하는 건 옵저버 생성을 의미할 것이다.

• useAtom은 리액트를 위한 것으로, 상태 변경이 일어나면 리렌더링을 트리거 시키기 위한 용도다.
• React 에서는 UI 업데이트를 위해선 상태를 써야 한다.
• 따라서 내부적으로는 useEffect와 useState를 사용할 것이다.
• 옵저버의 listener에 setState를 시키는 로직을 등록할 것이다.

• atom 값이 바뀌면 등록된 listner를 모두 호출하게 되고, 그러면 각 컴포넌트 내부의 상태가 업데이트 되는 거니까 컴포넌트도 리렌더링 될 것이다.

주요 컨셉

Atom

TwitterDaishi Kato on Twitter / X

Daishi Kato on Twitter / X

Demystifying the internal of jotai https://t.co/16ydaf1Mxv, an atomic state management solution for @ReactJS.Here's a simplified version of the core. Is it readable? pic.twitter.com/TBpFDYNkeP— Daishi Kato (@dai_shi) January 22, 2022

Twitter

const atomState = {
  value: atom.init,
  listeners: new Set(),
}

derived Atom ( 리코일에서의 selector )

TwitterDaishi Kato on Twitter / X

Daishi Kato on Twitter / X

While nobody requested 😆, here's a longer version of the simplified jotai core. This version supports derived atoms. Hope someone enjoys it! https://t.co/gPLLSS5WCe pic.twitter.com/bpMLlyGAp6— Daishi Kato (@dai_shi) January 24, 2022

Twitter

• atom을 useAtom으로 상태로써 사용하고 있는 것은 listners에 포함되고,

• atom을 참조하고 있는 derived atom은 dependents에 포함된다.

const atomState = {
  value: atom.init,
  listeners: new Set(),
  dependents: new Set(),
}

// sudo code
notify = () => {
	dependents.forEach( 모두 실행 ) // for derived atom
	listners.forEach( 모두 실행 ) // for components that using atom as a state
}

• 버튼은 countAtom 이랑만 연동되어 있다.

• doubleAtom은 countAtom을 이용해서 파생된 derived atom이다.

• 버튼을 눌러 countAtom이 수정되면 dependents가 쭉 호출되며 notify 되고, 따라서 doubleAtom의 값도 수정되어 DobleCounter 컴포넌트도 업데이트 되는 것을 확인할 수 있다.

진짜 뜯어보기 (for Beginner)

궁금1. Store가 있겠지? (해쉬맵으로 관리되니까)

type Store = {
  get: <Value>(atom: Atom<Value>) => Value
  set: <Value, Args extends unknown[], Result>(
    atom: WritableAtom<Value, Args, Result>,
    ...args: Args
  ) => Result
  sub: (atom: AnyAtom, listener: () => void) => () => void
}

const buildStore = (
  // store state
  atomStateMap: AtomStateMap = new WeakMap(),
  mountedMap: MountedMap = new WeakMap(),
  invalidatedAtoms: InvalidatedAtoms = new WeakMap(),
  changedAtoms: ChangedAtoms = new Set(),
  mountCallbacks: Callbacks = new Set(),
  unmountCallbacks: Callbacks = new Set(),
  storeHooks: StoreHooks = {},
  // atom intercepters
  atomRead: AtomRead = (atom, ...params) => atom.read(...params),
  atomWrite: AtomWrite = (atom, ...params) => atom.write(...params),
  atomOnInit: AtomOnInit = (atom, store) => atom.unstable_onInit?.(store),
  atomOnMount: AtomOnMount = (atom, setAtom) => atom.onMount?.(setAtom),
  // building-block functions
  ...buildingBlockFunctions: Partial<
    [
      ensureAtomState: EnsureAtomState,
      flushCallbacks: FlushCallbacks,
      recomputeInvalidatedAtoms: RecomputeInvalidatedAtoms,
      readAtomState: ReadAtomState,
      invalidateDependents: InvalidateDependents,
      writeAtomState: WriteAtomState,
      mountDependencies: MountDependencies,
      mountAtom: MountAtom,
      unmountAtom: UnmountAtom,
    ]
  >
): Store

• atom 상태와 관련된 atomStateMap - 이게 우리가 생각한 상태를 관리하는 해쉬맵

• 컴포넌트 마운트 여부를 관리하는 mountedMap
• mount, unmount 와 연관된다!

• invalidatedAtoms 로 업데이트가 필요한 atom 관리

• 상태가 변화했기 떄문에 전파가 필요한 상태들은 changedAtoms

궁금. 메모리 누수 관련 WeakMap

메모리 누수란?

메모리 누수(Memory Leak) : 더 이상 필요하지 않은 메모리가 해제되지 않고 계속 남아있는 현상

WeakMap이란?

WeakMap은 객체를 키로 사용하는 Map의 한 종류로, 키로 사용된 객체가 가비지 컬렉션(GC)의 대상이 되면 자동으로 제거된다. 더 이상 참조되지 않는 객체는 WeakMap에서도 자동으로 사라지므로 메모리 누수를 방지할 수 있다.

• 관리하는 쪽 : weak reference

• 사용하는 쪽 : strong reference

어떤 게 자동으로 GC 되는 걸까?

// 이런 atom은 보통 앱이 종료될 때까지 GC되지 않음
const globalAtom = atom(0)

function Component() {
  // 컴포넌트가 unmount되면 이 atom은 GC 대상이 됨
  const dynamicAtom = atom(0)
  useAtom(dynamicAtom)
}

const atomFamily = atomFamily((param: number) => atom(param))

function Component() {
  // 더 이상 참조되지 않는 atom은 GC 대상이 될 수 있음
  const someAtom = atomFamily(1)
  
  // atomFamily는 자체 cleanup 메커니즘도 제공
  atomFamily.setShouldRemove((createdAt, param) => {
    return Date.now() - createdAt > 1000 * 60
  })
}

궁금. 해쉬맵이면 key가 있어야 하는데 어떻게 생성할까?

// atom 객체가 생성됨
// 이 시점에는 아직 WeakMap에 등록되지 않음
const countAtom = atom(0)

const ensureAtomState =
  buildingBlockFunctions[0] ||
  ((atom) => {
    if (import.meta.env?.MODE !== 'production' && !atom) {
      throw new Error('Atom is undefined or null')
    }
    let atomState = atomStateMap.get(atom)
    if (!atomState) {
      atomState = { d: new Map(), p: new Set(), n: 0 }
      atomStateMap.set(atom, atomState) // 상태 해쉬맵에 추가!
      atomOnInit?.(atom, store)
    }
    return atomState as never
  })

export interface Atom<Value> {
  toString: () => string
  read: Read<Value>
  unstable_is?(a: Atom<unknown>): boolean
  debugLabel?: string
  /**
   * To ONLY be used by Jotai libraries to mark atoms as private. Subject to change.
   * @private
   */
  debugPrivate?: boolean
  /**
   * Fires after atom is referenced by the store for the first time
   * This is still an experimental API and subject to change without notice.
   */
  unstable_onInit?: (store: Store) => void
}

궁금2. atom의 생명주기

type Mounted = {
  // 리스너 함수들을 뜻한다. - listners
  readonly l: Set<() => void>
  
	// 내가 참조하고 있는 atom들을 뜻한다. - dependencies
  readonly d: Set<AnyAtom>
  
	// 나를 참조하고 있는 atom들을 뜻한다. - dependents
  readonly t: Set<AnyAtom>
  
  // unmount 될 때 호출할 함수를 뜻한다.
  u?: () => void
}

• 이 atom을 사용하고 있는 것은 l(리스너)에 모이므로 값이 바뀌면 전파해줘야 하고,

• 이 atom을 참조하고 있는 것은 t(dependents)에 담기고 역시 전파 해줘야 한다.

const mountAtom =
  buildingBlockFunctions[7] ||
  ((atom) => {
  
	  // atom이 있는지 확인하고 없으면 생성해서 atomStateMap 해쉬맵에 넣음
    const atomState = ensureAtomState(atom)
    
    // 현재 마운트 되어 사용중인지 확인
    let mounted = mountedMap.get(atom)
    
    // 마운트 되지 않았다면 마운팅 처리하면서 mountedMap에 넣음.
    if (!mounted) {
      // recompute atom state
      readAtomState(atom)
      // mount dependencies first
      for (const a of atomState.d.keys()) {
        const aMounted = mountAtom(a)
        aMounted.t.add(atom)
      }
      // mount self
      mounted = {
        l: new Set(),
        d: new Set(atomState.d.keys()),
        t: new Set(),
      }
      mountedMap.set(atom, mounted)
      storeHooks.m?.(atom)
      
      if (isActuallyWritableAtom(atom)) {
       ...
    }
    return mounted
  })

const unmountAtom =
  buildingBlockFunctions[8] ||
  ((atom) => {
    const atomState = ensureAtomState(atom)
    let mounted = mountedMap.get(atom)
    
    // 마운트 되어있던 상태고, 상태를 쓰는 곳과, 상태를 참조하는 곳이 없다는 것이 확인되면
    // unmount 처리한다.
    if (
      mounted &&
      !mounted.l.size &&
      !Array.from(mounted.t).some((a) => mountedMap.get(a)?.d.has(atom))
    ) {
      // unmount 될 때 호출할 함수 등록한 게 있다면 flush 할 때 호출될 수 있도록 추가
      if (mounted.u) {
        unmountCallbacks.add(mounted.u)
      }

      // unmount self
      mounted = undefined
      mountedMap.delete(atom)
      storeHooks.u?.(atom)

      // unmount dependencies
      // 현재 atom이 참조하고 있던 것들도 다 해제 시킴
      for (const a of atomState.d.keys()) {
        const aMounted = unmountAtom(a)
        aMounted?.t.delete(atom)
      }
      return undefined
    }
    return mounted
  })

궁금4. 전파 - flushCallbacks

const flushCallbacks =
  buildingBlockFunctions[1] ||
  (() => {
    let hasError: true | undefined
    let error: unknown | undefined
    
    // 중간에 에러가 발생하더라도 나머지 작업을 다 처리한 후에 에러를 다루기 위한 catch 처리
    const call = (fn: () => void) => {
      try {
        fn()
      } catch (e) {
        if (!hasError) {
          hasError = true
          error = e
        }
      }
    }
    
    // changed, unmount, mount 중에 호출해야 하는 게 있다면 모두 호출
    // 콜백 호출로 인해 또 바뀌는 changed가 있을 수 있기 떄문에
    // do-while로 계속적으로 반복
    do {
      if (storeHooks.f) {
        call(storeHooks.f)
      }
      
      // callbacks에 호출해야 하는 콜백들 모두 모으기
      const callbacks = new Set<() => void>()
      const add = callbacks.add.bind(callbacks)
      changedAtoms.forEach((atom) => mountedMap.get(atom)?.l.forEach(add))
      changedAtoms.clear()
      unmountCallbacks.forEach(add)
      unmountCallbacks.clear()
      mountCallbacks.forEach(add)
      mountCallbacks.clear()
      
      // callbacks 모두 호출! => 전파!
      callbacks.forEach(call)
      if (changedAtoms.size) {
        recomputeInvalidatedAtoms()
      }
    } while (
      changedAtoms.size ||
      unmountCallbacks.size ||
      mountCallbacks.size
    )
    
    // 작업 처리 과정에서 에러가 있었다면 이제 throw 해서 다룰 수 있도록 하기
    if (hasError) {
      throw error
    }
  })

setter가 상태 변화를 트리거 하잖아요

  const store: Store = {
    get: (atom) => returnAtomValue(readAtomState(atom)),
    set: (atom, ...args) => {
      try {
        return writeAtomState(atom, ...args)
      } finally {
        recomputeInvalidatedAtoms() // 의존성 체크
        flushCallbacks() // 전파!
      }
    },
    sub: (atom, listener) => {
      const mounted = mountAtom(atom)
      const listeners = mounted.l
      listeners.add(listener)
      flushCallbacks() // 전파!
      return () => {
        listeners.delete(listener)
        unmountAtom(atom)
        flushCallbacks() // 전파!
      }
    },
  }

더 깊게! ( TBD )

• 비동기 처리

• Promise 처리 방식

• Suspense 통합

• 업데이트 최적화

• 불필요한 재계산 방지

• 의존성 그래프 최적화

• React 통합

• useAtom의 구현

• 리렌더링 최적화