React Context 性能优化

引言

跨层级传递数据，通常会使用 React Context 作为媒介，hooks 提出之后，使用 context 前所未有的方便，代价也同样浮出水面。context payload 变更，会触发所有 consumer component 重新渲染，无论组件依赖的数据是否发生变化，仓库讨论参见 github.com/facebook/re…

案例分析

举个简单例子，假设需要持续跟踪 viewport 尺寸，用户点击次数，用户键盘敲击次数，基本数据结构如下：

interface Statistic {
  click: number;
  keypress: number;
}

interface Dimension {
  width: number;
  height: number;
}

type CompoundContextStructure = Statistic & Dimension;

数据置于一处管理，Provider 实现如下：

export const Provider: FunctionComponent = (props) => {
  const [state, setState] = useState(defaultCompoundContext);

  useEffect(() => {
    const click$ = fromEvent(document, 'click').pipe(
      scan((acc) => acc + 1, 0),
      map((click) => ({ click }))
    );
    const keypress$ = fromEvent(document, 'keypress').pipe(
      scan((acc) => acc + 1, 0),
      map((keypress) => ({ keypress }))
    );
    const dimension$ = fromEvent(window, 'resize').pipe(
      map(() => ({
        width: window.innerWidth,
        height: window.innerHeight,
      }))
    );

    const subscription = merge(click$, keypress$, dimension$).subscribe(
      (action) => {
        setState((acc) => ({
          ...acc,
          ...action,
        }));
      }
    );

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

  return (
    <CompoundContext.Provider value={state}>
      {props.children}
    </CompoundContext.Provider>
  );
};

Consumer 直接使用 useContext，实现如下：

function Device() {
  const { width, height } = useContext(CompoundContext);
  const message = `The browser viewport width: ${width}, height: ${height}, Celebration!`;

  return <Alert message={message} />;
}

function Statistics() {
  const { click, keypress } = useContext(CompoundContext);

  return (
    <Fragment>
      <Statistic title="Active Clicks" value={click} />
      <Statistic title="Active Keypress" value={keypress} />
    </Fragment>
  );

初步效果如下：

组件能够实时呈现用户操作数据，使用功能达成。使用 react devtool 进行 profile，鼠标点击两次，火焰图如下：

鼠标点击操作，仅更新 clicks 字段，最优解来说，应该仅触发 Statistics 组件重新渲染，但实际情况也触发 Device 组件重新渲染，暨引言所述的非必要渲染问题。

如果 context 承载数据量较大，consumer 重新渲染的成本不可忽视，需要进行优化。

优化方案

Context 拆分

从设计角度，context 变更本应该触发 consumer 重新渲染。实际情况中，如果不打算对实现大刀阔斧改动，为避免不必要的渲染，可以选择将相关数据置于同一 context 下，不相关数据置于不同 context 下，就能够达成渲染最小化目的。案例中状态，拆分为 DimensionContext、StatisticContext：

export const DimensionProvider: FunctionComponent = (props) => {
  const [state, setState] = useState(defaultDimensionContext);

  useEffect(() => {
    const dimension$ = fromEvent(window, 'resize').pipe(
      map(() => ({
        width: window.innerWidth,
        height: window.innerHeight,
      })),
    );

    const subscription = dimension$.subscribe((action) => {
      setState(action);
    });

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

  return (
    <DimensionContext.Provider value={state}>
      {props.children}
    </DimensionContext.Provider>
  );
};

export const StatisticProvider: FunctionComponent = (props) => {
  const [state, setState] = useState(defaultStatisticContext);

  useEffect(() => {
    const click$ = fromEvent(document, 'click').pipe(
      scan((acc) => acc + 1, 0),
      map((click) => ({ click })),
    );
    const keypress$ = fromEvent(document, 'keypress').pipe(
      scan((acc) => acc + 1, 0),
      map((keypress) => ({ keypress })),
    );

    const subscription = merge(click$, keypress$).subscribe(
      (action) => {
        setState((acc) => ({
          ...acc,
          ...action,
        }));
      },
    );

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

  return (
    <StatisticContext.Provider value={state}>
      {props.children}
    </StatisticContext.Provider>
  );
};

consumer 代码高度类似，粘贴如下：

function Statistics() {
  // Context 引用变更
  const { click, keypress } = useContext(StatisticContext);

  return (
    <Fragment>
      <Statistic title="Active Clicks" value={click} />
      <Statistic title="Active Keypress" value={keypress} />
    </Fragment>
  );
}

function Device() {
  // Context 引用变更
  const { width, height } = useContext(DimensionContext);
  const message = `The browser viewport width: ${width}, height: ${height}, Celebration!`;

  return <Alert message={message} />;
}

重复用户操作，火焰图如下：

统计数据变更不会影响 Device 组件重新渲染，目标达成。

Memoise

拆分 context 需要将管理逻辑分散，维护起来略显麻烦，增加开发成本。如果重渲染不可避免，那么将核心渲染进行缓存，也可以避免非必要资源消耗。 React 自带 React.memo、React.useMemo 记忆函数，使用上没有实质上的差别。

Provider 实现同原始版本实现，不再赘述，Consumer 代码如下：

function Device() {
  const { width, height } = useContext(CompoundContext);
  
  return useMemo(() => {
    const message = `The browser viewport width: ${width}, height: ${height}, Celebration!`;

    return <Alert message={message} />;
  }, [width, height]);
}

function Statistics() {
  const { click, keypress } = useContext(CompoundContext);

  return useMemo(() => {
    return (
      <Fragment>
        <Statistic title="Active Clicks" value={click} />
        <Statistic title="Account Balance (CNY)" value={keypress} />
      </Fragment>
    );
  }, [click, keypress]);
}

重复点击操作，火焰图如下：

底层的 Alert 组件没有重新渲染，目的达成。

订阅模式

前述方案简单粗暴，虽然能达成目的，感觉上多了份粗暴，少了份优雅，有没有更好的方案？

设计上来说，context payload 变更触发重新渲染，如果将原始 state 封装，且维持封装的稳定性，状态变更由封装对象推送到 consumer 端，由 consumer 端判断是否触发渲染。实质上来说，就是覆写 React 的变更分发逻辑，实现方式如下：

export const Provider: FunctionComponent = (props) => {
  const initialState = useMemo(
    () => ({
      click: 0,
      keypress: 0,
      width: window.innerWidth,
      height: window.innerHeight,
    }),
    []
  );
  const payload$ = useMemo(
    () => new BehaviorSubject<CompoundContextStructure>(initialState),
    []
  );

  useEffect(() => {
    const click$ = fromEvent(document, 'click').pipe(
      scan((acc) => acc + 1, 0),
      map((click) => ({ click }))
    );
    const keypress$ = fromEvent(document, 'keypress').pipe(
      scan((acc) => acc + 1, 0),
      map((keypress) => ({ keypress }))
    );
    const dimension$ = fromEvent(window, 'resize').pipe(
      map(() => ({
        width: window.innerWidth,
        height: window.innerHeight,
      }))
    );

    const subscription = merge(click$, keypress$, dimension$)
      .pipe(scan((acc, curr) => ({ ...acc, ...curr }), initialState))
      .subscribe((state) => payload$.next(state));

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

  return (
    <CompoundContext.Provider value={payload$}>
      {props.children}
    </CompoundContext.Provider>
  );
};

为保持 consumer 组件简洁性，需要实现自定义 hook 隐藏订阅逻辑，实现如下：

type Selector = (
  structure: CompoundContextStructure
) => Partial<CompoundContextStructure>;
type Compare = (prev: any, curr: any) => boolean;

export function useSelector(
  selector: Selector,
  compare: Compare = shallowEqual
) {
  const payload = useRef<Partial<CompoundContextStructure>>({});
  const payload$ = useContext(CompoundContext);
  const [, rerender] = useState(0);

  useEffect(() => {
    const subscription = payload$
      .pipe(
        map((state) => selector(state)),
        distinctUntilChanged((previous, current) => compare(previous, current))
      )
      .subscribe((state) => {
        payload.current = state;
        rerender((prev) => prev + 1);
      });

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

  return payload.current;
}

火焰图如下：

自定义 hook 名为 useSelector，眼熟的肯定能发现 react-redux 中也包含 useSelector hook。实际上，react-redux，react-mobx 同样采用订阅模式，差别在于对原始状态的封装方式，开发过程中推荐直接使用成熟类库，不要自造轮子。

注册模式

基本上，三板斧已经能解决重复渲染的问题，更进一步，consumer 参与渲染的数据可能有变化，若开发者懒癌附身，不愿意不断修改 selector，就喜欢随意使用字段、自动跟踪，体验跟使用原始对象一样顺滑，那该如何处理？

场景之下，可以采用注册模式。所谓注册模式，与订阅模式相似，核心差别在于重渲染判断逻辑由 provider 组件负责实现，consumer 端维持原始状态的可访问。实现方式上，采用拦截读写操作，拦截读操作，负责收集组件依赖，拦截写操作，负责触发渲染，实现起来有点绕，代码如下：

export class Provider extends Component<any, CompoundContextStructure> {
  private subscription: Subscription | undefined;
  private trackers: Tracker[];
  private readonly queues: Set<string>;
  private readonly payload: {
    [p: string]: Level1Proxy<CompoundContextStructure>;
  };

  constructor(props: {}) {
    super(props);

    this.state = {
      // statistics
      click: 0,
      keypress: 0,
      // dimension
      width: window.innerWidth,
      height: window.innerHeight,
    };

    // proxy related fields
    this.queues = new Set<string>();
    this.trackers = [];

    // private level1 proxy storage
    const store = new Map();
    const handlers: ProxyHandler<any> = {
      get: (_, identity: string) => {
        if (store.has(identity)) {
          return store.get(identity);
        } else {
          store.set(identity, this.createAccessStateProxy(identity));

          return store.get(identity);
        }
      },
    };

    // context wrapper
    this.payload = new Proxy({}, handlers);
  }

  createTracker(identity: string, accessKeys: string[], callback: () => void) {
    const tracker = (keys: string[]) => {
      keys.some((key) => {
        const match = accessKeys.includes(key);

        // tslint:disable-next-line:no-unused-expression
        match && callback();

        return match;
      });
    };

    tracker.identity = identity;

    return tracker as Tracker;
  }

  createAccessStateProxy(identity: string) {
    const cache: AccessStateProxyCache = {
      callback: () => {
        throw new Error('subscribe state callback function required');
      },
      accessKeys: [],
    };
    const subscribe = (render: () => void) => {
      // 添加订阅函数
      cache.callback = render;

      // reset empty keys
      cache.accessKeys = [];
    };
    // 更新 tracker
    const track = () => {
      this.trackers = this.trackers
        .filter((render) => render.identity !== identity)
        .concat(this.createTracker(identity, cache.accessKeys, cache.callback));
    };
    const accessStateProxyHandler: ProxyHandler<CompoundContextStructure> = {
      get: (_: any, property: string) => {
        // 记录 keys
        cache.accessKeys.push(property);

        // 返回原始值
        return Reflect.get(this.state, property);
      },
    };

    return {
      subscribe,
      track,
      // use this references directly, avoid target snapshot trap
      // @ts-ignore
      state: new Proxy({}, accessStateProxyHandler),
    };
  }

  componentDidMount() {
    const click$ = fromEvent(document, 'click').pipe(
      scan((acc) => acc + 1, 0),
      map((click) => ({ click }))
    );
    const keypress$ = fromEvent(document, 'keypress').pipe(
      scan((acc) => acc + 1, 0),
      map((keypress) => ({ keypress }))
    );
    const dimension$ = fromEvent(window, 'resize').pipe(
      map(() => ({
        width: window.innerWidth,
        height: window.innerHeight,
      }))
    );

    this.subscription = merge(click$, keypress$, dimension$).subscribe(
      (action) => {
        this.setState((acc) => {
          Reflect.ownKeys(action)
            .filter((key) => Reflect.get(action, key) !== Reflect.get(acc, key))
            .forEach((key) => this.queues.add(key as string));

          return {
            ...acc,
            ...action,
          };
        });
      }
    );
  }

  shouldComponentUpdate() {
    const changedKeys = Array.from(this.queues);

    // 清空队列
    this.queues.clear();
    // 执行订阅回调
    this.trackers.forEach((tracker) => tracker(changedKeys));

    return false;
  }

  componentWillUnmount() {
    this.subscription!.unsubscribe();
  }

  render() {
    return (
      <CompoundContext.Provider value={this.payload}>
        {this.props.children}
      </CompoundContext.Provider>
    );
  }
}

export function useTrackContext() {
  const identity = useMemo(
    () =>
      '_' +
      Math.random()
        .toString(36)
        .substr(2, 9),
    []
  );
  const [, rerender] = useReducer((acc) => acc + 1, 0);
  const context = useContext(CompoundContext);
  const payload = context[identity];

  // 订阅更新
  payload.subscribe(rerender);

  useEffect(() => {
    payload.track();
  });

  return payload.state;
}

组件的实现方式略有调整，代码如下：

function Device() {
  const payload = useTrackContext();

  const message = `The browser viewport width: ${payload.width}, height: ${payload.height}, Celebration!`;

  return <Alert message={message} />;
}

功能基本实现，火焰图如下：

目测效果与前述方案相同。

总结

私以为直接使用 useContext 的场景，在于少量数据的跨层级传递，脱离“小”场景，重渲染问题便不可忽视。社区状态管理方案 redux，mobx 重渲染问题不大，且可以传参调优，但上手成本不低，更适用于大批量数据管理。数据量不大不小的场景如何处理，就显得有点尴尬。

至于本文提及的方案，只推荐 split context 方案，一旦难以平衡过多的 context 与 consumer 维护成本，直接上 redux 、mobx，不要考虑太多，千万不要考虑太多。

上述提及的优化方式，具体的 benchmark 还没有实现，对性能的影响目测提升，程度尚未可知。仅作为了解，不建议生产环境使用。如果有其他想法，欢迎告知。示例所有代码，详见 github.com/huang-xiao-… 。