[Recoil]Jotai의 구현으로 알아보는 atom 종속성 추적

https://jotai.org/docs/guides/core-internals

Core Internals — Jotai

jotai의 마스코트

Jotai는 Recoil에서 영감을 받은 원자 모델을 사용하여 React 상태 관리에 대한 상향식 접근 방식을 취합니다.
원자를 결합하여 상태를 구축할 수 있으며 원자 종속성을 기반으로 렌더링이 최적화됩니다.
이는 React 컨텍스트의 추가 리렌더링 문제를 해결하고 메모이제이션 기술의 필요성을 제거합니다.

 

 

Jotai 핵심 구현의 완전한 예가 아니라 단순화된 버전입니다.

(Recoil의 코드베이스는 React와 연계하여 최적화 하기 위한 코드가 너무 많이 들어가 쉽게 분석할 수가 없습니다.)

아래의 짧은 코드로 핵심 구동 원리를 이해해 봅시다.

 

첫 번째 버전

쉬운 예부터 시작하겠습니다.

atom는 config 객체를 반환하는 함수일 뿐입니다.

실제로 사용하는 대상은 이를 통해 생성한 아톰입니다.

 

WeakMap을 사용하여 원자를 상태와 매핑합니다.

WeakMap은 키를 객체로 활용하는 자바스크립트의 Map 객체입니다.

WeakMap은 키를 메모리에 보관하지 않으므로 atom이 가비지 수집되면 상태도 가비지 수집됩니다.

이것은 메모리 누수를 방지하는 데 도움이 됩니다.

 

atomState의 listeners는 해당 아톰을 여러 컴포넌트에서 사용할 수 있게 해줍니다.

(하나의 아톰은 여러 컴포넌트에서 사용할 수 있는 공유 자원입니다.)

import { useState, useEffect } from "react";


// atom 함수는 초기 값을 포함하는 config 객체를 반환합니다.
export const atom = (initialValue) => ({ init: initialValue });


// 사용 시. 
const priceAtom = atom(10) // ({init : 10})

// atom의 상태 변화를 추적합니다.
// 메모리 누수를 피하기 위해 WeakMap을 사용합니다.
const atomStateMap = new WeakMap();
const getAtomState = (atom) => {
  let atomState = atomStateMap.get(atom);
  // 없으면 새로만듭니다.
  if (!atomState) {
    atomState = { value: atom.init, listeners: new Set() };
    atomStateMap.set(atom, atomState);
  }
  return atomState;
};

// useAtom hook은 [현재값,updator] 튜플을 반환합니다. 
export const useAtom = (atom) => {
  const atomState = getAtomState(atom);
  // useState를 활용합니다.
  const [value, setValue] = useState(atomState.value);
  useEffect(() => {
    const callback = () => setValue(atomState.value);   
    // 동일한 아톰을 여러 컴포넌트에서 사용할 수 있도록 mount시 리스너를 부착합니다.
    atomState.listeners.add(callback);
    callback();
    return () => atomState.listeners.delete(callback);
    // atomState 변화 시에만 변경을 알립니다.
  }, [atomState]);

  const setAtom = (nextValue) => {
    atomState.value = nextValue;
    // 모든 구독자 컴포넌트에 변경을 알립니다. (자기 자신을 포함하여...)
    atomState.listeners.forEach((l) => l());
  };

  return [value, setAtom];
};

API도 Recoil과 동일한 모습을 볼 수 있습니다.

 

 

두 번째 버전

Jotai에서는 파생 atom을 만들 수 있습니다. 파생 atom은 다른 atom에 의존하는 atom입니다.

Recoil의 Selector에 대응하는 개념입니다.

selector

 

// primitive atom
const priceAtom = atom(10)


// derived atom
const readOnlyAtom = atom((get) => get(priceAtom) * 2) // ({read : (get) => get(priceAtom) *  2, write : '사용시 에러'})

const writeOnlyAtom = atom(
  null, // it's a convention to pass `null` for the first argument
  (get, set, args) => {
    set(priceAtom, get(priceAtom) - args)
  }
) // ({read : '사용 시 에러', get : '전달한 함수'})

const readWriteAtom = atom(
  (get) => get(priceAtom) * 2,
  (get, set, newPrice) => {
    set(priceAtom, newPrice / 2)
    // you can set as many atoms as you want at the same time
  }
)

 // ({read :  '전달한 함수', get : '전달한 함수'})

 

모든 의존성을 추적하려면 atom 상태에 속성을 하나 더 추가해야 합니다.

원자 X가 원자 Y에 의존(depends on, X->Y)한다고 가정해 봅시다.

원자 Y를 업데이트하면 원자 X도 업데이트 해야 합니다.

이것을 종속성 추적(dependency tracking)이라고 합니다.

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

 

이제 종속 atom의 상태 업데이트를 처리할 수 있는 함수를 만들어야 합니다.

1. Atom

import { useState, useEffect } from "react";

export const atom = (read, write) => {
  // 파생 아톰 : 공식 문서에 오류가 있는 것 같아 살작 수정하였습니다.
  if (typeof read === "function" || typeof write === "function") {
    return { read, write };
  }
  // primitive atom. init 있음 {read,write}
  const config = {
    // init아 겂임
    init: read,

    // read 함수의 get은 atom 값을 읽는데 사용합니다.
    // reactive하며 읽기 종속성을 추적합니다.
    // get 함수를 인자로 받아 config를 넘기는데, config는 자기 자신입니다. (this)
    read: (get) => get(config),
    // write 함수의 get은 원자 값을 읽지만 종속성을 주적하지 않습니다.
    // set 함수는 atom 값을 쓰고, 해당 atom의 write 함수를 호출합니다.
    // 적당히 인자로 받아 config(자기 자신인 아톰)를 넘깁니다.

    write:
      write ||
      ((get, set, arg) => {
        if (typeof arg === "function") {
          set(config, arg(get(config)));
        } else {
          set(config, arg);
        }
      })
  };
  return config;
};

1. primitive 아톰일 경우 init에 초기값(read)이 들어옵니다.

read함수는 파생 아톰에 의해 전달됩니다만,

write는 default 함수가 있으며 따로 전달할 수도 있습니다.

const priceAtom = atom(10)

해당 값을 이용해 뒤에 나올 atomStateMap에 초기 상태를 저장합니다. (config.init)

2. 파생 상태일 경우 init(초기값)이 없으며, read, write 함수만 존재합니다.

파생 상태의 경우 recoil의 셀렉터와 다르게 read(get), write(set) 함수만 받습니다.

둘 중 하나만 있으면 read/write only 입니다.

const readOnlyAtom = atom((get) => get(priceAtom) * 2);

// it's a convention to pass `null` for the first argument
const writeOnlyAtom = atom(
  null, 
  (get, set, args) => {
    set(priceAtom, get(priceAtom) - args)
  }
)

const readWriteAtom = atom(
  (get) => get(priceAtom) * 2,
  (get, set, newPrice) => {
    set(priceAtom, newPrice / 2)
    // you can set as many atoms as you want at the same time
  }
)

2. atomStateMap

weakMap은 키를 객체로 받는 Map입니다.

키가 garbage collecting되는 것을 허용하며, 그에 따라 gc되는 메모리 효율적인 맵입니다.

키는 이전의 atom =(( read, write, init} ) 객체입니다.

// 첫번째 예제와 동일하지만 종속성(dependents)이 하나 추가되었습니다.
const atomStateMap = new WeakMap();

2.1 getAtomState

아톰(config)을 키로 상태를 조회하는 함수입니다.

const getAtomState = (atom) => {
  let atomState = atomStateMap.get(atom);
  if (!atomState) {
    atomState = {
      value: atom.init, // 위 config의 init입니다.
      listeners: new Set(),
      dependents: new Set()
    };
    atomStateMap.set(atom, atomState); // atom을 키로 atom State를 매핑합니다.
  }
  return atomState;
};

3. readAtom

primitive atom을 조회할 때까지 종속성 그래프를 따라 종속성을 추가하고, 아래로 내려가며 사용하는 primtive atom을 조회하는 함수입니다.

상태 종속성은 위에서 아래(primitive)로 향합니다.

// 파생한 원자가 아니면 원자의 원자값을 반환합니다.
// 파생한 원자면 부모의 원자값을 읽습니다.
// 그리고 현재 아톰을 부모의 의존성 집합에 추가합니다. (부모가 자식을 관리 - 부모가 파생상태일 경우)
// 이는 재귀적 과정입니다.
const readAtom = (atom) => {
  const atomState = getAtomState(atom);
  const get = (a) => {
  	// 파생한 원자가 아님
    if (a === atom) {
      return atomState.value;
    }
    // 다중파생
    // 부모의 상태값을 읽음
    const aState = getAtomState(a);
    // state는 atomStateMap에서 추적하며, atom은 직렬회되어 유일한 키로 동작합니다.
    // 부모에 의존성을 추가
    aState.dependents.add(atom); // XXX add only
    // 재귀적으로 부모에 대해 수행
    return readAtom(a); // XXX no caching
  };
  // 파생아톰에 전달한 함수를 실행함
  const value = atom.read(get);
  atomState.value = value;
  return value;
};

1. 먼저, atom.read(get)부분을 봅시다.

read(get)은 atom 내부의 함수였습니다. 파생 아톰(selector)인 경우 {?read, ?write} = config, 아니면 {init,write,read} = config 모양입니다.

// prmitive 아톰인 경우.
// !!!!!!!!!!!!!! config가 자기 자신임을 유의합니다. !!!!!!!!!!!!!!
read: (get) => get(config),

파생아톰의 경우 위와 모양이 다릅니다 예제를 보고 따라가봅시다.

2. get의 인자로 config(=== atom)이 전달됩니다. atom은 클로저입니다.

아래 함수를 봅시다.

const doubleAtom = atom((get) => get(countAtom) * 2);

즉 위의 함수는 read가 get=>get(countAtom) *2 입니다.

해당 함수에 get을 전달하면 get으로 전달되는 파라미터는 atom입니다. 이 경우는 primitive입니다.

const countAtom = atom(0); // prmitive atom

 

atom은 atomStateMap의 키입니다.

즉, (( ?read, ?write, init} ) 형태의 객체로, 불변입니다. (상태가 맵의 값, 키는 아톰으로 불변)

클로저 atom은 doubleAtom(파생아톰)이며, a는 conutAtom입니다. 둘은 다릅니다.

따라서 countAtom의 상태를 조회하고 (aState)

atom의 countAtom에 대한 종속성을 추가합니다. (부모(primitive)쪽에 자식을 추가).

그리고 countAtom에 대해 재귀적으로 reaadAtom을 수행합니다.

primitive atom은 1번에서 보았다시피 read함수의 인자로 자기 자신을 전달합니다. 따라서 a===atom으로 값을 리턴합니다.

 

// atom은 이전 컨텍스트의 atom
const get = (a) => {
  	// 파생한 원자가 아님
    if (a === atom) {
      return atomState.value;
    }
    // 다중파생
    // 부모의 상태값을 읽음
    const aState = getAtomState(a);
    // state는 atomStateMap에서 추적하며, atom은 직렬회되어 유일한 키로 동작합니다.
    // 부모에 의존성을 추가
    aState.dependents.add(atom); // XXX add only
    // 재귀적으로 부모에 대해 수행
    return readAtom(a); // XXX no caching
  };

4. notify

이제 어려운 코드들은 다 끝났습니다.

atomState가 수정(write)되면 모든 종속 아톰에 (재귀적으로) 알려야 합니다. 
이제 이 아톰에 종속된 모든 컴포넌트에 대해 콜백을 실행합니다.

  useEffect(() => {
    const callback = () => setValue(readAtom(atom)); // 아톰 값을 다시 읽어와서 값을 설정하도록 함
    const atomState = getAtomState(atom);
    atomState.listeners.add(callback);
    callback();
    return () => atomState.listeners.delete(callback);
  }, [atom]);

부모(primitive) 쪽에서 정보를 갖고 있던 것을 기억합시다.

즉 조회의 방향은 자식 > 부모, noti의 방향은 부모 > 자식입니다.

const notify = (atom) => {
  const atomState = getAtomState(atom);
  atomState.dependents.forEach((d) => {
    if (d !== atom) notify(d); // 자기 자신은 제외.
  });
  atomState.listeners.forEach((l) => l());
};

5. writeAtom

writeAtom은 필요한 매개변수로 atom.write를 호출하고  notify function을 트리거합니다.

write함수가 없으면 파생아톰은 쓰지 못합니다. (읽기전용)

primitive atom은 useAtom 훅을 이용해 쓸 수 있습니다.

primitive atom은 write함수를 전달받거나, default 함수를 사용합니다. (primitive 아톰 생성 시, 혹은 파생 아톰 생성 시 전달)

    write:
      write ||
      ((get, set, arg) => {
      	// arg가 값으로 전달되면 아톰과 arg 값을 set 함수에 전달합니다.
        // arg가 함수이면 아톰에 get함수를 호출하여 값을 얻어온 뒤에 arg 함수에 적용합니다.
        if (typeof arg === "function") {
          set(config, arg(get(config))); // config는 primitive atom
        } else {
          set(config, arg);
        }
      })

 

쓰기 전용 아톰을 봅시다. write 함수가 전달됩니다.

• set의 왼쪽 인자에 값을 쓸 대상 아톰을 지정합니다.
• get함수로 아톰을 감싸 값을 가져옵니다.
• set의 오른쪽 인자에 쓸 값을 전달합니다.

아래 아톰의 write는 인자로 전달된 값을 아톰에서 뺍니다.

const writeOnlyAtom = atom(
  null, 
  (get, set, args) => {
    set(priceAtom, get(priceAtom) - args)
  }
)

 

1. useAtom 훅에 의해 const callback = () => setValue(readAtom(atom))으로 종속성 관계가 이미 등록된 상태입니다.

(즉 쓰려면 먼저 읽어둬야 합니다)

따라서 getAtomState로 value를 읽어올 수 있습니다.

 

2. get함수는 항상 동일합니다. (no closure)

set 함수는 이전의 readOnly 안의 get와 마찬가지로 primitive atom을 찾을 때까지 writeAtom을 이용해 재귀합니다.

(여기에서 atom은 클로저입니다.)

primitve atom을 찾으면 값을 쓰고, 변경된 atom을 이용해 notify합니다.

const writeAtom = (atom, value) => {
  const atomState = getAtomState(atom);

  // 'a' is some atom from atomStateMap
  const get = (a) => {
    const aState = getAtomState(a);
    return aState.value;
  };

  // 'a'가 원자와 같으면 값을 업데이트하고 해당 원자에 알리고 반환 
  // 그렇지 않으면 'a'에 대해 writeAtom을 호출합니다(재귀적으로).
  const set = (a, v) => {
    if (a === atom) {
      atomState.value = v;
      notify(atom);
      return;
    }
    writeAtom(a, v);
  };
  // atom 객체의 함수를 호출함.
  atom.write(get, set, value);
};

파생 아톰은 단지 primitive 값을 조합해서 보여줄 뿐입니다.
수정 대상은 항상 primitive atom입니다.

6. useAtom

이제 다 끝났습니다!

// 이전과 동일합니다.
export const useAtom = (atom) => {
  const [value, setValue] = useState();
  useEffect(() => {
    const callback = () => setValue(readAtom(atom));
    const atomState = getAtomState(atom);
    atomState.listeners.add(callback);
    callback();
    return () => atomState.listeners.delete(callback);
  }, [atom]);
  const setAtom = (nextValue) => {
    writeAtom(atom, nextValue);
  };
  return [value, setAtom];
};

 

참고

해당 라이브러리 개발자는 zustand, valito를 포함한 수많은 오픈소스 라이브러리를 작성한 사람입니다.

최근에 자신이 개발한 상태관리 라이브러리들에 대한 책을 한권 냈더라구요. 읽어볼 시간이 있을런지 모르겠지만 기회가 되면 읽어볼 예정입니다.

https://www.amazon.com/Micro-State-Management-React-Hooks/dp/1801812373

 

                

상당히 쓰는데 힘이 들었던 글입니다. 맘에드셨으면 커피 한잔 부탁드립니다.

https://buymeacoffee.com/hyeoki 

hyeoki is creating sw