Zustand 等状态管理库能实现，全靠 React 提供的 useSyncExternalStore 钩子 😃😃😃

Moment
2,707 阅读16分钟

Zustand 是一个用于 React 的轻量级状态管理库。它的名字源自德语，意思是“状态”。Zustand 的设计目标是提供一种简单、快速且高效的方式来管理应用程序的状态。它通过避免复杂的设置和繁琐的样板代码，使开发者能够更专注于应用的逻辑和功能。

它的核心理念是使用一个单独的状态容器来管理应用的状态，并通过钩子函数（hooks）将状态和组件连接起来。它利用了现代 React 的特性，如 hooks 和上下文（context），以确保高性能和灵活性。

基本使用

首先我们要在 src 目录下创建一个 store 文件，并创建一个 index.ts 文件，编写如下代码：

import { create } from "zustand";

export interface State {
  count: number;
}

export interface Actions {
  increment: () => void;
  decrement: () => void;
}

const useStore = create<State & Actions>((set) => ({
  count: 0,
  increment: () => set((state) => ({ count: state.count + 1 })),
  decrement: () => set((state) => ({ count: state.count - 1 })),
}));

export default useStore;

在上面的这些代码中，使用 create 函数创建一个存储：

  1. set 是一个函数，用于更新状态。

  2. 初始状态定义为一个对象，包含 count 属性，初始值为 0。

  3. increment 函数：调用 set 函数来更新状态，将 count 的值加 1。

  4. decrement 函数：调用 set 函数来更新状态，将 count 的值减 1。

细心的你会发现，它不需要像 redux 那样使用 ... 解构语法，因为 zustand 会将新的部分状态对象与现有状态合并，这个过程中的合并是自动完成的。

这会我们就可以直接在项目中使用了：

import { create } from "zustand";

export interface State {
  count: number;
}

export interface Actions {
  increment: () => void;
  decrement: () => void;
}

const useStore = create<State & Actions>((set) => ({
  count: 0,
  increment: () => set((state) => ({ count: state.count + 1 })),
  decrement: () => set((state) => ({ count: state.count - 1 })),
}));

export default useStore;

20240702111940

Zustand 主体流程

zustand 的核心是将外部 store 和组件 view 的交互，交互的核心流程如下图：

20240923183416

先使用 create 函数基于注入的 initStateCreateFunc 创建一个闭包的 store，并暴露对应的 subscribe、setState、getState 这几个 api

借助于 react 官方提供的 useSyncExternalStoreWithSelector 可以将 store 和 view 层绑定起来，从而实现使用外部的 store 来控制页面的展示。

zustand 还支持了 middleware 的能力，采用 create(middleware(...args)) 的形式即可使用对应的 middleware

useSyncExternalStore

在 React 中，我们所说的状态通常分为三种：

  1. 组件内部的 State/Props

  2. 上下文 Context

  3. 组件外部的独立状态 Store(Redux/Zustand)

前两种状态实际上都是 React 内部维护的 Api，自然也会跟随着 React 版本的迭代而进行相对应的优化。

但是组件外部的状态，对于 React 来说并不可控，如果需要更好的契合 React 本身，我们需要去写一些与本身业务逻辑无关的胶水代码

例如：

  • 订阅外部状态

  • 外部状态更新时，对组件进行重渲染。

在 React 18 引入 Concurrent Mode 后，外部状态订阅模式可能会引发一个被称为“撕裂问题”的 bug。让我们考虑以下场景：

当页面触发更新并进行重新渲染时，Concurrent Mode 会根据任务优先级对更新进行划分，优先级低的任务可能会被打断。假设任务 A 和任务 B 同时依赖于外部状态中的某个值。在重新渲染开始时，该值为 1。任务 A 执行完毕后，React 将线程控制权交还给浏览器，此时，浏览器的某些操作可能会将该值更新为 2。当任务 B 重新恢复执行时，它读取到的值可能与任务 A 执行时的值不一致，从而导致渲染结果出现差异。

请看如下代码示例：

export const useOutSideStore = (store) => {
  const [state, setState] = useState(store.getState());

  useEffect(() => {
    const unsubscribe = store.subscribe(() => {
      setState(store.getState());
    });

    return () => {
      unsubscribe();
    };
  }, []);

  return state;
};

在这段代码中，useState 初始化状态为 store 的当前值，随后 store.subscribe 监听状态变化。当状态更新时，setState 将会触发重新渲染。然而，由于 Concurrent Mode 的特性，如果在任务 A 执行期间，store 的值发生变化，任务 B 可能会在恢复执行时读取到一个不一致的值，造成渲染结果的差异。

针对上述的一些外部状态与 React 本身不契合的情况，React 提供了一个名为 useSyncExternalStore 的 Hook，这个 hook 可以让我们更加方便的去订阅外部的 Store，并且避免发生撕裂问题。

useSyncExternalStore 是 React 18 引入的一个钩子，用于安全地与外部存储（如 Redux 或其他状态管理库）进行交互。它的设计考虑了 Concurrent Mode，以确保在并发渲染期间状态的一致性。

它允许你从外部存储中同步读取状态。这意味着组件在渲染时总是可以获得最新的状态。并且能够处理状态在多次渲染之间发生变化的情况，避免了“撕裂问题”。当状态更新时，组件可以安全地反映这些变化。

useSyncExternalStore 接受三个参数：

  1. subscribe：一个函数，用于订阅外部存储的变化。当外部状态发生变化时，该函数会被调用，通知组件进行更新。

  2. getSnapshot：一个函数，用于获取当前的状态快照。每次渲染时，React 会调用这个函数以获取最新的状态。

  3. getServerSnapshot（可选）：在服务器渲染时使用的函数。它提供服务器端的状态快照。

下面是一个使用 useSyncExternalStore 的示例：

import { useSyncExternalStore } from "react";

const useStore = (store) => {
  const state = useSyncExternalStore(
    store.subscribe,
    store.getState,
    // 可选的服务器快照函数
    store.getServerState // 如果有的话
  );

  return state;
};

在 react 组件初次挂载的时候对应的是 mountSyncExternalStore、re-render 更新渲染的时候对应的是 updateSyncExternalStore。

mountSyncExternalStore

function mountSyncExternalStore<T>(
  subscribe: (() => void) => () => void,
  getSnapshot: () => T,
  getServerSnapshot?: () => T
): T {
  // 当前正在处理的 fiber node
  const fiber = currentlyRenderingFiber;
  // 给当前fiber创建hook对象，如果是fiber的第一个hook对象，存放在fiber的memoizedState上，如果当前fiber存在hook对象用next连成单链表
  const hook = mountWorkInProgressHook();

  let nextSnapshot;
  const isHydrating = getIsHydrating();

  // 判断当前协调是否是 hydrate
  if (isHydrating) {
    if (getServerSnapshot === undefined) {
      throw new Error(
        "Missing getServerSnapshot, which is required for " +
          "server-rendered content. Will revert to client rendering."
      );
    }
    nextSnapshot = getServerSnapshot();
  } else {
    // 取出快照state值
    nextSnapshot = getSnapshot();

    const root: FiberRoot | null = getWorkInProgressRoot();

    if (root === null) {
      throw new Error(
        "Expected a work-in-progress root. This is a bug in React. Please file an issue."
      );
    }

    if (!includesBlockingLane(root, renderLanes)) {
      pushStoreConsistencyCheck(fiber, getSnapshot, nextSnapshot);
    }
  }

  // 将快照state值存在hook的memoizedState属性上
  hook.memoizedState = nextSnapshot;
  const inst: StoreInstance<T> = {
    value: nextSnapshot,
    getSnapshot,
  };
  hook.queue = inst;

  // 在useEffect hook里面去订阅store，前面api对象的subscribe其实是在这儿用的
  mountEffect(subscribeToStore.bind(null, fiber, inst, subscribe), [subscribe]);

  fiber.flags |= PassiveEffect;

  // 创建effect对象，
  pushEffect(
    HookHasEffect | HookPassive, // 给effect对象打tag标记，注意useLayout的tag是HookHasEffect | HookLayout
    updateStoreInstance.bind(null, fiber, inst, nextSnapshot, getSnapshot),
    undefined,
    null
  );

  return nextSnapshot;
}

在上面的函数中，他接收三个参数：

  1. subscribe：用于订阅外部状态。当外部状态发生变化时，该函数会触发 React 更新。

  2. getSnapshot：用于获取外部状态的快照。每次渲染时，React 调用此函数来读取当前的外部状态。

  3. getServerSnapshot（可选）：在服务端渲染中获取外部状态的函数。它在 SSR 阶段被使用，用于同步服务端和客户端的状态。

最后再简单梳理下 mountSyncExternalStore 里面的逻辑：

  1. mountSyncExternalStore 同其他 hook 一样先创建 hook 对象然后挂在当前 fiber.memoizedState 的 hook 链表上；

  2. 调用 getSnapshot 取得值赋值给 nextSnapshot 并将其存放在 hook.memoizedState 上；

  3. 赋值 hook.queue 为 inst 对象；

  4. 调用 mountEffect 添加一个 useEffect hook，useEffect 的 create 即为 subscribeToStore.bind(null, fiber, inst, subscribe)，deps 即为[subscribe]。

  5. 调用 pushEffect 创建一个带有 HookHasEffect | HookPassive 即 HasEffect | Passvie 标记的 effect 对象；

  6. 返回 nextSnapshot；

我们再来看一下 subscribeToStore、pushStoreConsistencyCheck、updateStoreInstance 的实现。

subscribeToStore

function subscribeToStore(fiber, inst, subscribe) {
  // handleStoreChange 方法在我们通过 store 的 dispatch 方法修改 store 时会触发
  var handleStoreChange = function () {
    if (checkIfSnapshotChanged(inst)) {
      // 如果 store 发生变化，采用阻塞模式渲染
      forceStoreRerender(fiber);
    }
  };
  // 使用 store 提供的 subscribe 方法去订阅
  return subscribe(handleStoreChange);
}

function checkIfSnapshotChanged(inst) {
  var latestGetSnapshot = inst.getSnapshot;
  // 之前的 store 值
  var prevValue = inst.value;

  try {
    // 新的 store 值
    var nextValue = latestGetSnapshot();
    // 浅比较 prevValue， nextValue
    return !objectIs(prevValue, nextValue);
  } catch (error) {
    return true;
  }
}

checkIfSnapshotChanged 函数通过浅比较外部状态的旧值和新值，判断外部状态是否发生变化。如果状态发生了变化，返回 true，以便 React 做出相应的更新处理。

function forceStoreRerender(fiber) {
  scheduleUpdateOnFiber(fiber, SyncLane, NoTimestamp);
}

forceStoreRerender 的作用是在外部状态（如 Redux store）发生变化时，强制触发与该状态相关的组件进行同步更新和重新渲染。它通过 React 的内部调度机制强制将该更新任务设为同步阻塞任务，使 React 立即处理该更新。

因此，在 subscribeToStore 中，React 通过外部状态管理库（如 Redux）的 subscribe 方法订阅了状态的变化。当通过 dispatch 方法修改 store 状态时，store 会遍历已注册的订阅者并按顺序执行订阅的回调函数，此时 handleStoreChange 会被调用。由于状态发生了变化，handleStoreChange 检测到变化后调用 forceStoreRerender，强制触发同步阻塞渲染，以确保组件渲染与最新的外部状态保持一致

pushStoreConsistencyCheck

function pushStoreConsistencyCheck(fiber, getSnapshot, renderedSnapshot) {
  fiber.flags |= StoreConsistency;
  var check = {
    getSnapshot: getSnapshot,
    value: renderedSnapshot,
  };
  var componentUpdateQueue = currentlyRenderingFiber$1.updateQueue;

  if (componentUpdateQueue === null) {
    componentUpdateQueue = createFunctionComponentUpdateQueue();
    currentlyRenderingFiber$1.updateQueue = componentUpdateQueue;
    // 收集 check 对象
    componentUpdateQueue.stores = [check];
  } else {
    var stores = componentUpdateQueue.stores;
    // 收集 check 对象
    if (stores === null) {
      componentUpdateQueue.stores = [check];
    } else {
      stores.push(check);
    }
  }
}

为了保证 store 状态的一致性，React 在 mountSyncExternalStore 方法中首先通过 pushStoreConsistencyCheck 将 check 对象添加到组件节点（Fiber）的 updateQueue 中进行状态追踪。然后，在协调（reconciliation）过程完成后，React 会再次遍历整个 Fiber 树，基于节点中的 check 对象进行状态一致性检查。如果发现 store 状态与渲染时不一致，React 将通过 renderRootSync 方法进行一次同步阻塞渲染，以确保最终的 UI 和 store 状态保持一致。

updateStoreInstance

function updateStoreInstance(fiber, inst, nextSnapshot, getSnapshot) {
  inst.value = nextSnapshot;
  inst.getSnapshot = getSnapshot;
  if (checkIfSnapshotChanged(inst)) {
    // 在 commit 阶段，检查 store 是否发生变化，如果发生变化，触发同步阻塞渲染
    forceStoreRerender(fiber);
  }
}

updateStoreInstance 函数在 React 渲染过程中更新外部状态的快照值，并检查外部状态是否发生了变化。如果检测到外部状态（如 Redux store）与当前渲染时的快照不一致，它会强制触发同步阻塞渲染，确保 UI 与外部状态保持一致。这种机制有效避免了 Concurrent Mode 下可能出现的 UI 和状态不匹配的“撕裂问题”。

useSyncExternalStore

看完 mountSyncExternalStore 的实现之后，我们再来看一下 useSyncExternalStore 在 update 阶段要执行的 updateSyncExternalStore 的实现。

function updateSyncExternalStore<T>(subscribe, getSnapshot, getServerSnapshot) {
  // 当前正在构建的fiber
  const fiber = currentlyRenderingFiber;
  // 更新hook，更新的时候其实是对原来的hook对象复制一份，有更新的更新便是
  const hook = updateWorkInProgressHook();
  // Read the current snapshot from the store on every render. This breaks the
  // normal rules of React, and only works because store updates are
  // always synchronous.
  // 同样的调用getSnapshot取得最新的值并赋值给nextSnapshot
  const nextSnapshot = getSnapshot();
  // 省略部分开发环境代码
  // 将之前存在hook.memoizedState的上一次的nextSnapshot取出来
  const prevSnapshot = hook.memoizedState;
  // 通过is比对看snapshot是否变化
  const snapshotChanged = !is(prevSnapshot, nextSnapshot);
  if (snapshotChanged) {
    // 变化了的话，把最新的snapshot赋值给hook.memoizedState，并标记更新
    hook.memoizedState = nextSnapshot;
    markWorkInProgressReceivedUpdate();
  }
  // 取出之前存放在hook.queue上的inst对象
  const inst = hook.queue;
  // 后面useEffect相关逻辑不再展开了，简单说就是判断是否去useEffect里面重新添加listener
  updateEffect(subscribeToStore.bind(null, fiber, inst, subscribe), [
    subscribe,
  ]);

  // Whenever getSnapshot or subscribe changes, we need to check in the
  // commit phase if there was an interleaved mutation. In concurrent mode
  // this can happen all the time, but even in synchronous mode, an earlier
  // effect may have mutated the store.
  if (
    inst.getSnapshot !== getSnapshot ||
    snapshotChanged ||
    // Check if the susbcribe function changed. We can save some memory by
    // checking whether we scheduled a subscription effect above.
    (workInProgressHook !== null &&
      workInProgressHook.memoizedState.tag & HookHasEffect)
  ) {
    fiber.flags |= PassiveEffect;
    pushEffect(
      HookHasEffect | HookPassive,
      updateStoreInstance.bind(null, fiber, inst, nextSnapshot, getSnapshot),
      undefined,
      null
    );

    // Unless we're rendering a blocking lane, schedule a consistency check.
    // Right before committing, we will walk the tree and check if any of the
    // stores were mutated.
    const root: FiberRoot | null = getWorkInProgressRoot();

    if (root === null) {
      throw new Error(
        "Expected a work-in-progress root. This is a bug in React. Please file an issue."
      );
    }

    if (!includesBlockingLane(root, renderLanes)) {
      pushStoreConsistencyCheck(fiber, getSnapshot, nextSnapshot);
    }
  }

  return nextSnapshot;
}

在 updateSyncExternalStore 中，主要执行了以下任务：

  1. 调用 getSnapshot 方法获取当前 store 的状态值，并将其存储在 hook 中；

  2. 利用 updateEffect，相当于 useEffect 在更新阶段执行的方法，在组件更新完成后调用 store 提供的 subscribe 方法（如果 subscribe 没有变化，这一步不会执行）；

  3. 标记 PassiveEffect 副作用，在 commit 阶段对外部状态进行一致性检查；

  4. 设置一致性检查机制，在渲染结束时检查 store 状态的一致性。

通过分析源码，可以看出 useSyncExternalStore 通过以下三道机制来确保 store 状态的一致性：

  1. 当通过 dispatch 修改 store 状态时，强制使用 Sync 模式进行同步、不可中断的渲染；

  2. 在 Concurrent 模式下，协调（reconciliation）结束后进行一致性检查，如果状态不一致，强制重新进行一次同步渲染；

  3. 在 commit 阶段，再次进行一致性检查，如果状态仍不一致，强制触发一次同步渲染。

updateSyncExternalStore 通过调用 getSnapshot 获取外部状态的最新快照并将其存储在 hook 中，同时设置副作用（PassiveEffect）在 commit 阶段进行一致性检查。如果在渲染过程中检测到 subscribe 或 selector 发生变化，会通过同步渲染强制更新组件，确保 UI 和外部状态保持一致。此外，它在 Concurrent Mode 下通过多次状态检查和不可中断的同步渲染机制，确保外部状态与渲染结果的一致性，即使在并发渲染的情况下，也能避免状态不一致的风险。

zustand 执行流程

假设我们有以下的这段 react 代码：

function BearCounter() {
  const bears = useStore((state) => state.bears);
  return <h1>{bears} around here...</h1>;
}

他的执行流程如下图所示：

20241005190733

在 zustand 中，createStoreImpl 导出的 subscribe 方法会在 subscribeToStore 函数中被绑定为新的函数，并在 useEffect 中执行。具体流程如下：

组件挂载与 useSyncExternalStore 调用

当 BearCounter 组件渲染时，useSyncExternalStore 通过 useStore 获取当前外部状态（state.bears），并执行 getSnapshot 函数获取最新的 store 快照。

useSyncExternalStore 为当前 Fiber 节点创建一个新的 hook，将其连接到 Fiber 树上，并存储在 fiber.memoizedState 中。此时，nextSnapshot（即 bears 状态值）也被存储在 hook.memoizedState 上。

订阅外部状态变化

在 useEffect 中，subscribeToStore 被执行，subscribe 方法用于将 handleStoreChange 作为监听器添加到 zustand 的 store 中。

React 在这里创建了一个 effect，带有 HasEffect | Passive 标记，表示该副作用将在 commit 阶段执行。这个副作用会将 zustand 的状态监听器添加到 store 中。

zustand 状态更新触发

当 zustand 的 store 状态通过 dispatch 发生变化时，handleStoreChange 会被触发，进入状态检查流程。

图中的 checkIfSnapshotChanged：checkIfSnapshotChanged(inst) 会检查当前的 store 状态快照与之前保存的快照是否一致。图中展示了 getSnapshot 获取的最新快照与之前的 inst.value 进行对比，如果状态发生了变化，将返回 true。

如果 checkIfSnapshotChanged 返回 true，则通过 forceStoreRerender(fiber) 强制触发同步渲染。forceStoreRerender 会通过 scheduleUpdateOnFiber 来触发 SyncLane 模式下的同步更新。这保证了组件的状态与外部 store 状态保持一致，立即渲染最新的状态。

在并发模式下，React 在 mountSyncExternalStore 中调用了 pushStoreConsistencyCheck，来确保外部状态和组件渲染结果的一致性。

Fiber 节点的 updateQueue 中收集了 check 对象。在渲染结束时，React 会遍历这些 check 对象，确保 store 的状态没有发生变化。如果 store 状态与渲染时不一致，将通过 renderRootSync 进行同步渲染，以保证一致性。

当组件卸载或 useEffect 的依赖项发生变化时，useEffect 的清理函数会被调用，移除之前绑定的监听器。当依赖项或组件生命周期变化时，effect 的清理函数会移除之前添加的 listener，从而防止内存泄漏。这也是 subscribe 方法返回清理函数的原因。

小结

最终流程总结：

  1. 状态快照获取与存储：通过 getSnapshot 获取外部状态 state.bears，并将其存储在 hook.memoizedState 中。

  2. 监听器添加：useEffect 中执行 subscribeToStore，将 handleStoreChange 作为监听器添加到 zustand store 中。

  3. 状态变化检测：当 store 状态变化时，handleStoreChange 会通过 checkIfSnapshotChanged 检测是否发生状态改变。

  4. 同步阻塞渲染：如果状态改变，则通过 forceStoreRerender 强制组件同步重新渲染，以保持 UI 与状态一致。

  5. 一致性检查：在并发模式下，React 在渲染结束后通过 pushStoreConsistencyCheck 进行状态一致性检查，确保状态与 UI 同步。

  6. 清理机制：当组件卸载或依赖项变化时，useEffect 清理函数会移除之前的 listener，防止重复订阅和内存泄漏。

参考资料

总结

useSyncExternalStore 是 React 18 引入的一个钩子，用于安全地与外部状态管理系统（如 Redux 或 Zustand）交互。它通过订阅外部状态的变化，确保组件在渲染时始终读取最新的状态快照，并在状态变化时强制组件重新渲染。该钩子设计考虑了并发模式的特性，在协调完成后会进行一致性检查，确保 UI 与外部状态一致。通过结合 subscribe 和 getSnapshot 函数，useSyncExternalStore 提供了一个可靠的机制来管理外部状态在 React 应用中的一致性。

结合 useSyncExternalStore，zustand 通过内部的 useStore 方法实现外部状态与 React 组件的同步。具体来说，zustand 利用 useSyncExternalStore 来订阅 store 状态的变化并获取状态快照，从而保证组件渲染时能够读取最新的状态。通过 subscribe 函数，zustand 在状态发生变化时通知 React，触发 useSyncExternalStore 强制组件重新渲染。这样，zustand 提供了一个高效的状态管理解决方案，结合 useSyncExternalStore 来确保在并发渲染模式下状态与 UI 的一致性。

zustand 的执行流程可以简化为以下步骤：

  1. 创建 store：调用 create 函数生成一个 store，这个 store 实际上是一个自定义的 React Hook，它基于 zustand 提供的 API 对象和 React 官方的 useSyncExternalStore Hook 实现。

  2. API 对象：store 返回的 API 对象包含用于管理状态的几个关键方法：更新状态的 setState、获取状态的 getState、添加监听器的 subscribe，以及清除所有监听器的 destroy。

  3. 状态 state：zustand 的状态由用户在创建 store 时定义的 createState 函数生成，当调用 create(createState) 时，state 初始化为用户自定义的初始状态。

  4. 订阅的时机：subscribe 方法在 useEffect 内部执行，用于添加 listener。尽管没有显式调用 useEffect，zustand 通过 useSyncExternalStore 的内部机制实现了对状态变化的订阅和清理。

  5. 状态更新与渲染触发：当调用 setState 更新状态时，zustand 会遍历所有已注册的 listener 并触发它们。每个 listener 实际上就是 handleStoreChange，它会检查前后状态是否变化，如果变化，则通过 useSyncExternalStore 触发组件重新渲染。

最后分享两个我的两个开源项目,它们分别是:

这两个项目都会一直维护的，如果你想参与或者交流学习，可以加我微信 yunmz777 如果你也喜欢,欢迎 star 🚗🚗🚗

avatar
文章
阅读
粉丝