Jotai 深入

原子化状态管理：基础使用、派生 Atom、异步集成、核心原理与源码解析

概述

Jotai（日语"状态"）是由 Pmndrs 团队开发的原子化状态管理库。灵感来自 Recoil，但 API 更简洁、无需 key 字符串。核心思想是：状态是一个个独立的原子（atom），组件按需订阅，派生自动追踪依赖。

---

一、优势与劣势

优势

• 原子化模型：每个状态独立定义，组件只订阅用到的 atom，天然细粒度重渲染
• 原生派生状态：atom((get) => ...) 自动追踪依赖，类似 Vue computed / MobX derivation
• 无 Provider（v2+）：默认 store 自动创建，无需包裹 Provider
• Suspense 原生集成：异步 atom 自动配合 Suspense，无需手动管理 loading 状态
• 极小体积：核心约 2KB（gzip），按需引入扩展
• 组合灵活：atom 可以自由组合、拆分、嵌套，适合复杂派生关系
• 丰富的官方扩展：jotai/utils（atomWithStorage、atomWithDefault 等）、与 TanStack Query / Immer / TRPC 的集成

劣势

• 概念门槛：atom / derived atom / writable derived atom 需要理解函数式组合思维
• 调试体验：atom 是匿名函数，DevTools 中不如 Redux 直观（可用 debugLabel 改善）
• 分散定义：大型应用中 atom 散落各处，缺少 Redux slice / Zustand store 那样的中心化组织
• 异步复杂场景：多个异步 atom 相互依赖时，Suspense 边界和错误处理需要仔细设计

---

二、基础使用

安装

npm install jotai

原始 Atom

import { atom, useAtom, useAtomValue, useSetAtom } from 'jotai';

// 原始 atom（类似 useState 的全局版）
const countAtom = atom(0);
const nameAtom = atom('');
const todosAtom = atom([]);

// 在组件中使用
function Counter() {
  const [count, setCount] = useAtom(countAtom);      // 读 + 写
  const name = useAtomValue(nameAtom);                 // 只读（不会因为 setName 重渲染）
  const setName = useSetAtom(nameAtom);                // 只写（不会因为 name 变化重渲染）

  return (
    <div>
      <p>计数: {count}</p>
      <button onClick={() => setCount(c => c + 1)}>+1</button>
      <input value={name} onChange={e => setName(e.target.value)} />
    </div>
  );
}

派生 Atom（只读）

// 自动追踪依赖：countAtom 变化时 doubleCountAtom 自动重新计算
const doubleCountAtom = atom((get) => get(countAtom) * 2);

const todoStatsAtom = atom((get) => {
  const todos = get(todosAtom);
  return {
    total: todos.length,
    done: todos.filter(t => t.done).length,
    pending: todos.filter(t => !t.done).length,
  };
});

function Stats() {
  const stats = useAtomValue(todoStatsAtom);
  // 只有 todosAtom 变化时才重渲染，countAtom 变化不影响
  return <p>总计 {stats.total}，已完成 {stats.done}</p>;
}

可写派生 Atom

const uppercaseNameAtom = atom(
  (get) => get(nameAtom).toUpperCase(),           // 读取（派生）
  (get, set, newName: string) => set(nameAtom, newName) // 写入（代理到原始 atom）
);

// 复杂写入逻辑
const addTodoAtom = atom(
  null, // 只写 atom，读取值为 null
  (get, set, text: string) => {
    const todos = get(todosAtom);
    set(todosAtom, [...todos, { id: Date.now(), text, done: false }]);
  }
);

function AddTodo() {
  const addTodo = useSetAtom(addTodoAtom);
  const [text, setText] = useState('');

  return (
    <form onSubmit={e => { e.preventDefault(); addTodo(text); setText(''); }}>
      <input value={text} onChange={e => setText(e.target.value)} />
      <button type="submit">添加</button>
    </form>
  );
}

异步 Atom

// 异步读取 — 自动配合 Suspense
const userAtom = atom(async () => {
  const res = await fetch('/api/user');
  return res.json();
});

// 依赖其他 atom 的异步读取
const userPostsAtom = atom(async (get) => {
  const user = await get(userAtom); // 等待 userAtom 解析
  const res = await fetch(`/api/users/${user.id}/posts`);
  return res.json();
});

// 使用（必须包裹在 Suspense 中）
function UserInfo() {
  const user = useAtomValue(userAtom); // 自动 suspend
  return <span>{user.name}</span>;
}

<Suspense fallback={<Spinner />}>
  <UserInfo />
</Suspense>

---

三、实用工具（jotai/utils）

atomWithStorage —— 持久化

import { atomWithStorage } from 'jotai/utils';

// 自动读写 localStorage
const themeAtom = atomWithStorage('theme', 'light');
// key: 'theme'，默认值: 'light'

// 支持 sessionStorage
const tokenAtom = atomWithStorage('token', '', sessionStorage);

atomWithDefault —— 懒初始化

import { atomWithDefault } from 'jotai/utils';

const currentUserAtom = atomWithDefault(async () => {
  const res = await fetch('/api/me');
  return res.json();
});
// 首次读取时执行，之后可通过 set 覆盖

atomFamily —— 参数化 Atom

import { atomFamily } from 'jotai/utils';

// 根据参数创建不同的 atom 实例
const todoAtomFamily = atomFamily((id: string) =>
  atom(async () => {
    const res = await fetch(`/api/todos/${id}`);
    return res.json();
  })
);

function TodoItem({ id }: { id: string }) {
  const todo = useAtomValue(todoAtomFamily(id));
  return <span>{todo.text}</span>;
}

selectAtom —— 细粒度选择

import { selectAtom } from 'jotai/utils';

const userNameAtom = selectAtom(userAtom, (user) => user.name);
// 只有 user.name 变化时才触发使用 userNameAtom 的组件重渲染

---

四、核心原理

Atom 的本质

// atom() 返回的不是值，而是一个"配置对象"
function atom(read, write) {
  const config = {
    read,    // 读取函数 或 初始值
    write,   // 写入函数（可选）
  };
  return config;
}

// 原始 atom：read = initialValue（非函数）
const countAtom = atom(0);
// → { read: 0, write: defaultWrite }

// 派生 atom：read = (get) => ...
const doubleAtom = atom((get) => get(countAtom) * 2);
// → { read: (get) => get(countAtom) * 2, write: undefined }

Store 的核心结构

// Store 内部用 WeakMap 存储每个 atom 的状态
const atomStateMap = new WeakMap(); // atom config → atomState

interface AtomState {
  value: any;            // 当前值
  dependencies: Set;     // 依赖的其他 atom
  dependents: Set;       // 被哪些 atom 依赖
  listeners: Set;        // 订阅的组件
}

读取流程

function readAtom(atom) {
  // 1. 如果是原始 atom（read 不是函数），直接返回存储的值
  if (typeof atom.read !== 'function') {
    return atomStateMap.get(atom).value;
  }

  // 2. 如果是派生 atom，执行 read 函数
  //    传入 get 函数，追踪依赖
  const dependencies = new Set();
  const get = (depAtom) => {
    dependencies.add(depAtom); // 记录依赖
    return readAtom(depAtom);   // 递归读取
  };

  const value = atom.read(get);

  // 3. 更新依赖关系
  atomStateMap.get(atom).dependencies = dependencies;
  return value;
}

写入与通知流程

function writeAtom(atom, newValue) {
  // 1. 更新值
  const atomState = atomStateMap.get(atom);
  atomState.value = newValue;

  // 2. 通知所有依赖此 atom 的派生 atom 重新计算
  for (const dependent of atomState.dependents) {
    const newDerivedValue = readAtom(dependent);
    // 如果派生值变化了，继续通知它的 dependent...
  }

  // 3. 通知所有订阅的组件重渲染
  for (const listener of atomState.listeners) {
    listener(); // 触发 React 重渲染
  }
}

与 React 的连接

// useAtomValue 简化实现
function useAtomValue(atom) {
  const store = useStore(); // 获取 Jotai store

  return useSyncExternalStore(
    (callback) => store.subscribe(atom, callback),
    () => store.get(atom)
  );
}

关键设计

1. 依赖图：派生 atom 通过 get() 调用自动建立依赖关系，形成有向无环图（DAG）
2. 按需计算：只有被组件订阅的 atom 才会计算，未订阅的派生 atom 不执行
3. 自动最小订阅：组件用哪个 atom 就只订阅哪个，无需手写 selector
4. WeakMap 存储：atom 被垃圾回收时，对应状态自动清理

---

面试高频题

Q: Jotai 的 atom 和 useState 有什么区别？

• useState 是组件级状态，状态生命周期与组件绑定，跨组件共享需要 props 或 Context
• atom 是全局级状态，存储在 Store 中，任何组件都可直接订阅，且支持派生（一个 atom 依赖另一个 atom）

Q: Jotai 如何实现细粒度重渲染？

每个 atom 独立维护订阅列表。组件通过 useAtomValue(someAtom) 只订阅特定 atom。当 atom A 变化时，只有订阅了 A（或依赖 A 的派生 atom）的组件会重渲染。不像 Context 那样所有消费者都重渲染。

Q: Jotai 和 Recoil 的区别？

• Jotai atom 无需字符串 key（Recoil 的 atom({ key: 'unique' }) 容易冲突）
• Jotai 体积更小（~2KB vs ~20KB）
• Jotai 无需 RecoilRoot Provider（v2+）
• Recoil 已停止维护（Meta 内部项目），Jotai 社区活跃