谈谈 React Context 的性能优化

在 react 应用中，大家一定使用过 Context 来管理状态以解决 props drilling 等问题。在使用 Context 时也会带来一些问题，下面内容给大家介绍在 react 应用中，如何更好地使用 Context 来管理状态以及解决使用它带来的重复渲染问题。

基础使用 Context

通过 react 官方对 Context 的定义，我们了解 Context 为我们提供了一种，当需要在组件中多层传递我们 props 时，可以避免每一层都向下传递同样的 props。像下面 👇🏻

像上图中的组件结构，或许需要如下代码示例：(不使用 Context 时)

// App
const [query, setQuery] = useState({
  name: '',
  team: '',
  age: undefined,
  score: undefined,
});

const handleChange = (val) => {
  setQuery({ ...query, ...val });
};

// FormContainer
<FormContainer value={query} onChange={handleChange} />;
// ...

可以看的出，同样的东西传递了很多层，写法既繁琐也不利于后期维护。

当使用了 Context 时，我们则不必每一层传递 props，而是通过 Context 直接在对应的组件内直接获取和修改即可。

示例：（使用 Context）

// App
<AppContext.Provider value={state}>
  <AppContainer />
</AppContext.Provider>;

// 这样子我们可以通过useContext, 直接获取和修改我们需要的

// SearchForm
const [query, setQuery] = useContext(AppContext);

<Select
  value={query.age}
  onChange={(val) => {
    setQuery({ age: val });
  }}
/>;

Context 重复渲染问题

上面 Context 的使用减少了我们传递 props 的繁琐，但是会引入一个不必要的重复渲染问题。

当我们通过 Context 修改了 App 组件中的 state 时，整个组件树自上而下都会重新渲染，也就是说很多组件没有依赖到我们 Context 中的 state，但是在我们修改 Context 的值时，却也引起了这些组件的 渲染。

我们可以通过 react devtools 高亮，来观察重新渲染的组件，如下图：

可以看出当修改 name 的值时，相应的会修改 Context 中 state，但是那些没有依赖的组件 （图 1 中未标红的组件）也发生了重新渲染，这个是一个性能牺牲，我们不答应。

更好地应用 Context

要解决 Context 重复渲染问题，也就是当修改 Context 的值时，我们只要重新渲染那些依赖 Context 的组件即可，不需要从顶层的整个组件树都渲染一遍。

也就是说我们需要一个数据源，然后组件可以订阅数据源的变化，数据源变化时，我们通知订阅的组件变化，触发组件重新渲染即可。

所以说我们可以简单实现一个发布订阅模型即可。

发布订阅

一个简单的发布订阅模型如下所示：

export type Listener<T> = (state: T) => void;

export type Store<T> = {
  setState: (partial: Partial<T>) => void;
  getState: () => T;
  subscribe: (listener: Listener<T>) => () => void;
};

export function createStore<T>(initState: T): Store<T> {
  let state = initState;

  const listeners = new Set<Listener<T>>();

  const setState = (partial: Partial<T>) => {
    state = { ...state, ...partial };

    listeners.forEach((listener) => {
      listener(state);
    });
  };

  const getState = () => state;

  const subscribe = (listener: Listener<T>) => {
    listeners.add(listener);
    return () => listeners.delete(listener);
  };

  return {
    getState,
    setState,
    subscribe,
  };
}

接下来我们需要把它应用到 react 的应用中。首先我们可以通过 Context 将这个初始化 store 传递下去；在每个需要的组件中，我们可以通过（getState）获取我们需要的数据和通过（setState）来修改这个数据，然后通过（subscribe）来订阅这个 store 的数据源变化，然后当数据源变化时各个订阅组件来重新渲染即可。

初始化 store，作为 Context 传递

export type QueryState = {
  name?: string;
  team?: string;
  age?: string;
  score?: string;
};

export const AppContext = createContext<Store<QueryState> | null>(null);

const initState: QueryState = {};
const store = createStore(initState);

function App() {
  return (
    <AppContext.Provider value={store}>
      <AppContainer />
    </AppContext.Provider>
  );
}

组件中添加订阅，由于组件订阅的逻辑基本一致，那我们可以定义一个 hook 来实现这部分逻辑

const useAppStore = (): [QueryState, (p: Partial<QueryState>) => void] => {
  const storeCtx = useContext(AppContext)!;

  const [state, setState] = useState(storeCtx.getState());

  // add subscribe

  useEffect(() => {
    const unsubscribe = storeCtx.subscribe((s: QueryState) => {
      setState(s);
    });

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

  // 如果在 react-18中 我们可以使用新的api useSyncExternalStore
  // useSyncExternalStore(store.subscribe, store.getState);
  // https://zh-hans.reactjs.org/docs/hooks-reference.html#usesyncexternalstore
  return [state, storeCtx.setState];
};

组件获取或者修改

使用 useAppStore 这个 hook，从中订阅数据的变化。

function SearchForm() {
  const [query, setQuery] = useAppStore();
  // ......省略
}

按照类似的逻辑，修改之前的实现。最后我们查看经过优化后，我们的应用渲染的效果。

可以通过高亮看到，差不多已经达到了我们预期的思考，只有依赖数据源的组件发生了渲染，其他的依旧如初，很好但是还可以更好！

性能优化

在我们的 demo 中 AgeSelect 和 ScoreSelect 只需要 query 中个别字段，我们可以更进一步实现一下我们获取 state 的方法，从 state 中只获取我们需要的。可以加入一个 selector function 选取需要的数据。

const useAppStore = (
  selector?: (state: QueryState) => any,
): [Partial<QueryState>, (p: Partial<QueryState>) => void] => {
  const storeCtx = useContext(AppContext)!;

  const defaultSelectState = selector
    ? selector?.(storeCtx.getState())
    : storeCtx.getState();

  const [state, setState] = useState(defaultSelectState);

  // add subscribe
  useEffect(() => {
    const unsubscribe = storeCtx.subscribe((s: QueryState) => {
      const selectState = selector ? selector(s) : s;
      setState(selectState);
    });

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

  return [state, storeCtx.setState];
};

在调用 useAppStore 时传入 selector，从而做到只有当 score 变化时，我们[只关注 score]的组件才会重新渲染，其他的变化不会引起渲染，又做到了进一步优化。

const ScoreSelect = () => {
  const [score] = useAppStore((state) => state.score);
  // ...
};

渲染效果如下：

最后重新组织代码

// createStoreContext
import { createStore, Store } from './createStore';

function createStoreContext<T>(initState: T) {
  const StoreContext = createContext<Store<T> | null>(null);

  const StoreProvider = ({ children }: PropsWithChildren) => {
    const storeRef = useRef<Store<T>>();
    if (!storeRef.current) {
      storeRef.current = createStore(initState);
    }

    return (
      <StoreContext.Provider value={storeRef.current}>
        {children}
      </StoreContext.Provider>
    );
  };

  const useStore = (
    selector?: (state: T) => any,
  ): [Partial<T>, (p: Partial<T>) => void] => {
    const storeCtx = useContext(StoreContext)!;

    const defaultSelectState = selector
      ? selector?.(storeCtx.getState())
      : storeCtx.getState();

    const [state, setState] = useState(defaultSelectState);

    // add subscribe
    useEffect(() => {
      const unsubscribe = storeCtx.subscribe((s: T) => {
        const selectState = selector ? selector(s) : s;

        setState(selectState);
      });

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

    return [state, storeCtx.setState];
  };

  return {
    StoreProvider,
    useStore,
  };
}

//  AppContext
export type QueryState = {
  name?: string;
  team?: string;
  age?: string;
  score?: string;
};

const initState: QueryState = {
  name: '',
  team: '',
  age: undefined,
  score: undefined,
};

const { StoreProvider: AppContextProvider, useStore: useAppStore } =
  createStoreContext(initState);

export { AppContextProvider, useAppStore };

// App
import { AppContextProvider, useAppStore } from './AppContext';

function App() {
  return (
    <AppContextProvider>
      <AppContainer />
    </AppContextProvider>
  );
}

// other
const ScoreSelect = () => {
  const [score, setQuery] = useAppStore((state) => state.score);
  //...省略
};

总结

我们通过简单实现一个发布订阅模型，通过 Context 传递下去，然后在每个需要的组件中，通过 getter 方法将需要的 state 映射到 react 组件的状态上，然后通过添加订阅，在相应的数据更新后，即可引起重新渲染以达到我们的目的，并且也减少了重复渲染的问题。一些状态管理库，例如 zustand 的核心原理就与本文示例代码差不多。