likes
comments
collection
share

从零开始写一个简单高效的 React 状态管理库！

作者站长头像
站长
· 阅读数 2

前言

React 社区拥有丰富的状态管理库，但是多样的选择也给开发者造成了一定的困扰：选择一个适合自己的状态管理库对于像笔者这样，拥有选择困难症来的开发者来说是一件痛苦万分的事情。所以笔者萌生这样的想法：与其面对痛苦的选择，不如自己写一个符合自己要求的状态管理库。

在写库之前，先明确需求。对于笔者而言，一个好的状态管理库有3点要求：

  1. 简单：概念简单，API 数量少。
  2. 好用：支持 TS 并且类型完备。
  3. 高效：性能要好，包体积要小。

想法

首先，社区中状态管理库主要是解决以下两项问题：

  1. 组件间状态共享（避免 Props 层层传递）
  2. 状态的集中管理（逻辑层抽离）

而 React 本身就提供了这两项功能： React Hooks 可以提供抽离逻辑和集中管理状态的功能，Context 可以提供组件间状态共享的功能。基于「概念简单」的要求，笔者不想额外引入其他概念。因此状态管理库的核心概念可以基于 React 自身提供的 Context + Hooks，即用户通过 自定义 Hooks 定义状态和方法，然后通过 Context.Provider 提供给组件，组件内部通过 useContext 进行消费。基于以上的想法，状态管理库大致的形状如下：

1、通过给 createStore 函数传入一个 自定义Hook 参数，创建 Store。

// 通过传入一个自定义 hook 创建一个 Store
const CounterStore = createStore(() => {
  const [count, setCount] = useState(0);
  const increase = () => {
    setCount((v) => v + 1);
  };
  return {
    count,
    increase,
  };
});

2、通过 Store.Provider 提供。

const App = () => {
  // 提供 Store
  return (
    <CounterStore.Provider>
      <Child1 />
      <Child2 />
    </CounterStore.Provider>
  );
};

3、通过 Store.useStore 使用。

const Child1 = () => {
  // 获取 Store 数据
  const { count } = CounterStore.useStore();
  return <div>{count}</div>;
};
const Child2 = () => {
  // 获取 Store 数据
  const { increase } = CounterStore.useStore();
  return <button onClick={increase}>Increase</button>;
};

实现

实现 createStore 的原理很简单，只需按照以下步骤：

  1. 创建 Context。
  2. 定义 Provider，执行 自定义 Hook，然后将执行的结果 value 传给 Context.Provider。
  3. 定义 useStore，通过 useContext(Context) 获取 value。

具体的代码如下：

const createStore = <Value, Props>(useHook: (props?: Props) => Value) => {
  // 创建 Context
  const Context = createContext<Value>(null);
  // 定义 Provider
  const Provider: FC<PropsWithChildren<{ props?: Props }>> = ({ props, children }) => {
    // 执行自定义 Hook 并获取结果 value
    const value = useHook(props as Props);
    // 将 value 传给 Context.Provider
    return <Context.Provider value={value}>{children}</Context.Provider>;
  };
  // 定义 useStore
  const useStore = () => {
    // 通过 useContext 获取 value
    const value = useContext(Context);
    return value;
  };
  return { Provider, useStore };
};

至此，我们已经基本实现了一个简单的 React 状态管理库。回到文章开头的目标：简单（概念简单），好用（类型完备）已经基本满足，但是高效（渲染性能），受限于 React Context 自身的性能问题，还没有实现。所以，下文我们继续探索一下 Context 存在什么性能问题，以及我们该如何解决。

优化

Context 的性能问题

React Context 的性能问题在于：每次 Context 的状态 value 更新时，所有订阅该 Context 的组件都会重新全量渲染，即使组件中使用的状态并未发生改变。因此在状态更新时，Context 会造成很多组件的无效重复渲染。

由于我们的状态管理库是基于 Context 的，所以也会面临相同的性能问题：

  1. 在每次更新状态时，React 会重新执行 自定义Hook 返回一个新的对象 value；
  2. Context.Provider 每次都接受新的 value 对象，所有订阅这个 Context 的组件都跟着一起更新，造成无意义的重复渲染。

React 通过 Object.is 方法对新旧 value 对象进行「浅比较」，如果相等则不触发更新，如果不等则触发更新。

用我们之前写好的状态管理库举个栗子：

import { useCallback, useState } from "react";
import { createStore } from "./create-store";
// 创建 Store
const CounterStore = createStore(() => {
  const [count, setCount] = useState(0);
  const increase = useCallback(() => {
    setCount((v) => v + 1);
  }, []);
  const decrease = useCallback(() => {
    setCount((v) => v - 1);
  }, []);
  return {
    count,
    increase,
    decrease,
  };
});

const Child1 = () => {
  // 使用 Store
  const { count } = CounterStore.useStore();
  return <div>count:{count}</div>;
};

const Child2 = () => {
  // 使用 Store
  const { increase } = CounterStore.useStore();
  return <button onClick={Increase}>increase</button>;
};

const Child3 = () => {
  // 使用 Store
  const { decrease } = CounterStore.useStore();
  return <button onClick={Decrease}>decrease</button>;
};

const App = () => {
  // 提供 Store
  return (
    <CounterStore.Provider>
      <Child1 />
      <Child2 />
      <Child3 />
    </CounterStore.Provider>
  );
};

export default App;

上述例子中，我们创建一个 CounterStore 并提供给 Child 组件们使用，其中 increase 和 decrease 方法使用了 useCallback 进行优化，使其引用不会发生改变。对于子组件：

  • Child1 使用 count 状态；
  • Child2 使用 increase 方法；
  • Child3 使用 decrease 方法；

那么我们在点击 increase 按钮更新 count 时，因为只有 count 的值变动了，所以理想情况是只渲染使用了 count 的 Child1。但是实际情况如下图： 点击 Increase 按钮后，Child1、Child2、Child3 都重新渲染了。（图中使用了 React Devtools Profiler 的火焰图）

从零开始写一个简单高效的 React 状态管理库！

性能优化的思路与实现

根据上述对 Context 渲染机制的分析：由于 自定义Hook 每次都生成新的对象 value 传给 Context.Provider，导致 React 每次都会重新渲染订阅 Context 的子组件，没有办法根据组件的需要实现「按需更新」。如下图示意： 从零开始写一个简单高效的 React 状态管理库！ 那么，有什么方法可以避免上述问题呢？或许我们可以换一个思路，既然 Context 的更新机制无法实现按需更新，那我们可以自己实现一个 「发布-订阅」 事件系统（通过维护一个事件列表实现）：子组件「订阅」所需状态并注册事件，当更新时事件系统再将更新「派发」到子组件内，以实现按需更新，如下图示意： 从零开始写一个简单高效的 React 状态管理库！ 整体思路有了，接下来我们来具体实现：

1、传递

首先我们要解决的问题是：如何传递「事件列表」和「状态」且不触发更新？

我们可以将「事件列表」和「状态」放在一个对象中，然后通过 useRef 将对象的引用恒定化，再将该对象传递给 Context.Provider 来避免 React 触发更新。

interface ContextObject<Value> {
  // 状态
  value: Value;
  // 事件列表
  events: Set<(value: Value) => void>;
}
  // 创建 Context
  const Context = createContext<ContextObject<Value>>({} as unknown as ContextObject<Value>);
  function Provider({ children, props }: PropsWithChildren<{ props?: Props }>) {
    // 获取状态
    const value = useHook(props as Props);
    // 创建 context ：包含「事件列表」和「状态」的对象（恒定引用）
    const context = useRef<ContextObject<Value>>({ value, events: new Set() }).current;
    // 传递 context 对象（恒定引用）
    return <Context.Provider value={context}>{children}</Context.Provider>;
  }

2、派发

在 自定义Hook 更新时，React 会重新执行 Provider 函数，因此我们可以借此更新「状态」和「派发」事件；

 function Provider({ children, props }: PropsWithChildren<{ props?: Props }>) {
    const value = useHook(props as Props);
    const context = useRef<ContextObject<Value>>({ value, events: new Set() }).current;
    // 更新「状态」
    context.value = value;
    useEffect(() => {
      // 每次更新时「派发」事件（派发事件对于 React 属于「外部效应」，需要放在 useEffect 内执行）
      context.events.forEach((event) => {
        event(value);
      });
    });
    return <Context.Provider value={context}>{children}</Context.Provider>;
  }

3、订阅与注册

关于订阅，首先需要知道组件订阅了哪些状态。我们可以像 redux 那样：让用户通过 Store.useStore 传入一个 selector 函数来订阅组件需要的状态。如下所示：

const Child1 = () => {
  // selector 函数：(value) => value.count
  // 订阅 count 状态
  const count = CounterStore.useStore((value) => value.count);
  return <div>{count}</div>;
};

接下来我们来看订阅与注册的具体实现：

  1. 首先需要在组件挂载时注册事件，并在组件卸载时注销事件；
  2. 在「事件」内部：需要获取「当前（订阅）状态」和「下次（订阅）状态」，判断是否需要更新组件；
  3. 我们可以在组件内部获取「当前（订阅）状态」，但是需要通过 ref 传递到事件内部，避免闭包陷阱；
  4. 组件内部的更新，我们可以通过 useReducer 触发组件的自更新。
  function useStore<SelectedValue = Value>(selector?: (value: Value) => SelectedValue): SelectedValue {
    // 组件自更新触发器
    const [, forceUpdate] = useReducer(() => ({}), {});
    // 获取「状态」和「事件列表」
    const { value, events } = useContext(Context);
    // 获取「订阅」的状态（如果有的话）
    const selectedValue = selector ? selector(value) : value;
    // 「当前状态」
    const current = {
      value,
      selectedValue,
      events,
      selector,
    };
    // 通过 ref 将「当前状态」对象传递到事件内部，避免闭包陷阱
    const ref = useRef(current);
    // 更新「当前状态」对象
    ref.current = current;
    useEffect(() => {
      // 创建「事件」
      function event(nextValue: Value) {
        // 获取「当前状态」和「当前订阅状态」
        const { value, selectedValue, selector } = ref.current;
        // 如果「当前状态」与「下次状态」相同，不更新组件
        if (value === nextValue) return;
        // 「下次订阅状态」
        const nextSelectedValue = selector ? selector(nextValue) : nextValue;
        // 如果「当前订阅状态」与「下次订阅状态」相同，不更新组件
        if (selectedValue === nextSelectedValue) return;
        // 走到这里，说明当前状态和下次状态不同，更新组件
        forceUpdate();
      }
      // 组件挂载时注册事件
      ref.current.events.add(event);
      return () => {
        // 组件卸载时注销事件
        ref.current.events.delete(event);
      };
    }, []);
    // 返回订阅状态
    return selectedValue as SelectedValue;
  }

完整代码实现如下：

import { PropsWithChildren, createContext, useContext, useEffect, useReducer, useRef } from "react";

interface ContextObject<Value> {
  // 状态
  value: Value;
  // 事件列表
  events: Set<(value: Value) => void>;
}

export function createStore<Value, Props>(useHook: (props: Props) => Value) {
  // 创建 Context
  const Context = createContext<ContextObject<Value>>({} as unknown as ContextObject<Value>);
  function Provider({ children, props }: PropsWithChildren<{ props?: Props }>) {
    // 状态
    const value = useHook(props as Props);
    // 包含「事件列表」和「状态」的对象（恒定引用）
    const context = useRef<ContextObject<Value>>({ value, events: new Set() }).current;
    // 更新「状态」
    context.value = value;
    useEffect(() => {
      // 每次更新时「派发」事件
      context.events.forEach((event) => {
        event(value);
      });
    });
    return <Context.Provider value={context}>{children}</Context.Provider>;
  }
  function useStore<SelectedValue = Value>(selector?: (value: Value) => SelectedValue): SelectedValue {
    // 自更新触发器
    const [, forceUpdate] = useReducer(() => ({}), {});
    // 获取「状态」和「事件列表」
    const { value, events } = useContext(Context);
    // 获取「订阅」的状态（如果有的话）
    const selectedValue = selector ? selector(value) : value;
    // 「当前状态」
    const current = {
      value,
      selectedValue,
      events,
      selector,
    };
    // 通过 ref 将「当前状态」对象传递到事件内部，避免闭包陷阱
    const ref = useRef(current);
    // 更新「当前状态」对象
    ref.current = current;
    useEffect(() => {
      // 创建「事件」
      function event(nextValue: Value) {
        // 获取「当前状态」和「当前订阅状态」
        const { value, selectedValue, selector } = ref.current;
        // 如果「当前状态」与「下次状态」相同，不触发更新
        if (value === nextValue) return;
        // 「下次订阅状态」
        const nextSelectedValue = selector ? selector(nextValue) : nextValue;
        // 如果「当前订阅状态」与「下次订阅状态」相同，不触发更新
        if (selectedValue === nextSelectedValue) return;
        // 走到这里，说明当前状态和下次状态不同，触发组件自更新
        forceUpdate();
      }
      // 组件挂载时注册事件
      ref.current.events.add(event);
      return () => {
        // 组件卸载时注销事件
        ref.current.events.delete(event);
      };
    }, []);
    // 返回订阅状态
    return selectedValue as SelectedValue;
  }
  return { Provider, useStore };
}

优化成果

最后我们来验证一下性能优化的效果：

import { useState, useCallback } from "react";
import { createStore } from "./create-store";

const CounterStore = createStore(() => {
  const [count, setCount] = useState(0);
  const increase = useCallback(() => {
    setCount((v) => v + 1);
  }, []);
  const decrease = useCallback(() => {
    setCount((v) => v - 1);
  }, []);
  return {
    count,
    increase,
    decrease,
  };
});

const Child1 = () => {
  const count = CounterStore.useStore((value) => value.count);
  return <div>{count}</div>;
};

const Child2 = () => {
  const increase = CounterStore.useStore((value) => value.increase);
  return <button onClick={increase}>Increase</button>;
};

const Child3 = () => {
  const decrease = CounterStore.useStore((value) => value.decrease);
  return <button onClick={decrease}>Decrease</button>;
};

const App = () => {
  return (
    <CounterStore.Provider>
      <Child1 />
      <Child2 />
      <Child3 />
    </CounterStore.Provider>
  );
};

export default App;

可以看到，点击 Increase 按钮后，只有 Child1 重新渲染了，Child2、Child3 都未渲染。成功避免了 Context 导致的重复渲染问题。（图中使用了 React Devtools Profiler 的火焰图。） 从零开始写一个简单高效的 React 状态管理库！

总结

最后我们来看一下笔者最初制定的目标有没有达成：

  1. 简单：
    1. 概念简单：通过 自定义Hook + Context 实现，全部都是 React 自身概念。（✅）
    2. API数量少：只有 createStore 一个API。（✅）
  2. 好用：全部使用 TypeScript 实现，类型完备。（✅）
  3. 高效：
    1. 性能好：通过 selector函数和「发布-订阅」事件系统，避免 Context 导致的重复渲染性能问题。（✅）
    2. 包体积小：总共 60 行代码，mini bundle 打包仅 600B 大小。（✅）

最终笔者的目标全部达成！

根据本文思路实现的状态管理库已发布 hostore（起名含义为 hooks + store），完整源码可见：github/hostore（觉得有帮助的可以点个 star 🙏）。

参考文章