React Suspense のソースコード解析

説明#

表紙の画像は我が家の猫の妹です。
React バージョン v18.1.0
React.lazy で導入されたコンポーネントを lazyComponent と呼び、Suspense の fallback 内のコンポーネントを fallbackComponent と呼びます。

以下のコードに基づいて分析します。

/**
 * src/components/suspense/OtherComponent.js
*/
import React from 'react';

function OtherComponent() {
  return (
    <div>
      OtherComponent
    </div>
  );
}
  
export default OtherComponent;

/**
 * src/components/suspense/index.js
*/
import React, { Suspense } from 'react';

const OtherComponent = React.lazy(() => import('./OtherComponent'));

function SuspenseTest() {
  return (
    <div>
      <Suspense fallback={<div>Loading...</div>}>
        <OtherComponent />
      </Suspense>
    </div>
  );
}

export default SuspenseTest;

TLNR#

beginWork で fiber を作成する過程で初めて lazyComponent に遭遇した際、pending 状態の promise が throw されます。この promise オブジェクトは外部の try...catch...（sync モードでは renderRootSync 内、concurrent モードでは renderRootConcurrent 内）で捕捉され、その後上に遡って通過したすべての fiber ノードに Incomplete タグが付けられ、Suspense ノードに達するまで（completeUnitOfWork）行われます。その後、そのノードから再び下に向かって beginWork で fiber を作成し、この時に Suspense の fallback ノードが作成されます。その後、mutation フェーズ（commitMutationEffects）で以前に throw された promise に成功コールバックが追加され、非同期コンポーネントが成功裏に読み込まれた後に再レンダリングをサポートします。

React.lazy#

React.lazy の Suspense に関する分析はそれほど重要なポイントではないため、ts 宣言を参考にして引数と戻り値の型を理解するだけで十分です。

function lazy<T extends ComponentType<any>>(
    factory: () => Promise<{ default: T }>
): LazyExoticComponent<T>;

サンプルコードの OtherComponent に打点を入れると、その値をより直感的に見ることができ、その中で $$typeof にのみ注目すればよいです。

最初のパス#

タイトルは updateSuspenseComponent 内のコメントから取られています。この関数は beginWork が Suspense ノードに達したときに呼び出され、その内部ロジックは First pass と Second pass でそれぞれ異なるロジックを持っています（The second pass では LazyComponent の fallback に対応するノードが作成されます）。しかし、The first pass ではまず offScreen を作成します。文中の例は First pass フェーズを経た後、以下のような構造の fiber ツリーを得ることができます。

この中で lazy という fiber は React.lazy で導入された OtherComponent に対応しており、この時点ではまだプレースホルダーに過ぎず、コンポーネントとは呼べません。beginWork が lazy に達したときに初めて導入コンポーネントのロジックが呼び出されます（コードは mountLazyComponent にあります）。

function mountLazyComponent(
  _current,
  workInProgress,
  elementType,
  renderLanes,
) {
  resetSuspendedCurrentOnMountInLegacyMode(_current, workInProgress);

  const props = workInProgress.pendingProps;
  const lazyComponent: LazyComponentType<any, any> = elementType;
  const payload = lazyComponent._payload;
  const init = lazyComponent._init;
  let Component = init(payload);
  ...
  // 以下は気にしなくて良い、init 関数内で throw によって中断されます。以下のコードを参照してください。
}

// mountLazyComponent 内で呼び出される init
function lazyInitializer<T>(payload: Payload<T>): T {
  if (payload._status === Uninitialized) {
    // この時の payload._result は import OtherComponent 部分で、その呼び出しは pending の promise を返します。
    const ctor = payload._result;
    const thenable = ctor();
    // ここで promise インスタンスにコールバックが追加されますが、読み込み成功後の再レンダリングのロジックはここにはありません。
    thenable.then(
      moduleObject => {
        if (payload._status === Pending || payload._status === Uninitialized) {
          // 次の状態に遷移します。
          const resolved: ResolvedPayload<T> = (payload: any);
          resolved._status = Resolved;
          resolved._result = moduleObject;
        }
      },
      error => {
        if (payload._status === Pending || payload._status === Uninitialized) {
          // 次の状態に遷移します。
          const rejected: RejectedPayload = (payload: any);
          rejected._status = Rejected;
          rejected._result = error;
        }
      },
    );
    if (payload._status === Uninitialized) {
      // first pass の時にこのロジックに入ります。状態を Uninitialized から Pending に変更します。
      const pending: PendingPayload = (payload: any);
      pending._status = Pending;
      pending._result = thenable;
    }
  }
  if (payload._status === Resolved) {
    const moduleObject = payload._result;
    ...dev 内容
    return moduleObject.default;
  } else {
    // first pass では pending 状態の promise を投げ出します。
    // そして外部の try catch に捕捉されます。
    throw payload._result;
  }
}

上記のコードブロックで ctor に対応する関数

ctor

対応するコールスタック

エラーキャッチ —— throwException#

エラーキャッチは renderRoot の handleError 内で発生し、この関数内の throwException はまず lazy fiber に Incomplete のマークを付け、最近の Suspense ノードを見つけて shouldCapture のマークを付けます（このマークは The second pass 内で作用します）。 lazy fiber から上に遡って通過したすべての fiber ノードに Incomplete マークが付けられ、Suspense ノードに達するまで続きます（Suspense ノードは Incomplete 以外に shouldCapture マークも付けられます）。

handleError 全体を分析するのは長くなるため、以下では重要なステップのみを抜粋して証明します。

lazyComponent に Incomplete マークを付ける#

最近の SuspenseComponent を見つけて shouldCapture マークを付ける#

throwException 内のいくつかの小さな詳細#

ConcurrentMode 下の attachPingListener#

attachPingListener は ConcurrentMode 下でのみ呼び出されます。非同期コンポーネントのリクエストは、React が commit フェーズに入るときに返される可能性があるため、ConcurrentMode モードでは高優先度のタスクが割り込むことが許可されています。リクエストに対応する優先度が進行中のタスクよりも高い場合、割り込んで処理されます。

function attachPingListener(root: FiberRoot, wakeable: Wakeable, lanes: Lanes) {
  // 以下は pingCache を維持するロジックであり、読むべき重点ではありません。
  let pingCache = root.pingCache;
  let threadIDs;
  if (pingCache === null) {
    pingCache = root.pingCache = new PossiblyWeakMap();
    threadIDs = new Set();
    pingCache.set(wakeable, threadIDs);
  } else {
    threadIDs = pingCache.get(wakeable);
    if (threadIDs === undefined) {
      threadIDs = new Set();
      pingCache.set(wakeable, threadIDs);
    }
  }
  if (!threadIDs.has(lanes)) {
    // 冗長なリスナーを防ぐためにスレッド ID を使用してメモ化します。
    threadIDs.add(lanes);
    // ping 関数内で ensureRootIsScheduled を呼び出してタスクスケジューリングをトリガーします。
    // ensureRootIsScheduled の解析は私の [useState](https://github.com/IWSR/react-code-debug/issues/3) の解析を参照してください。
    const ping = pingSuspendedRoot.bind(null, root, wakeable, lanes);
    if (enableUpdaterTracking) {
      if (isDevToolsPresent) {
        // まだ保留中の作業がある場合、元のアップデーターを復元します。
        restorePendingUpdaters(root, lanes);
      }
    }
    // ここでの wakeable は以前の pending の promise であり、つまり ConcurrentMode モード下で非同期コンポーネントの成功リクエストも ensureRootIsScheduled をトリガーする可能性があります。
    wakeable.then(ping, ping);
  }
}

attachRetryListener#

attachRetryListener は主に pending 状態の promise を Suspense fiber の updateQueue に掛けるもので、他の fiber とは少し異なります（結局のところ、すべて update インスタンスを装填するための連結リストです）。

function attachRetryListener(
  suspenseBoundary: Fiber,
  root: FiberRoot,
  wakeable: Wakeable,
  lanes: Lanes,
) {
  // リトライリスナー
  //
  // fallback がコミットされる場合、異なるタイプのリスナーをアタッチする必要があります。このリスナーは、Suspense 境界の更新をスケジュールして fallback 状態をオフにします。
  //
  // コミットフェーズでアクセスできるように、境界 fiber に wakeable を保存します。
  //
  // wakeable が解決されると、境界を再度レンダリングしようとします（「リトライ」）。
  const wakeables: Set<Wakeable> | null = (suspenseBoundary.updateQueue: any);
  if (wakeables === null) {
    const updateQueue = (new Set(): any);
    updateQueue.add(wakeable);
    suspenseBoundary.updateQueue = updateQueue;
  } else {
    wakeables.add(wakeable);
  }
}

エラーキャッチ —— completeUnitOfWork#

throwException の後に completeUnitOfWork が呼び出されます。この関数の役割は、lazy コンポーネントから上に遡って通過したすべての fiber ノードに Incomplete マークを付けることですが、ShouldCapture マークの付いたノードに遭遇すると、そのノードに関するエラー処理のマークを消去します（next.flags &= HostEffectMask）し、そのノードから beginWork に入ります。これは completeWork のエントリ関数でもあり、React の重要な関数です。

function completeUnitOfWork(unitOfWork: Fiber): void {
  let completedWork = unitOfWork;
  do {
    const current = completedWork.alternate; // WIP に対応する current ノード
    const returnFiber = completedWork.return; // 親ノード

    if ((completedWork.flags & Incomplete) === NoFlags) {
      // 現在のノードにエラーがない場合（Incomplete マークがない場合）、このロジックに入ります。
			...
      let next;
      ... 不要な部分は削除しました。
      // completeWork の役割は主に DOM を生成し、props を更新し、イベントをバインドすることです。詳細は他の記事を参照してください。
      next = completeWork(current, completedWork, subtreeRenderLanes);

      if (next !== null) {
        // この fiber の完了により新しい作業が生成されました。その次の作業を行います。
        workInProgress = next;
        return;
      }
    } else {
      // ここに到達した場合、ノードには Incomplete マークが付いています。
      // unwindWork が ShouldCapture マークの付いたノードに遭遇すると、そのノードを DidCapture にマークして返します。
      // unwindWork 内の (flags & ~ShouldCapture) | DidCapture という部分が DidCapture のマークを付けます。
      // ShouldCapture マークの付いたノードは通常エラーバウンダリと呼ばれ、unwindWork の役割はこの関数です。
      const next = unwindWork(current, completedWork, subtreeRenderLanes);

      // この fiber は完了しなかったため、そのレーンをリセットしません。

      if (next !== null) {
        /**
         * この例では SuspenseComponent が ShouldCapture を持っているため、SuspenseComponent に遭遇するとそのノードに関するエラー処理のマークが消去されます。
         * 
         * マーク消去のロジックは以下を参照してください。
         *
         * 0b0000000000001000000000000000 Incomplete      
         * 0b0000000000010000000000000000 ShouldCapture  
         * 0b0000000000011000000000000000 Incomplete と ShouldCapture の両方のマークを持つノード
         * 0b0000000000000111111111111111 HostEffectMask
         * 0b0000000000000000000000000000 上記の二つのマークをビットごとに AND（=&）すると、明らかに Incomplete と ShouldCapture マークが消えます。
        */
        next.flags &= HostEffectMask;
        // ここで suspense ノードを workInProgress に割り当てます。
        workInProgress = next;
        // return すると再び workLoopSync/workLoopConcurrent に入り、その後 workInProgress から beginWork を開始します。
        // つまり、ここから再び suspense ノードに対して beginWork を行うことになり、これが The second pass の始まりを意味します。
        return;
      }

      ... 不要なロジックは省略します。

      if (returnFiber !== null) {
        // 現在の fiber の親ノードに Incomplete マークを付け、サブツリーのフラグをクリアします。
        // 現在の fiber の親ノードに Incomplete マークを付けます。
        returnFiber.flags |= Incomplete;
        returnFiber.subtreeFlags = NoFlags;
        returnFiber.deletions = null;
      } else {
        // ルートノードまで戻りました。
        workInProgressRootExitStatus = RootDidNotComplete;
        workInProgress = null;
        return;
      }
    }

    const siblingFiber = completedWork.sibling;
    if (siblingFiber !== null) {
      // この returnFiber にまだ作業がある場合、次にそれを行います。
      // 兄弟ノードが存在する場合、兄弟ノードに対して beginWork を行います。workLoopSync/workLoopConcurrent 内の while の条件を参照してください。
      workInProgress = siblingFiber;
      return;
    }
    // そうでなければ、親に戻ります。
		// workInProgress を現在のノードの親ノードに指し示し、do...while... ループを続けます。
    completedWork = returnFiber;
    // 何かが throw される場合に備えて、次に作業するものを更新します。
    workInProgress = completedWork;
  } while (completedWork !== null);

  // すでにルートノードに到達したため、workInProgressRootExitStatus を完了状態にマークします。
  if (workInProgressRootExitStatus === RootInProgress) {
    workInProgressRootExitStatus = RootCompleted;
  }
}

The second pass#

completeUnitOfWork の処理を経た後、Suspense ノードには DidCapture マークが付けられ、The second pass では beginWork が再び Suspense ノードから fiber ノードを作成します（updateSuspenseComponent）。DidCapture の影響により、updateSuspenseComponent は first pass でスキップされたロジックに入り、fallback の内容を処理します。

function updateSuspenseComponent(current, workInProgress, renderLanes) {
  const nextProps = workInProgress.pendingProps;

  let suspenseContext: SuspenseContext = suspenseStackCursor.current;

  let showFallback = false;
  // 現在の fiber に DidCapture マークが付いているかどうかを判断します。
  const didSuspend = (workInProgress.flags & DidCapture) !== NoFlags;

  if (
    didSuspend ||
    shouldRemainOnFallback(
      suspenseContext,
      current,
      workInProgress,
      renderLanes,
    )
  ) {
		// DidCapture マークが存在するため、showFallback が true になります。
    showFallback = true;
    // DidCapture をノードのマークからクリアします。
    workInProgress.flags &= ~DidCapture;
  } else {
    ... The second pass では入らないので削除しました。
  }

	...
  
  if (current === null) {
    ... hydrate に関連するコードを削除します。
    // これは脱水された suspense コンポーネントであった可能性があります。
    const suspenseState: null | SuspenseState = workInProgress.memoizedState;
    ...

    const nextPrimaryChildren = nextProps.children;
    const nextFallbackChildren = nextProps.fallback;

    if (showFallback) {
      // second pass の時 showFallback は true です。
      // ここで返されるのは fallback の fiber で、以下のスクリーンショットには内部の大まかなロジックが紹介されています。
      const fallbackFragment = mountSuspenseFallbackChildren(
        workInProgress,
        nextPrimaryChildren,
        nextFallbackChildren,
        renderLanes,
      );
      const primaryChildFragment: Fiber = (workInProgress.child: any);
      primaryChildFragment.memoizedState = mountSuspenseOffscreenState(
        renderLanes,
      );
      workInProgress.memoizedState = SUSPENDED_MARKER;
      ...
			// 次に fallback のノードを作成します。
      return fallbackFragment;
    } else if (
      enableCPUSuspense &&
      typeof nextProps.unstable_expectedLoadTime === 'number'
    ) {
      ...
    } else {
     	...
    }
  } else {
    ... 
}

mountSuspenseFallbackChildren 内の大まかなロジック

上記のロジックを経た後、その構造は以下のようになります。

The second pass が終了した後、その fiber ツリー構造は以下のようになります。

promise resolve 後にトリガーされる再レンダリングについて#

非同期コンポーネントの読み込みが終了した後、つまり以前の pending の promise インスタンスの状態が変わった後、更新をトリガーするために対応する成功コールバックを追加する必要があります。このコールバック関数は commit フェーズ（commitRootImpl）内の mutation 小フェーズ（commitMutationEffects）内で promise に登録され、具体的な位置は attachSuspenseRetryListeners 内に存在します。

図中の retry 関数は resolveRetryWakeable（位置は ReactFiberWorkLoop.old.js 内）です。

ensureRootIsScheduled に関連するものは私の React Hooks: useState 分析の解析を参照してください。