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）で先に投げられた 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 Component から上に遡り、通過したすべての 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;
    // 何かがスローされる可能性があるため、次に作業するものを更新します。
    workInProgress = completedWork;
  } while (completedWork !== null);

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

第二のパス#

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）内に登録されます。具体的な位置は attachSuspenseRetryListeners 内に存在します。

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

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