Context API 效能問題 - use-context-selector 解析

前言

最近經手的一個專案採用 React Hooks 與 Context API 實作類似 Redux 的狀態管理，也就是利用 useReducer、createContext 等 API 來實作全域的 Store 與 Dispatch Actions。

這樣做其實挺方便的，在狀態管理的流程上跟 Redux 的思維一樣，但設置上更為簡單。

不過有個問題是，ㄧ但任何 context 的值更新，所有使用 useContext 的 component 都會被通知到，並且進行 render，即便該 component 需要的 state 可能根本沒有變動?。

簡單看個範例(modified from here)：

從上圖中 devtool 中的 flamegraph 可以明顯看出當點選 Counter 時，TextBox 也會觸發 render，因為他們共享同一個 Context。

附上 codesandbox 供參考（另外，這邊提到的 render 主要是 VDOM 的 render，範例中為了凸顯效果，在其中放了 Math.random() 讓 DOM 一定會更新，否則實際上 TextBox 在值都不變的狀態下，DOM 是不會更新的）：

先不論頁面複雜時可能會有的潛在效能問題，光是想到會有這種無謂的 render，應該很多人就會覺得不舒服。

而實際上，Context API 一開始就不是拿來給你作用在更新頻率高的狀態上的。

官方文件雖然沒有明講這件事，但從他們給的範例圍繞在 theme 與 user data 就可略知一二，另外在 react-redux v6 版本推出時的討論中也有提到。

所以我們應該要就此打住，改回用 react-redux 嗎？

也不一定，創造出問題然後解決，就是工程師的職責啊，怎麼能逃避！

玩笑話，實務上當然自己斟酌，如果是公司內部專案或是你自己的 side project，當然是能多嘗試就嘗試，我並不覺得一昧遵守 best practice 是好的。

另外，官方團隊也是有意識到這件事情

We indeed have observed performance problems when propagating context to large trees. @joshcstory is doing some great research on how to make it better. We do have a plan, but it will require a significant refactor so it might take a while to land. https://t.co/gtpLEyfgU9

— Andrew Clark (@acdlite) April 14, 2020

並且在 RFC: Context selectors 中曾有蠻熱絡的討論，雖然依照現況來說沒有明確的計劃針對這個問題去做改善，但 RFCs 提出的概念已經有類似實作了，而今天我就是想要來解析ㄧ下到底是怎麼在不更動架構，利用現有 API 下去解決這個問題。

解決方法 - useContextSelector

除了在頁面不複雜的狀態下可以透過組合多個 context 來解決，同事找到的這套 lib - use-context-selector 實作了 RFCs 中的概念，提供了 selector 給 Context 使用。

以先前同樣的範例來看看使用後的效果：

從圖中的 flamegraph 可以看到，在一樣的操作下，TextBox 在所有的 commits 中都沒有被觸發 render，只有 Counter 有執行 render。

若是再仔細看一點，你也會發現，跟原本的版本比起來，Commits 數量多了一倍，並多了一個 Anonymous (memo) 的 component。

而這多出來的部分就是 use-context-selector 能 bail out of rendering 的原因，接下來我們就從程式碼來理解實作原理！

（題外話，bail out of rendering 是我在查詢相關資訊時，常常看到的句子，覺得是很貼切的描述，所以保留原文，加上我也找不到合適的中文翻譯...）

程式碼解析

use-context-selector 的程式碼很短，就 100 多行而已，所以要直接看也是 ok，但我一般都習慣先從 lib 的使用方式下手，觀察出我們應該先閱讀哪部分的程式碼。

我們只取上面範例中的 Counter 來觀察：

import {
  createContext,
  useContextSelector,
} from './use-context-selector';

const context = createContext(null);

const Counter = () => {
  const count = useContextSelector(context, (v) => v[0].count);
  const dispatch = useContextSelector(context, (v) => v[1]);
  return (
    <div>
      {Math.random()}
      <div>
        <span>Count: {count}</span>
        <button type="button" onClick={() => dispatch({ type: 'increment' })}>+1</button>
        <button type="button" onClick={() => dispatch({ type: 'decrement' })}>-1</button>
      </div>
    </div>
  );
};

const Provider = ({ children }) => {
  const [state, dispatch] = useReducer(reducer, initialState);
  return (
    <context.Provider value={[state, dispatch]}>
      {children}
    </context.Provider>
  );
};

const App = () => (
  <StrictMode>
    <Provider>
      <h1>Counter</h1>
      <Counter />
      <Counter />
    </Provider>
  </StrictMode>
);

跟我們一般使用 Context API 的方式相同，需要用 createContext 來創建 context，只不過這邊用到的並不是 React 原生的 createContext，而是 use-context-selector 提供的。

另外就是與一般 useContext 不同，在 Component 中使用 useContextSelector 來取得 context 中的 state 與 dispatch 函式（React.useReducer() 產生的）。

useContextSelector 很好理解，就是多傳一個 selector 參數進去選取我們需要的 context value，但為什麼這邊他要我們使用它提供的 createContext 呢？

看來關鍵就在這邊，所以我們直接先從 use-context-selector 中的 createContext 函式看起：

export const createContext = (defaultValue) => {
  // make changedBits always zero
  const context = React.createContext(defaultValue, () => 0);
  // shared listeners (not ideal)
  context[CONTEXT_LISTENERS] = new Set();
  // hacked provider
  context.Provider = createProvider(context.Provider, context[CONTEXT_LISTENERS]);
  // no support for consumer
  delete context.Consumer;
  return context;
};

可以看出他其實也是使用 React.createContext 來創建 Context，只是他多傳了一個參數進去。

🤔 什麼時候 React.createContext 有第二個參數選項了？

從上面的註解來看，傳入的第二個參數會回傳一個叫做 changedBits 的值，Google 一下後發現原來是沒有寫在文件上的 API，而且兩年前新的 Context API 出來時就已經有不少人在討論了（原來只是自己學識淺薄😅）

在先前提到的 RFC: Context selectors 中也是想要利用這個 API。

這第二個參數叫做 calculateChangedBits，他會接受 Context 的新值與舊值作為 input，最後 return changedBits，如果 changedBits 為 0，Context Provider 就不會觸發更新；而Context Consumer 中也能傳入一個叫做 unstable_observedBits 的 props，若是 unstable_observedBits & changedBits !== 0，Consumer 也不會更新。

雖然 observedBits 是 unstable 的，但在 react-reconciler 的 NewContext test 中，他們就是利用 changedBits 與 observedBits 來做更新的測試。

這邊再羅列幾篇講解得比較詳細的文章供大家參考：

• 不一樣的 React context
• A Secret parts of React New Context API
• React tips — Context API (performance considerations)

總而言之，我們是可以客製化一個函式來決定 Context 的值更動時，需不需要觸發更新。

但這個函式是在 createContext 時就得傳入的，而不是 useContext，我們的 Component 沒辦法動態去傳各自的 Selector。

也正是如此，use-context-selector 就直接以 () => 0 作為 calculateChangedBits 函式，讓 React Context Provider 拿到的 changedBits 永遠為 0。

這樣做會讓 Provider 永遠不會跟隨著 Context 變動而觸發 render，而是由我們自己來判斷何時要做更新，為此，use-context-selector 實作了另一個 context.Provider：

const createProvider = (OrigProvider, listeners) => React.memo(({ value, children }) => {
  if (process.env.NODE_ENV !== 'production') {
    // we use layout effect to eliminate warnings.
    // but, this leads tearing with startTransition.
    // eslint-disable-next-line react-hooks/rules-of-hooks
    React.useLayoutEffect(() => {
      listeners.forEach((listener) => {
        listener(value);
      });
    });
  } else {
    // we call listeners in render for optimization.
    // although this is not a recommended pattern,
    // so far this is only the way to make it as expected.
    // we are looking for better solutions.
    // https://github.com/dai-shi/use-context-selector/pull/12
    listeners.forEach((listener) => {
      listener(value);
    });
  }
  return React.createElement(OrigProvider, { value }, children);
});

createProvider 除了包裹 React 原生的 Context Provide 外，額外接收一個 listeners 參數，而這就是 Custom Provider 的主要目的。

剛剛提到由於 changedBits 都會是零，所以需要我們主動觸發更新，而觸發的方式就是直接將 listener 註冊到 Customer Provder 中，而 listener 就是每個 Component 用來針對目前最新的 context value 做 select 以決定要不要更新的函式，詳細實作等等就會說明。

現在重新拿範例程式碼來檢視一下目前為止的邏輯：

const Provider = ({ children }) => {
  const [state, dispatch] = useReducer(reducer, initialState);
  return (
    <context.Provider value={[state, dispatch]}>
      {children}
    </context.Provider>
  );
};

const App = () => (
  <StrictMode>
    <Provider>
      <h1>Counter</h1>
      <Counter />
      <Counter />
    </Provider>
  </StrictMode>
);

將 useReducer 回傳的 state 與 dispatch 當作 Context Value 傳入 Provider，當 Counter 裡面透過 dispatch 去更新 Context 內的 state 時，由於此時的 Provider 是客製化後的 Provider，他會進行 render，並在 render 的過程中，呼叫所有與他直接 subscribe 的 listener，由 listener 來判斷與執行 component 的 re-render 與否。

這層客製化的 Provider 也就是我們先前在 flamegraph 中看到多出來的一層 Anonymous (memo) component，也解釋了為什麼 commits 數量會多了一倍，就是因為這個 Anonymous component 所進行的 render。

最後我們來看看 listener 是怎麼產生與運作的，我們拆三個部分來說明：

export const useContextSelector = (context, selector) => {
  const listeners = context[CONTEXT_LISTENERS];
  if (process.env.NODE_ENV !== 'production') {
    if (!listeners) {
      throw new Error('useContextSelector requires special context');
    }
  }
  // ...
};

在一開始 createContext 時，其實有在 context 中塞一個 Set()：

context[CONTEXT_LISTENERS] = new Set();

而在 useContextSelector 中的一開始，我們就會取出這個 set，目的在於要放入呼叫 useContextSelector 的 component 的 listener。

// ...
const [, forceUpdate] = React.useReducer((c) => c + 1, 0);
const value = React.useContext(context);
const selected = selector(value);
const ref = React.useRef(null);
React.useLayoutEffect(() => {
  ref.current = {
    f: selector, // last selector "f"unction
    v: value, // last "v"alue
    s: selected, // last "s"elected value
  };
});
// ...

接著準備一些 listener 需要的東西：

1. 當執行完 Selector，確認 Component 需要更新後，我們得有個 forceUpdate 函式來觸發 render，這邊的實作方式是額外使用 React.useReducer 產生一個不斷 +1 的 reducer，來達到效果。
2. 我們還是需要一個真正的 React.context 來紀錄 Globle state。
3. 透過 React.useRef 紀錄當下的 selector function、context value 與 selector 選出的值。

// ...
React.useLayoutEffect(() => {
  const callback = (nextValue) => {
    try {
      if (ref.current.v === nextValue
        || Object.is(ref.current.s, ref.current.f(nextValue))) {
        return;
      }
    } catch (e) {
      // ignored (stale props or some other reason)
    }
    forceUpdate();
  };
  listeners.add(callback);
  return () => {
    listeners.delete(callback);
  };
}, [listeners]);
return selected;

再來實作 listener，listener function 接受的 nextValue 就是 Custom Provider 取得的最新的 Context value，listener function 就能夠利用這個 nextValue 與我們先前存放在 ref 中的值做比較，若是 Context Value 完全相等，或是 Selected 的值也沒有變動（用 ref 中存好的 selector function 對 nextValue 做選取），那就不用 render。

反之，若發現值不同，需要更新，就會呼叫 forceUpdate() 強制讓這個 useContextSelector 進行 render，也就會跟著觸發使用 useContextSelector 的 Component 進行 render，更新 ref 內的值，並回傳最新的 selected value。

而這邊建立的 listener 會放入一開始從 Context 取出的 Set() 中，Custom Provider 在 render 時，就能取出運行。

總結一遍流程

use-context-selector 替 Context API 的效能問題所找到的 escape hatch 流程如下：

1. 利用 Custom Provider 與 Custome createContext 迫使 changedBits 總是回傳 0，停止所有 Context 使用者的自動更新。
2. 建立一個 global listeners 的 Set 在 Context 中，讓 Components 直接 subscribe 到 (Custom Provider)
3. 有使用 useContextSelector 的 components 會建立 listener，放入 Context Set 中進行 subscribe。
4. 當 re-renders 時， 觸發所有 subscribers。
5. listener 執行，檢查 Selector，檢查 Context Value，只針對有需要更新的 Component 做 forceUpdate。

這就是 use-context-selector 所找到的出路，讓你在 global context update 時，bail out of rendering。

結論

use-context-selector 的作者自己也說了這個套件有很多限制與不足

即便他有 v2 版本的實作，是建立在比較有機會實作的 RFC 上，但整體來說還是不能算一個穩定的解決方案。

但是作為使用在內部或是個人專案上來說，是個還不錯的選擇。尤其是簡單易懂的實作，就算是出了什麼問題，只要理解他的原理，也是能找得出問題所在。

這次也是透過閱讀其程式碼，才對 Context API 有更多了解，從中延伸閱讀了很多包含 react-redux v6 當初的效能 issue、RFCs 上的討論、關於 calculateChangedBits 的知識，或甚至是 react scheduling 的一些內部實作。

這也回應到我最開始所說的，有時候太過於遵循 best practice，會讓你失去研究一些有趣問題或是學習的機會，甚至透過走這些旁門走道，會讓你對於 best practice 之所以為 best practice 的原因更加深刻。

分析程式碼的文章有點冗長鬆散，如果你有看到這邊，感謝你的閱讀，若有任何問題也歡迎指教討論！

資料來源

1. use-context-selector
2. When to use Context
3. React-Redux Roadmap: v6, Context, Subscriptions, and Hooks
4. RFC: Context selectors
5. Why calculateChangedBits = () => 0
6. 不一樣的 React context
7. A Secret parts of React New Context API
8. React tips — Context API (performance considerations)