状态管理没那么复杂：手写实现 Zustand 核心逻辑

ETA8
5 阅读8分钟

在 React 生态中，状态管理是一个绕不开的话题。从早期的 Flux、Redux，到后来的 MobX、Recoil，再到如今流行的 Zustand，每一款库的出现都试图解决组件间通信的痛点。

Zustand 之所以能迅速获得大量关注，核心在于它的“极简”。它不需要包裹 Provider，没有繁琐的 Action 和 Reducer，API 设计直观且易于上手。但作为开发者，仅仅学会使用 API 是不够的。理解其背后的实现原理，不仅能让我们更放心地在生产环境使用它，还能提升我们设计架构的能力。

今天，我们将剥离掉 Zustand 复杂的边界情况处理，从核心逻辑出发，手写一个简化版的 Zustand。通过这个过程，我们将深入探讨状态管理的本质：状态存储、发布订阅模式以及与 React 渲染机制的结合。

一、核心设计思路

在动手写代码之前，我们需要明确 Zustand 的几个核心设计目标：

  1. 状态外置：Store 不应该依赖于 React 组件树的生命周期，它应该是一个独立的单例对象。
  2. 发布订阅（Pub/Sub）：当状态发生变化时，需要通知所有订阅者。这是实现响应式更新的基础。
  3. 按需订阅：组件不应该因为 Store 中任何微小的变化而重新渲染，它只应该关心自己用到的那部分状态。
  4. 无 Provider：利用 Hook 的特性，直接在组件内部调用，减少样板代码。

基于以上目标，我们可以将实现分为两层：

  • 基础 Store 层：负责状态的存储、修改和通知（与 React 无关）。
  • React 绑定层：负责订阅 Store 变化，并触发组件更新。

二、基础 Store 的实现

首先，我们需要一个函数来创建 Store。这个 Store 需要维护当前的状态（state），以及一组监听器（listeners）。

1. 状态与闭包

为了保持状态的私有性，我们利用 JavaScript 的闭包特性。state 变量不直接暴露，而是通过 getState 和 setState 来访问。

const createStore = (createState) => {
    let state; // 状态存储在闭包中
    const listeners = new Set(); // 使用 Set 存储监听器，避免重复

    const getState = () => state;

    const setState = (partial, replace = false) => {
        // 支持函数式更新，如 setState(prev => ({ count: prev.count + 1 }))
        const nextState = typeof partial === 'function' ? partial(state) : partial;
        
        // 只有状态真正发生变化时，才通知订阅者
        if (!Object.is(state, nextState)) {
            const prevState = state;
            // 默认行为是浅合并，除非指定 replace 为 true
            if (!replace) {
                state = typeof nextState !== 'object' 
                    ? partial 
                    : Object.assign({}, state, nextState);
            } else {
                state = nextState;
            }
            // 通知所有监听器
            listeners.forEach(listener => listener(state, prevState));
        }
    };

    const subscribe = (listener) => {
        listeners.add(listener);
        // 返回取消订阅的函数
        return () => listeners.delete(listener);
    };

    // 初始化状态
    state = createState(setState, getState);

    return { getState, setState, subscribe };
};

实现细节分析：

  • 浅合并策略：代码中使用了 Object.assign。这是 Zustand 默认的行为。它的优点是简单高效，但缺点是对于深层嵌套的对象，修改深层属性可能会导致引用不变，从而无法触发更新。在实际使用中，我们通常建议保持状态扁平化，或者配合 Immer 等库使用。
  • 变化检测：Object.is(state, nextState) 用于判断状态是否真的变了。如果引用相同，则跳过更新和通知，这是一种基础的性能优化。
  • 监听器管理：使用 Set 数据结构可以保证监听器的唯一性。subscribe 返回一个清理函数，这对于防止内存泄漏至关重要。

三、与 React 的结合

有了基础的 Store，下一步是如何让 React 组件“感知”到状态的变化。这需要一个自定义 Hook。

1. 强制渲染机制

React 组件的更新依赖于 State 或 Props 的变化。但我们的 Store 在 React 外部，它的变化不会自动触发组件更新。因此，我们需要在 Hook 内部维护一个局部的 State，当 Store 通知变化时，修改这个局部 State 来“欺骗”React 进行重渲染。

const useStore = (api, selector) => {
    // 用于触发组件重新渲染的状态
    const [, forceRender] = useState(0);

    useEffect(() => {
        // 订阅 Store 的变化
        const unsubscribe = api.subscribe((state, prevState) => {
            // 核心优化：只有选中的状态发生变化时才更新
            const newObj = selector(state);
            const oldObj = selector(prevState);
            
            // 浅比较引用
            if (newObj !== oldObj) {
                forceRender(Date.now()); 
            }
        });

        // 组件卸载时取消订阅
        return unsubscribe;
    }, [api, selector]);

    // 返回当前选中的状态
    return selector(api.getState());
};

实现细节分析：

  • 选择器（Selector）的重要性：这是 Zustand 性能优化的关键。如果不传 selector，组件会订阅整个 Store 的变化。通过传入 selector，组件只关心自己需要的数据。
  • 引用相等性检查：newObj !== oldObj 是判断是否重渲染的依据。这意味着，如果 selector 返回的是一个新对象（例如在 selector 中创建了新的对象字面量），即使内容相同，引用也会不同，导致不必要的重渲染。因此，最佳实践是 selector 直接返回状态中的某个属性，或者使用 shallow 比较钩子。
  • forceRender 的技巧：代码中使用了 forceRender(Date.now())。这是一种强制更新的手段。在更严谨的实现中，通常会使用一个布尔值切换或者计数器递增，但核心目的都是为了改变 Hook 内部的状态引用，触发 React 的渲染流程。
  • 依赖项：useEffect 的依赖数组中包含了 selector。如果 selector 在每次渲染时都重新定义（例如在组件体内直接写匿名函数），会导致订阅关系不断重建。在实际使用中，建议将 selector 提取到组件外部，或使用 useCallback 包裹。

四、工厂函数 create

为了提供更友好的 API，我们需要一个高阶函数 create，它将 Store 的创建和 Hook 的绑定封装在一起。

export const create = (createState) => {
    const api = createStore(createState);

    // 返回的 Hook 函数
    const useBoundStore = (selector) => {
        // 如果没有传 selector，默认返回整个状态
        const selectedState = selector ? selector(api.getState()) : api.getState();
        // 这里为了简化演示，实际逻辑需结合 useStore 的订阅机制
        // 在完整实现中，这里会调用 useStore(api, selector || identity)
        return useStore(api, selector || ((state) => state));
    };

    // 将 Store 的方法挂载到 Hook 函数上
    // 这样既可以使用 useBoundStore()，也可以使用 useBoundStore.getState()
    Object.assign(useBoundStore, api);

    return useBoundStore;
};

设计考量：

  • 函数即对象：在 JavaScript 中，函数也是对象。通过 Object.assign，我们将 getState、setState 等方法直接挂载到了 useBoundStore 上。这使得开发者可以在组件外部修改状态（例如在事件处理函数或异步逻辑中直接调用 useStore.setState(...)），而不需要进入组件内部。
  • 默认选择器：当用户不传递参数调用 Hook 时，默认返回整个状态树。这虽然方便，但也失去了细粒度更新的优势，因此在大型应用中，建议始终传递 selector。

五、优缺点与改进思考

通过上述手写过程，我们可以更客观地评价 Zustand 这类库的优劣。

优点

  1. 架构简单：没有 Provider 包裹，减少了组件树的嵌套层级，代码更清爽。
  2. 灵活性强：状态逻辑可以在组件外部复用，不依赖 React 上下文。
  3. 性能可控：通过 selector 机制，可以精确控制组件的渲染范围，避免全局状态变更导致的全局重绘。

潜在问题与改进

  1. 浅合并的局限：如前所述，Object.assign 只能处理第一层属性的合并。如果状态结构复杂，修改深层属性需要展开整个路径，代码会变得冗长。
    • 改进方案：集成 Immer，支持直接修改“可变”语法的状态更新。
  2. Stale Closure（过时的闭包）风险：在 setState 接收函数的场景中，如果依赖了外部变量，可能会捕获旧的值。
    • 改进方案：始终推荐使用函数式更新 set((prev) => ...) 来确保获取最新状态。
  3. 并发渲染兼容性：在 React 18 的并发模式下，渲染可能会被中断。简单的 useState 触发更新在某些极端边缘情况下可能需要更细致的处理（例如使用 useSyncExternalStore）。
    • 改进方案：Zustand 官方新版本已经迁移到 useSyncExternalStore，以更好地支持并发特性。我们手写的版本使用的是 useEffect + useState，是兼容旧版本的经典实现方式。
  4. 调试困难：由于状态在外部，且更新可能在任何地方发生，追踪状态变化来源比 Redux 困难。
    • 改进方案：接入 Redux DevTools 中间件，记录状态变更历史。

六、总结

手写 Zustand 的过程，本质上是对“状态”与“视图”分离思想的一次实践。

我们看到了一个轻量级状态管理库是如何通过闭包隐藏实现细节，通过发布订阅解耦状态与视图，通过选择器优化渲染性能的。

理解这些底层原理，对我们日常开发有直接的指导意义：

  • 在设计全局状态时，尽量保持结构扁平，利于浅合并和选择器提取。
  • 在编写 selector 时，注意引用稳定性，避免不必要的渲染。
  • 明白组件外修改状态的可行性，但也需警惕由此带来的数据流追踪难度。

状态管理没有银弹，Zustand 的流行是因为它在简单和功能之间找到了一个很好的平衡点。通过源码剖析，我们不仅掌握了一个工具，更掌握了一种处理复杂数据流的思维模式。希望这篇文章能帮助你在使用 Zustand 时更加得心应手。

avatar
文章
阅读
粉丝