從原始碼入手探究一個因useImperativeHandle引起的Bug

今天本來正在工位上寫著一段很普通的業務程式碼，將其簡化後大致如下：

function App(props: any) {		// 父元件
  const subRef = useRef<any>(null)
  const [forceUpdate, setForceUpdate] = useState<number>(0)

  const callRef = () => {
    subRef.current.sayName()	// 呼叫子元件的方法
  }

  const refreshApp = () => {	// 模擬父元件重新整理的方法
    setForceUpdate(forceUpdate + 1)
  }

  return <div>
    <SubCmp1 refreshApp={refreshApp} callRef={callRef} />
    <SubCmp2 ref={subRef} />
  </div>
}

class SubCmp1 extends React.Component<any, any> {	// 子元件1
  constructor(props: any) {
    super(props)
    this.state = {
      count: 0
    }
  }

  add = () => {
    this.props.refreshApp()		// 會導致父元件重渲染的操作

    // 修改自身資料，並在回撥函數中呼叫外部方法
    this.setState({ count: this.state.count + 1 }, () => {
      this.props.callRef()
    })
  }

  render() {
    return <div>
      <button onClick={this.add}>Add</button>
      <span>{this.state.count}</span>
    </div>
  }
}

const SubCmp2 = forwardRef((props: any, ref) => {	// 子元件2

  useImperativeHandle(ref, () => {
    return {
      sayName: () => {
        console.log('SubCmp2')
      }
    }
  })

  return <div>SubCmp2</div>
})

程式碼結構其實非常簡單，一個父元件包含有兩個子元件。其中的元件2因為要在父元件中呼叫它的內部方法，所以用forwardRef包裹，並通過useImperativeHandle向外暴露方法。元件1則是通過props傳遞了兩個父元件的方法，一個是用於間接地存取元件2中的方法，另一個則是可能導致父元件重渲染的方法（當然這種結構的安排明顯是不太合理的，但由於專案歷史包袱的原因咱就先不考慮這個問題了\doge）。

然後當我滿心歡喜地Click元件時，一片紅色的Error映入眼簾：

在幾個關鍵位置加上列印：

const callRef = (str) => {
    console.log(str, ' --- ', subRef.current)
}

add = () => {
    this.props.callRef('列印1')

    this.props.refreshApp()
    this.setState({ count: this.state.count + 1 }, () => {
		this.props.callRef('列印2')

        setTimeout(() => {
            this.props.callRef('列印3')
        }, 0)
    })
}

結果：

有點amazing啊。在呼叫前ref.current是有正確值的，在setState的回撥中ref.current變為null了，而在setState的回撥中加上一個非同步後，立即又變為正確值了。

要debug這個問題，一個非常關鍵的位置就在setState的回撥函數。熟悉React內部渲染流程的同學，應該知道，在React觸發更新之後的commit階段，也就是在React更新完DOM之後，針對fiber節點的型別分別做不同的處理（位於commitLifeCycles方法）。例如class元件中，會同步地執行setState的回撥；函陣列件的話，則會同步地執行useLayoutEffect的回撥函數。

帶著這個前提知識的情況下，我們給useImperativeHandle加個斷點。因為對於其他常見的hook和class元件生命週期在React更新渲染中的執行時機都是比較熟悉的，唯獨這個useImperativeHandle內部機制還不太瞭解，然我們看看程式碼在進入該斷點時的執行棧是怎樣的：

首先，在左側的callstack面板裡看到了commitLifeCycles方法，說明 useImperativeHandle這個hook也是在更新渲染後的commit同步執行的。接著我們進去impreativeHandleEffect，也就是useImperativeHandle回撥函數的上一層：

方法體裡先判斷父元件傳入的ref的型別。如果是一個函數，則將執行useImperativeHandle回撥函數執行後的物件傳入去並執行；否則將物件賦值到ref.current上。但這兩種情況都會返回一個清理副作用的函數，而這個清理函數的任務就是——把我的ref.current給置為null ！？

抓到這個最重要的線索了，趕緊給這個清理函數打個斷點，然後再觸發一次更新看下：

這個清理函數是在commitMutationEffects時期執行的；commitMutationEffects裡做的主要工作就是就是fiber節點的型別執行需要操作的副作用（位於commitWork方法），例如對DOM的增刪改，以及我們熟知的useLayoutEffect的清理函數也是在這時候完成的。

到目前為止，引發報錯問題的整條鏈路就清晰了：

在觸發更新後，在commit階段的commitMutationEffects部分會先執行useImperativeHandle的清理函數，自這之後ref.current就被置為了null。

⬇

接著才到commitLayoutEffects，該部分會執行setState，useLayoutEffect和useImpreativeHandle這些方法的回撥。

⬇

依據React以深度優先遍歷方式生成fiber樹且邊生成邊收集副作用的規則，子元件1中setState回撥會比useImpreativeHandle的回撥先執行，那麼此時ref.current仍然還為null。