Jotai 리렌더링 최적화, atom 분해가 먼저다

들어가며

Jotai structural sharing vs selectAtom (Peter Piekarczyk) 를 읽고 설명을 더해 정리한 글입니다.

저자는 같은 쿼리, 같은 응답, 같은 렌더 결과인데도 jotai-tanstack-query가 TanStack Query useQuery보다 44배 많은 리렌더링을 발생시키는 현상을 발견했다고 합니다. 왜 이런 일이 생기는지, 그리고 리렌더링 문제를 만났을 때 selectAtom을 사용하는게 왜 대부분 정답이 아닌지를 다뤄보려 합니다.

Jotai가 변경을 전파하는 방식

Jotai는 파생 atom의 값이 바뀌었는지를 Object.is로만 판단합니다. 얕은 비교나 구조적 비교를 하지 않고, 원시값은 값으로 비교하고 객체는 참조값을 비교해요. 이 때, 문제가 되는 건 읽기 함수가 매번 새 값을 조립하는 경우입니다.

const summaryAtom = atom((get) => ({
  count: get(cartAtom).items.length,
  total: sum(get(cartAtom).items),
}))

cartAtom이 갱신될 때마다 읽기 함수가 재실행되고, 객체 리터럴은 실행될 때마다 새 객체를 만듭니다. Object.is는 얕은 비교조차 하지 않기 때문에, count와 total이 이전과 같아도 새 객체는 무조건 “다르다”고 판정되는 겁니다. 그래서 구독자는 값이 안 바뀌었는데도 리렌더링돼요.

반대로 원시값이나 소스에 이미 존재하는 참조를 반환하는 파생 atom은 아무 문제가 없습니다.

const countAtom = atom((get) => get(cartAtom).items.length)
const userAtom = atom((get) => get(rawAtom).user)

읽기 함수가 재실행되는 건 똑같지만, 같은 숫자와 보존된 참조는 Object.is가 “같다”고 인식해서 알림이 멈춥니다. 다만 이 안전함은 rawAtom을 갱신하는 쪽이 { ...prev, settings: newSettings }처럼 user 참조를 보존해줄 때만 성립해요. 결국 재실행 결과를 같다고 인식할 수 있느냐가 갈림길인 거예요.

반사적인 처방: selectAtom

이 문제를 만나면 보통 selectAtom을 꺼내게 됩니다.

import { selectAtom } from "jotai/utils"

const summaryAtom = selectAtom(
  cartAtom,
  (cart) => ({ count: cart.items.length, total: sum(cart.items) }),
  shallowEqual,
)

selectAtom은 이전에 선택한 값을 기억해뒀다가, 원본 atom이 업데이트될 때마다 비교 함수를 실행해서 같으면 이전 참조를 그대로 반환합니다. 그러면 Object.is 검사가 통과해서 리렌더링을 방지하는 원리예요.

그런데 저자는 이게 대부분의 경우 올바른 해결책이 아니라고 말합니다.

selectAtom 대신 atom을 쪼개라

위의 summaryAtom은 애초에 이렇게 쪼개는 게 낫습니다.

const countAtom = atom((get) => get(cartAtom).items.length)
const totalAtom = atom((get) => sum(get(cartAtom).items))

리프가 원시값이 되는 순간 Object.is가 공짜로 중복 제거를 해주기 때문에, 동등성 함수도 selectAtom도 필요 없어져요. 저자가 제시하는 일반 규칙은 이겁니다.

원시값이나 소스에 이미 존재하는 참조를 반환한다면 일반 파생 atom을 쓰고, 새 래퍼 객체를 조립하고 있다면 여러 개의 원시 atom으로 쪼개라.

객체를 다루는 atom도 같은 원리로 필드별 파생 atom을 만들면 됩니다.

const userAtom = atom({ name: "", email: "", age: 0 })
const nameAtom = atom((get) => get(userAtom).name)
const emailAtom = atom((get) => get(userAtom).email)

selectAtom이 제값을 하는 경우

그렇다고 selectAtom이 쓸모없는 건 아니에요. 저자가 인정하는 정당한 사용처는 두 가지입니다.

첫 번째는 배열의 .filter, .map입니다.

const activeIdsAtom = selectAtom(
  usersAtom,
  (users) => users.filter((u) => u.active).map((u) => u.id),
  shallowEqual,
)

요소 하나하나가 참조 안정적이어도 filter와 map은 매번 새 배열을 만들기 때문에, atom을 쪼개는 걸로는 해결이 안 됩니다. 이럴 땐 얕은 비교가 유일한 선택지예요.

두 번째는 정체성이 중요한 집계입니다. 결과 객체를 메모이제이션된 자식 컴포넌트나 컨텍스트에 통째로 넘겨야 해서 참조 안정성 자체가 필요한 경우요. 이 둘이 아니라면 분해를 우선하라는 게 저자의 입장입니다.

selectAtom으로 못 푸는 문제: 원본 atom을 고쳐라

처음의 44배 문제로 돌아가볼게요. 서버 응답을 처음 담는 원본 atom이 업데이트될 때마다 새 객체를 내보내는 상황에서, 읽는 쪽마다 selectAtom으로 감싸면 어떻게 될까요? 읽는 곳이 N곳이면 같은 비교를 N번 반복하고, selectAtom 보일러플레이트도 사방에 깔립니다. 값을 읽는 곳마다 방어하는 대신, 값이 들어오는 입구에서 한 번에 고치는 게 좋아요.

여기서 나오는 개념이 structural sharing입니다. TanStack Query가 기본으로 켜두는 동작인데요. 새 데이터가 도착하면 이전 값과 비교하면서 트리를 순회하고, 내용이 같은 부분은 이전 객체를 그대로 재사용하고 실제로 바뀐 부분만 새 객체로 교체합니다. 그러면 이 atom을 읽는 쪽은 자기가 보는 부분의 참조가 안정적으로 유지되니까, 별다른 장치 없이 Object.is만으로 불필요한 리렌더링이 걸러져요.

TanStack Query 없이 직접 만들 수도 있습니다. TanStack Query 내부 유틸인 replaceEqualDeep을 가져와서 값을 쓰는 시점에 적용하면 돼요.

function atomWithStructuralSharing<T>(initial: T) {
  const inner = atom(initial)
  return atom(
    (get) => get(inner),
    (get, set, next: T) => set(inner, replaceEqualDeep(get(inner), next)),
  )
}

값이 들어올 때 한 번만 비용을 치르면, 그 아래의 모든 파생 atom이 공짜로 혜택을 봅니다. WebSocket으로 데이터를 받는 상황이라면 메시지를 atom에 넣는 지점에 적용하는 식이죠.

이제 남은 건 읽는 쪽입니다. structural sharing이 data의 참조를 지켜줘도, 원본 atom을 통째로 구독하면 래퍼 객체가 새로 만들어질 때마다 리렌더링되거든요.

const queryAtom = atom({ data: null, isLoading: false, updatedAt: 0 })

// 래퍼를 통째로 구독: 메타 정보만 바뀌어도 리렌더링
const result = useAtomValue(queryAtom)

// 실제로 읽는 필드만 파생해서 구독: data 참조가 바뀔 때만 리렌더링
const dataAtom = atom((get) => get(queryAtom).data)
const data = useAtomValue(dataAtom)

원본 atom을 useAtomValue로 직접 읽지 말고, 실제로 읽는 필드를 파생 atom으로 만들어서 구독하는 것. 이게 벤치마크의 44배를 1배로 만든 전부입니다.

Jotai가 이걸 기본 제공하지 않는 이유

“이렇게 좋은 거면 왜 Jotai에 내장 안 해줘?”라는 의문이 들 수 있는데요. 저자는 두 가지 이유를 듭니다.

먼저 atom은 의도적으로 최소한의 프리미티브입니다. 구조적 공유를 내장하려면 일반 객체, 클래스 인스턴스, Date 같은 값들을 어떻게 비교할지에 대한 입장을 라이브러리가 정해야 하고, 그 비용을 모든 소비자가 쓰기마다 치르게 됩니다.

그리고 올바른 답이 데이터 계층에 따라 완전히 달라집니다. 소스가 TanStack Query라면 이미 해주고 있으니 또 하면 낭비고, WebSocket이라면 수신 지점에서 replaceEqualDeep을 쓰는 게 맞습니다. 합성 가능한 프리미티브로 두는 쪽이 유연하다는 거죠.

마무리

제 나름대로 우선순위를 정리하면 이렇습니다. 리렌더링 문제를 만나면 selectAtom부터 꺼내는 대신, 먼저 파생 atom이 새 객체를 조립하고 있지 않은지 보고 원시값 단위로 쪼갭니다. 원본 atom이 매번 새 객체를 내보내는 게 문제라면 읽는 쪽이 아니라 값이 들어오는 입구에서 structural sharing으로 고칩니다. selectAtom은 filter/map 같은 파생 컬렉션과 참조 정체성이 필요한 집계, 딱 그 자리에만 남겨두고요. Object.is라는 단순한 규칙 위에서 상태 설계가 어떻게 달라지는지 보여주는 좋은 글이었습니다.