10分钟 从源码解读 Zustand

420 阅读5分钟

Zustand 是什么?

Zustand 是一个小型、快速的状态管理库,提供十分便捷的 Hook 的状态钩子以解决 react 数据流中的各种问题。

Github 地址戳这里:github.com/pmndrs/zust…

为什么选择 Zustand?

Context 很好。事实上 Context 能够做到 react 业务需求中的大部分事情。只要你把 context 内容拆分的足够细,那你就不需要 selector 来避免频繁的 rerender。但是也同时也意味着顶层的 context provider 结构会变得非常复杂。而 zustand 很好的解决的这一问题。

Redux 很好。Redux 已经久经战阵,证明了它能够承接庞大的规模的业务,满足各种数据流的需求。但是 Redux 本身过于笨重了。说实在的,reduce 的那套 store/action/reducer 的单项数据源本身逻辑并不复杂,但是如果你面对的是并没有那么笨重的业务需求,那为什么不考虑更加简洁的表达方式呢?

总结一下,zustand 的关键词就是:清晰 & 轻量

Show Me The Code

Zustand 的使用 demo 方式非常简单,一眼就能看明白。你只需要声明 hook,然后在任何的函数组件中使用它:

import { create } from 'zustand'

// 创建 store hook
const useBearStore = create((set) => ({
  bears: 0,
  increasePopulation: () => set((state) => ({ bears: state.bears + 1 })),
  removeAllBears: () => set({ bears: 0 }),
}))

// 组件中应用 hook
function BearCounter() {
  const bears = useBearStore((state) => state.bears)
  return <h1>{bears} around here ...</h1>
}
function Controls() {
  const increasePopulation = useBearStore((state) => state.increasePopulation)
  return <button onClick={increasePopulation}>one up</button>
}

Zustand 原理

zustand 到底做了什么?一言以蔽之,zustand 将用户状态封装为外部的 state,并且通过 useSyncExternalStore hook 来提供给用户监听。

这样一来:

  • zustand 状态复用了 useSyncExternalStore 内置的 selector 功能,在合理使用的情况下能够去除不必要的更新逻辑;
  • zustand 通过外部的 store 来封装相关的逻辑,避免 useContext 的复杂嵌套。

useSyncExternalStore

Zustand 实现依赖于 react 的 useSyncExternalStore api。api 的作用参见:react.dev/reference/r…

一言以蔽之:useSyncExternalStore 允许开发者监听一个外部数据源,并且避免由于 React 18 引入的 concurrent mode 所导致的 撕裂问题(React 18 为更新任务设置了不同的优先级,如果外部状态在某一任务执行的过程中受到了变更,那么其他依赖这一数据源的任务读取到的结果就会出现不一致的情况)。

useSyncExternalStore 用法:

export function useSyncExternalStore<Snapshot>(
  subscribe: (onStoreChange: () => void) => () => void,
  getSnapshot: () => Snapshot,
  getServerSnapshot?: () => Snapshot,
): Snapshot;

export function useSyncExternalStoreWithSelector<Snapshot, Selection>(  
  subscribe: (() => void) => () => void,  
  getSnapshot: () => Snapshot,  
  getServerSnapshot: void | null | (() => Snapshot),  
  selector: (snapshot: Snapshot) => Selection,  
  isEqual?: (a: Selection, b: Selection) => boolean,  
  ): Selection

useSyncExternalStore 中包含了三个参数:

  • subscribe:监听外部数据源的变更;
  • getSnapshot:获取外部数据源的快照;
  • getServerSnapshot:服务端渲染所使用的参数,在 ssr 时获取数据源的快照。

Zustand 源码解析

zustand 的源码实现实际上很简单。核心函数包含两个:useStore & createStore。 如果觉得代码比较绕的话可以顺着传入的 createState 依次往下找,就能够把几个函数串起来。

入口函数 create 包含了两个核心步骤:

  • 通过 createStore 创建必要的接口内容;
  • 通过 useStore 将 store 注册到 react 的 useSyncExternalStore hook 中。
export function useStore<TState, StateSlice>(
  api: ReadonlyStoreApi<TState>,
  selector: (state: TState) => StateSlice = identity as any,
) {
  // 通过自定义的发布监听逻辑将 store 注册到 react useSyncExternalStore 中
  const slice = React.useSyncExternalStore(
    api.subscribe,
    // 应用用户传入的 selector
    () => selector(api.getState()),
    () => selector(api.getInitialState()),
  )
  React.useDebugValue(slice)
  return slice
}

const createImpl = <T>(createState: StateCreator<T, [], []>) => {
  // 通过核心的 createStore 函数创建 useStore 需要的接口内容
  const api = createStore(createState)

  const useBoundStore: any = (selector?: any) => useStore(api, selector)

  // 将 api 能力放到 useBoundStore 上允许用户自由调用
  Object.assign(useBoundStore, api)

  return useBoundStore
}

export const create = (<T>(createState: StateCreator<T, [], []> | undefined) =>
  createState ? createImpl(createState) : createImpl) as Create

Store 的封装包含几个部分:

  • setState 也就是用户 store 中函数所拿到的 set 参数(见示例代码)。setState 中会应用用户提前声明的状态更新逻辑,并且触发监听的回调函数;
  • getState:状态的 getter;
  • getInitialState:初始状态的 getter,提供给 useSyncExternalStore 的 ssr 参数;
  • subscribe:简单的发布订阅实现;

这些 api 再加上 create 的封装就足够支持 zustand 的状态管理了:

const createStoreImpl: CreateStoreImpl = (createState) => {
  type TState = ReturnType<typeof createState>
  type Listener = (state: TState, prevState: TState) => void
  let state: TState
  const listeners: Set<Listener> = new Set()

  // setState 也就是用户 strore 中函数所拿到的 set 参数
  const setState: StoreApi<TState>['setState'] = (partial, replace) => {
    const nextState =
      typeof partial === 'function'
        ? (partial as (state: TState) => TState)(state)
        : partial
    // 通过 Object.is 比较是否为同个对象/值
    if (!Object.is(nextState, state)) {
      const previousState = state
      // 应用更新的 state 并 assign 给当前 state
      state =
        (replace ?? (typeof nextState !== 'object' || nextState === null))
          ? (nextState as TState)
          : Object.assign({}, state, nextState)
      listeners.forEach((listener) => listener(state, previousState))
    }
  }

  const getState: StoreApi<TState>['getState'] = () => state

  const getInitialState: StoreApi<TState>['getInitialState'] = () =>
    initialState

  // 简单的发布订阅模式
  const subscribe: StoreApi<TState>['subscribe'] = (listener) => {
    listeners.add(listener)
    // Unsubscribe
    return () => listeners.delete(listener)
  }

  const api = { setState, getState, getInitialState, subscribe }
  const initialState = (state = createState(setState, getState, api))
  return api as any
}

export const createStore = ((createState) =>
  createState ? createStoreImpl(createState) : createStoreImpl) as CreateStore

再补充一个常用的 useShallow hook 吧:useShallow 通常在拆分 store 的部分状态时使用,其核心逻辑就是通过 ref 暂存状态对象,并且通过浅比较来避免不必要的状态更新过程:

// 在用户传入的 selector 中使用,通过 ref 存储浅比较没有更新的内容
export function useShallow<S, U>(selector: (state: S) => U): (state: S) => U {
  const prev = React.useRef<U>()
  return (state) => {
    const next = selector(state)
    return shallow(prev.current, next)
      ? (prev.current as U)
      : (prev.current = next)
  }
}

// useShallow 常见用法:
const { nuts, honey } = useBearStore(
  useShallow((state) => ({ nuts: state.nuts, honey: state.honey })),
)

写在最后

最后的最后其实没有什么可说的,zustand 本身的代码实现并不复杂。但是它就是能够通过如此精简的代码解决了 react 实际开发场景中遇到的诸多问题。大概这就是巧妙的 idea 所展现出的价值吧。