react 状态管理之 Zustand 使用与实现原理

在 react 这种数据驱动的框架中, 状态管理是应用的重要组成部分，我们需要跟踪和管理应用程序的状态，确保组件之间的数据同步和更新，从而驱动应用视图的变化。今天呢主要给大家介绍一下一个比较轻量且使用简单的状态管理库 Zustand，并且对比 zustand 和其他react状态管理库的异同。

Zustand 是德语中 状态 也就是 state 的意思, 它具有如下一些特点:

• 较少的模板语法, 对比 redux 不需要定义 Action, Reducer, combine reducer
• 状态集中管理，并通过 action 函数进行更新
• 使用 react hooks 在组件中使用 state
• Api 简单从而开发者使用简单
• 无需使用 React Context 这意味着不用使用 Provider 包裹根组件

Zustand 的使用

在使用状态管理的App中最简单的场景中, 我们通常需要以下几个最简单的几个功能:

• 能够创建一个 store 该 store 中有初始的 state
• 能够在视图中渲染 state 中的数据
• 在用户操作视图时能够改变 state 中的数据, 并且重新渲染视图

例如我们通过下面的例子来实现一个简单的计数器:

import React from 'react';
import styles from './Counter.module.css';
import { create } from 'zustand';

/**
 * 首先我们需要创建一个 `store`, 该 `store` 即整个应用的状态, 它通常由两部分组成, 一个是数据, 另一个是改变数据的函数
 */
const useCountStore = create((setState, getState, api) => ({
  value: 0,
  increment: () => setState((state) => ({ value: state.value + 1 })),
  decrement: () => setState((state) => ({ value: state.value - 1 })),
}));

export function Counter() {
  /**
   * 其次我们能够通过 create 函数返回的 hook, 传入 selector 函数获取 state 中的数据或者更新函数
   */
  const value = useCountStore((state) => state.value);
  const increment = useCountStore((state) => state.increment);
  const decrement = useCountStore((state) => state.decrement);

  return (
    <div>
      <div className={styles.row}>
        <button
          className={styles.button}
          aria-label="Decrement value"
          /**
           * 最后我们能够在用户操作时改变数据, 并且重新入 render
           */
          onClick={() => decrement()}
        >
          -
        </button>
        <span className={styles.value}>{value}</span>
        <button
          className={styles.button}
          aria-label="Increment value"
          onClick={() => increment()}
        >
          +
        </button>
      </div>
    </div>
  );
}

效果类似于:

通过 zustand 暴露出的 create 函数定义 store, create 函数接收一个 createState 函数。

createState 函数接收三个参数

• 函数 setState 用于更新 state, setStae 的类型大致是 (partial: (state: State) => Partial<State> | Partial<State>, replace: boolean): void, 它接收两个参数 partial 和 replace。 partial 可以是一个返回值是 state 子集的函数, 也可以是一个 state 子集对象。 replace 为 true 时会直接用前面计算出的子集替换 state, replace 默认是 false
• 函数 getState, 执行后返回值是 state
• api, api 是这么一个对象 { setState, getState, subscribe, destroy }, 其中 setState getState和前面的参数一样, 其中的 subscribe 是一个发布订阅模式中用于注册监听函数, 该监听函数会在 setState 执行后被执行并传入最新的 state, 内部就是靠这个 listener 来让组件重新 render。其实就是在 state 变更之后重新触发组件 render。clear 用于清空 listener。

createState 函数的类型定义大致入下:

type SetState = (
  partial: T | Partial<T> | { (state: T): T | Partial<T> },
  replace?: boolean | undefined,
): void;


type GetState = () => T;

export type StateCreator<T> = (
  setState: SetState;
  getState: GetState;
  store: {
    setState: SetState;
    getState: GetState;
    subscribe: (listener: (state: T, prevState: T) => void) => () => void;
    destroy: () => void;
  }
) => U

并且 create 函数返回一个 useBoundStore hooks函数, 用于在组件中获取状态中的数据, useBoundStore 接收两个参数:

• selector 一个可选的查询函数, 该函数接收 state 作为参数, 可以在 selector 函数中计算需要的数据并返回该值, 该值会成为 useBoundStore 的返回值, 如果 selector 不传, 那么默认会传入 getState 函数作为 selector 即 useBoundStore 返回整个 state
• isEqual(currentSlice: Partial<T>, nextSlice: Partial<T>): boolean, 该函数用于在比较 selector 的返回值在 state 变更时前后两次的值是否相等, 如果返回 true 意味着两次计算值相等那么就返回上一次的值, 反之亦然。如果 isEqual 为 undefined 那么除非前后两次的 state 相等(基于Object.is) 否则每次都返回最新的计算值

实现原理

根据 Zustand 4.x 版本的源码简化后的代码如下, 核心逻辑不变:

/**
 * createStoreImpl 用于创建 store, 本质上是一个单例, 并且是一个闭包, 通过暴露出 setState 和 getState 来更新和获取 store
 * 此外还通过发布订阅模式在使用 setState 即 store 状态变更时通知外部的监听函数执行也就是重新让 react 组件 render
 */
const createStoreImpl = (createState) => {
  let state

  // 发布订阅, 用于将所有监听函数收集起来
  const listeners = new Set()

  /**
   * 更新 store 中状态的函数, 接受 partial 作为参数, 如果 partial 是一个函数则会执行 partial 函数拿到结果来更新状态
   * 第二个参数 replace 用于标识是否全量更新
   */
  const setState = (partial, replace) => {
    // partial 可以是一个返回新状态的函数, 或者本身就是一个新状态对象
    const nextState =
      typeof partial === 'function'
        ? partial(state)
        : partial
    if (!Object.is(nextState, state)) {
      const previousState = state
      state =
        replace ?? typeof nextState !== 'object'
          ? nextState
          : Object.assign({}, state, nextState)

      // 让监听函数执行
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

  // getState 就是获取内部的状态
  const getState = () => state

  // 注册监听函数的方法
  const subscribe = (listener) => {
    listeners.add(listener)
    // Unsubscribe
    return () => listeners.delete(listener)
  }

  const destroy = () => {
    listeners.clear()
  }

  const api = { setState, getState, subscribe, destroy }

  // 初始化状态, 这里的 state 就是初始化的 state
  state = createState(setState, getState, api)
  return api
}

export const createStore = (createState) => createStoreImpl(createState);

export function useStore(
  api,
  selector,
  equalityFn,
) {
  // useSyncExternalStoreWithSelector 和 useSyncExternalStore 类似, 不同的地方在于可以通过传入 selector 对获取的状态进行一个处理, 并且可以传入 equalityFn 这个函数会比较前后两次的结果, 在两次结果不同时才会触发组件 render
  // 下面我会详细介绍 useSyncExternalStoreWithSelector
  const slice = useSyncExternalStoreWithSelector(
    api.subscribe,
    api.getState,
    api.getServerState || api.getState,
    selector,
    equalityFn
  )
  return slice;
}

/**
 * 用于创建store的 create 方法实现
 * createState 参数就是我们使用 create 方法传入的函数
 */
const createImpl = (createState) => {
  // 创建单例 store, 并暴露出 setState, getState, subscribe 等方法用于更新 store , 获取 store 中的状态 以及监听状态变化
  const api = createStore(createState);

  // useBoundStore 就是我们用于查询分状态的 hook 函数, 这里的 useStore hook 内部使用 useSyncExternalStoreWithSelector hook 让状态更新时重新 render 组件
  // 所以其实 selector 的粒度越细, 越能减少组件的更新
  const useBoundStore = (selector, equalityFn) => useStore(api, selector, equalityFn)

  Object.assign(useBoundStore, api)

  return useBoundStore
}

// 从 Zustand 中导出的 create 方式即为我们用来创建 store 的方法
// create 方法返回一个 hook, 我们可以将 selector 函数传入这个 hook 获取我们需要的状态
export const create = (createState) => createImpl(createState);

以上就是 Zustand 的核心逻辑, 和 redux 类似都是通过维护一个私有单例用于存储状态, 并且通过发布订阅暴露出状态变更的监听函数, 并通过 useSyncExternalStoreWithSelector 来监听状态变化从而重新 render 组件, 下面介绍一下 useSyncExternalStoreWithSelector 这个 hook 并不是 react 自带的 hook

首先看一下 useSyncExternalStore, 这个 hook 用于订阅外部的 store 变化, 在 store 变化时会执行传入的 getSnapshot 函数, react 会比较前后两次的结果如果结果不相同那么会触发组件 render, 它返回一个状态快照 , 并且接受两个参数:

• subscribe 函数用于订阅 store 变化
• getSnapshot 用于获取状态快照

const snapshot = useSyncExternalStore(subscribe, getSnapshot)

useSyncExternalStoreWithSelector 内部使用了 useSyncExternalStore, 相比下多接受 selector 和 isEqual 函数, 可以粒度更细地获取快照和控制组件是否重新 render, 它的函数签名如下:

function useSyncExternalStoreWithSelector(
  // 监听函数, 和 useSyncExternalStore 一样
  subscribe: (() => void) => () => void,
  // 获取快照的函数, 和 useSyncExternalStore 一样
  getSnapshot: () => Snapshot,
  getServerSnapshot: void | null | (() => Snapshot), // ssr 中使用
  // 查询函数, 接收 getSnapshot 的执行结果, selector 函数的返回结果也是 useSyncExternalStoreWithSelector 的返回结果
  selector: (snapshot: Snapshot) => Selection,
  // 比较前后两次 selector 的执行结果是否相通, 不相同才会重新 render 组件
  isEqual?: (a: Selection, b: Selection) => boolean,
): Selection

useSyncExternalStoreWithSelector 的实现如下, 简化了源码中的一些类型定义和一些调试逻辑

const { useRef, useEffect, useMemo, useSyncExternalStore } = React;

export function useSyncExternalStoreWithSelector(
  subscribe,
  getSnapshot,
  getServerSnapshot,
  selector,
  isEqual,
) {
  const instRef = useRef(null);

  // 使用 inst 来记录上一次的值, 这样即使在getSnapshot selector isEqual参数变化时也能够比较上一次的值来决定是否需要重新 render
  let inst;
  if (instRef.current === null) {
    inst = {
      hasValue: false,
      value: null,
    };
    instRef.current = inst;
  } else {
    inst = instRef.current;
  }

  // 通过 useMemo 来构造传给 useSyncExternalStore 的参数
  const [getSelection, getServerSelection] = useMemo(() => {
    let hasMemo = false; // 是否有缓存
    let memoizedSnapshot; // 上一次缓存的全部状态值
    let memoizedSelection; // 上一次缓存的经过 selector 计算的状态值

    /**
     * 记忆函数, 在前后两次 selection 相同时返回上一次的缓存值
     */
    const memoizedSelector = (nextSnapshot) => {
      // 第一次没有缓存时
      if (!hasMemo) {
        hasMemo = true;
        memoizedSnapshot = nextSnapshot;
        // 执行 selector 得到 selection
        const nextSelection = selector(nextSnapshot);
        if (isEqual !== undefined) {
          if (inst.hasValue) {
            // 即使 getSnapshot selector isEqual 变更了, 也要比较上一次的缓存值, 也就是和 inst.value 进行比较
            const currentSelection = inst.value;
            // 如果传入了自定义的比较方法则使用它
            if (isEqual(currentSelection, nextSelection)) {
              memoizedSelection = currentSelection;
              return currentSelection;
            }
          }
        }
        memoizedSelection = nextSelection;
        return nextSelection;
      }

      // 更新时和上一次的计算结果进行比较
      const prevSnapshot = memoizedSnapshot;
      const prevSelection = memoizedSelection;

      if (is(prevSnapshot, nextSnapshot)) {
        return prevSelection;
      }

      const nextSelection = selector(nextSnapshot);

      if (isEqual !== undefined && isEqual(prevSelection, nextSelection)) {
        return prevSelection;
      }

      memoizedSnapshot = nextSnapshot;
      memoizedSelection = nextSelection;
      return nextSelection;
    };

    const maybeGetServerSnapshot = getServerSnapshot === undefined ? null : getServerSnapshot;

    const getSnapshotWithSelector = () => memoizedSelector(getSnapshot());

    const getServerSnapshotWithSelector =
      maybeGetServerSnapshot === null
        ? undefined
        : () => memoizedSelector(maybeGetServerSnapshot());

    return [getSnapshotWithSelector, getServerSnapshotWithSelector];
  }, [getSnapshot, getServerSnapshot, selector, isEqual]);

  // 内部其实是用了 useSyncExternalStore 这个 hook 使 getSelection 的结果发生变更时触发组件 render
  // getSelection 本质上是将 getSnapShot() 的结果作为参数执行 selector 函数, 内部做了缓存以及比较前后两次执行结果是否相等
  const value = useSyncExternalStore(
    subscribe,
    getSelection,
    getServerSelection,
  );

  useEffect(() => {
    inst.hasValue = true;
    inst.value = value;
  }, [value]);

  return value;
}

zustand middleware 的使用

以 immer 为例子, immer 可以 balabal

import React from 'react';
import styles from './Counter.module.css';
import { create } from 'zustand';
import { immer } from 'zustand/middleware/immer';

const useCountStore = create(
  immer(
    (set) => ({
      value: 0,
      increment: () => set((state) => ({ value: state.value + 1 })),
      decrement: () => set((state) => ({ value: state.value - 1 })),

      todos: [
        {
          id: '82471c5f-4207-4b1d-abcb-b98547e01a3e',
          title: 'Learn Zustand',
          done: false,
        },
        {
          id: '354ee16c-bfdd-44d3-afa9-e93679bda367',
          title: 'Learn Jotai',
          done: false,
        },
        {
          id: '771c85c5-46ea-4a11-8fed-36cc2c7be344',
          title: 'Learn Valtio',
          done: false,
        },
      ],
      complete: (todoId) => {
        set((state) => {
          state.todos.find(({ id }) => todoId === id).done = true;
        });
      }
    }),
  ),
);

export function Counter() {
  const value = useCountStore((state) => state.value);
  const increment = useCountStore((state) => state.increment);
  const decrement = useCountStore((state) => state.decrement);

  const todos = useCountStore(state => state.todos);
  const complete = useCountStore(state => state.complete);

  return (
    <div>
      <div className={styles.row}>
        <button
          className={styles.button}
          aria-label="Decrement value"
          onClick={() => decrement()}
        >
          -
        </button>
        <span className={styles.value}>{value}</span>
        <button
          className={styles.button}
          aria-label="Increment value"
          onClick={() => increment()}
        >
          +
        </button>
      </div>
      <div className={styles.row}>
        <ul>
          {
            todos.map(({ id, title, done }) => (
              <li
                className={`${styles.totoItem} ${done ? styles.done : ''}`}
                key={id}
              >
                {title} <button className={styles.todoButton} onClick={() => complete(id)}>✔️</button>
              </li>
            ))
          }
        </ul>
      </div>
    </div>
  );
}

自定义中间件

不同于 redux 的中间件有固定的格式, 并且有 applyMiddleware 来注册中间件, zustand 的中间件比较自由, zustand 的中间件是 createState 函数的 enhancer 增强器, 这个 middleware 之需要最终返回一个 createState 函数即可, 逻辑可以在 middleware 内部进行处理

以一个 store logger middleware 为例, 它可以在每次 store 变更时即调用 setState 函数时, 打日志

import React from 'react';
import styles from './Counter.module.css';
import { create } from 'zustand';

const logger = (createStore) => (setState, getState, api) => {
  return createStore((...args) => {
    console.group(`set state`);
    console.info(`current state`, getState());
    setState(...args);
    console.info(`next state`, getState());
    console.groupEnd(`set state`);
  }, getState, api);
}

const useCountStore = create(
  logger(
    (set) => ({
      value: 0,
      increment: () => set((state) => ({ value: state.value + 1 })),
      decrement: () => set((state) => ({ value: state.value - 1 })),
    }),
  ),
);

export function Counter() {
  const value = useCountStore((state) => state.value);
  const increment = useCountStore((state) => state.increment);
  const decrement = useCountStore((state) => state.decrement);

  return (
    <div>
      <div className={styles.row}>
        <button
          className={styles.button}
          aria-label="Decrement value"
          onClick={() => decrement()}
        >
          -
        </button>
        <span className={styles.value}>{value}</span>
        <button
          className={styles.button}
          aria-label="Increment value"
          onClick={() => increment()}
        >
          +
        </button>
      </div>
    </div>
  );
}

效果如下:

zustand 的中间件比较自由没有什么格式, 但是实际上它的中间件仍然是一个洋葱模型, 例如以下使用的三个中间件

const useCountStore = create(
  immer(logger(devtools(
    (set) => ({
      value: 0,
      increment: () => set((state) => ({ value: state.value + 1 })),
      decrement: () => set((state) => ({ value: state.value - 1 })),
    }),
  ))),
)

他们的执行顺序是 immer 先执行 => 接着 logger 执行 => 最后是 devtool 执行, 接着真正地使用 setState 来改变store, 最后 devtool 执行结束, logger 执行结束, immer 执行结束

与 redux 对比

仅仅用 redux 的话, 模板代码会很多, 比如定义 reducer, action, 在组合 state slice 切片时也会有很多模板代码, 所以基本上不会直接使用 redux 而是会和 redux-toolkit 一起使用。

redux-toolkit 可以通过 configureStore 减少模板代码, 通过 createReducer 简化 reducer 函数, 并且内置了 immer reselect 等比较方便的库, 所以以下的比较主要是基于 redux-toolkit 和 zustand

	zustand	redux-tdk
单向数据流	是	是
不可变数据	是	是
是否支持react hooks	是	是
是否需要react context	否	是
学习曲线	api简单, 概念少, 可直接处理异步	概念较多, 需要额外的库来处理副作用比如 redux-thunk, redux-saga 增加了学习成本
特点	轻量化, 体检小, api简单使用简单	通过派发action来改变store, store集中 整个应用共享一个 store, 状态变更可追踪,功能强大 社区繁荣, 成熟的 devtools

zustand 和 redux 很像, 他们都是单向数据流且基于不可变数据, 这使得他们都能够做到Time Travel, 应用状态的回退和还原都可以很方便的做到。zustand 拥有轻量化, 使用方便的优点。而 redux 需要写大量的模板代码, 不过 redux-toolkit 解决了这个问题并内置了很多使用的库, 不过这同时又增加了学习成本。redux 拥有强大的社区支持, 通过派发Action这一唯一方式改变store使得redux应用的状态变更变得更加容易被观测和记录。比如需要对用户的行为进行埋点时，可以使用自定义的 redux 中间件在每一次 action 派发时进行数据埋点, 并分析action中的payload, 而 zustand 没发做到这一点, zustand 需要在每一个函数中去完成这个埋点的动作。