Cache, CacheContext, createCache

まずは cache が呼ばれた際に、そのデータがどこに格納されるのかを知りたい。

少し読んでみると CacheContext という名前で以下のコードが見つかる。

export type Cache = {
  controller: AbortController,
  data: Map<() => mixed, mixed>,
  refCount: number,
};

// ...

export const CacheContext: ReactContext<Cache> = enableCache
  ? {
      $$typeof: REACT_CONTEXT_TYPE,
      // We don't use Consumer/Provider for Cache components. So we'll cheat.
      Consumer: (null: any),
      Provider: (null: any),
      // We'll initialize these at the root.
      _currentValue: (null: any),
      _currentValue2: (null: any),
      _threadCount: 0,
      _defaultValue: (null: any),
      _globalName: (null: any),
    }
  : (null: any);

// ...

// Creates a new empty Cache instance with a ref-count of 0. The caller is responsible
// for retaining the cache once it is in use (retainCache), and releasing the cache
// once it is no longer needed (releaseCache).
export function createCache(): Cache {
  if (!enableCache) {
    return (null: any);
  }
  const cache: Cache = {
    controller: new AbortControllerLocal(),
    data: new Map(),
    refCount: 0,
  };

  return cache;
}

上記のコードのうちCache型のフィールドにある data というMap型の値が、実際にキャッシュを保持していると思われる。

Reactアプリケーションのライフサイクルの中ではルートのFiber生成時に初期化が行われるよう。あとはReactFiberTransitionと呼ばれるモジュールの中でキャッシュプールを生成する際にも使われているようだが、ここは複雑すぎてよく分からなかった。

なお、上記の ReactContext<T> はReact Contextそのものであり、内部的なインメモリキャッシュの実装もReact Contextを用いて行われていることが分かる。

getCacheForType

createCache で生成されたキャッシュはMapをキャッシュデータとして保持するが、実際にほかのモジュールがキャッシュの値を参照するにあたっては getCacheForType という関数を経由する。

// ...

function getCacheForType<T>(resourceType: () => T): T {
  if (!enableCache) {
    throw new Error('Not implemented.');
  }
  const cache: Cache = readContext(CacheContext);
  let cacheForType: T | void = (cache.data.get(resourceType): any);
  if (cacheForType === undefined) {
    cacheForType = resourceType();
    cache.data.set(resourceType, cacheForType);
  }
  return cacheForType;
}

export const DefaultCacheDispatcher: CacheDispatcher = {
  getCacheSignal,
  getCacheForType,
};

この関数は、Reactにおけるキャッシュが必要なあらゆる機能群で使われることが想定されているようである。

例えばReactによるfetch関数の拡張 *1にも使われいるし、ユーザーレベルにも unstable_getCacheForType としてexportされている。

cache

さて、とうとう核心部分であるが、以下が cache の関数の実装になる。

getCacheForType のレシーバとなっている dispatcher という変数は、上記で触れた CacheDispatcher インターフェイスを実装したもの。

export function cache<A: Iterable<mixed>, T>(fn: (...A) => T): (...A) => T {
  return function () {
    const dispatcher = ReactCurrentCache.current;
    if (!dispatcher) {
      // If there is no dispatcher, then we treat this as not being cached.
      // $FlowFixMe[incompatible-call]: We don't want to use rest arguments since we transpile the code.
      return fn.apply(null, arguments);
    }
    const fnMap: WeakMap<any, CacheNode<T>> = dispatcher.getCacheForType(
      createCacheRoot,
    );
    const fnNode = fnMap.get(fn);
    let cacheNode: CacheNode<T>;
    if (fnNode === undefined) {
      cacheNode = createCacheNode();
      fnMap.set(fn, cacheNode);
    } else {
      cacheNode = fnNode;
    }
    for (let i = 0, l = arguments.length; i < l; i++) {
      const arg = arguments[i];
      if (
        typeof arg === 'function' ||
        (typeof arg === 'object' && arg !== null)
      ) {
        // Objects go into a WeakMap
        let objectCache = cacheNode.o;
        if (objectCache === null) {
          cacheNode.o = objectCache = new WeakMap();
        }
        const objectNode = objectCache.get(arg);
        if (objectNode === undefined) {
          cacheNode = createCacheNode();
          objectCache.set(arg, cacheNode);
        } else {
          cacheNode = objectNode;
        }
      } else {
        // Primitives go into a regular Map
        let primitiveCache = cacheNode.p;
        if (primitiveCache === null) {
          cacheNode.p = primitiveCache = new Map();
        }
        const primitiveNode = primitiveCache.get(arg);
        if (primitiveNode === undefined) {
          cacheNode = createCacheNode();
          primitiveCache.set(arg, cacheNode);
        } else {
          cacheNode = primitiveNode;
        }
      }
    }
    if (cacheNode.s === TERMINATED) {
      return cacheNode.v;
    }
    if (cacheNode.s === ERRORED) {
      throw cacheNode.v;
    }
    try {
      // $FlowFixMe[incompatible-call]: We don't want to use rest arguments since we transpile the code.
      const result = fn.apply(null, arguments);
      const terminatedNode: TerminatedCacheNode<T> = (cacheNode: any);
      terminatedNode.s = TERMINATED;
      terminatedNode.v = result;
      return result;
    } catch (error) {
      // We store the first error that's thrown and rethrow it.
      const erroredNode: ErroredCacheNode<T> = (cacheNode: any);
      erroredNode.s = ERRORED;
      erroredNode.v = error;
      throw error;
    }
  };
}

細かい説明は面倒なので省くが、ざっくり引数で与えられた関数の arguments をキーにして、再帰的にネストするツリー構造のような形の CacheNode と呼ばれる型の値を生成している。

すでに関数が実行済み (CacheNode.s === TERMINATED) であれば CacheNode.v に保存されている値を返す。

自分はあまりキャッシュアルゴリズムなどの理論に詳しくないため、なぜこのようなキャッシュの構造になっているのかは実装されたPRのdescriptionを読んでもよく分からなかった...

refreshCache

cache 関数の実装は分かったので、次に unstable_useCacheRefresh のほうの実装を見ていく。

いきなり実装から行くが、以下のコードがキャッシュのinvalidationを行うものになる。

unstable_useCacheRefresh はhookであるため、hookを呼び出したコンポーネントの所属するCacheContextが createCache によって初期化される。

function refreshCache<T>(fiber: Fiber, seedKey: ?() => T, seedValue: T): void {
  if (!enableCache) {
    return;
  }
  // TODO: Does Cache work in legacy mode? Should decide and write a test.
  // TODO: Consider warning if the refresh is at discrete priority, or if we
  // otherwise suspect that it wasn't batched properly.
  let provider = fiber.return;
  while (provider !== null) {
    switch (provider.tag) {
      case CacheComponent:
      case HostRoot: {
        // Schedule an update on the cache boundary to trigger a refresh.
        const lane = requestUpdateLane(provider);
        const refreshUpdate = createLegacyQueueUpdate(lane);
        const root = enqueueLegacyQueueUpdate(provider, refreshUpdate, lane);
        if (root !== null) {
          scheduleUpdateOnFiber(root, provider, lane);
          entangleLegacyQueueTransitions(root, provider, lane);
        }

        // TODO: If a refresh never commits, the new cache created here must be
        // released. A simple case is start refreshing a cache boundary, but then
        // unmount that boundary before the refresh completes.
        const seededCache = createCache();
        if (seedKey !== null && seedKey !== undefined && root !== null) {
          if (enableLegacyCache) {
            // Seed the cache with the value passed by the caller. This could be
            // from a server mutation, or it could be a streaming response.
            seededCache.data.set(seedKey, seedValue);
          } else {
            if (__DEV__) {
              console.error(
                'The seed argument is not enabled outside experimental channels.',
              );
            }
          }
        }

        const payload = {
          cache: seededCache,
        };
        refreshUpdate.payload = payload;
        return;
      }
    }
    provider = provider.return;
  }
  // TODO: Warn if unmounted?
}

ReactCacheのテストコードを読むとこのあたりの挙動はわかりやすい。

おそらく Suspense がひとつのCache boundaryという扱いだと思われる。

  test('refresh a cache boundary', async () => {
    let refresh;
    function App() {
      refresh = useCacheRefresh();
      return <AsyncText showVersion={true} text="A" />;
    }

    // Mount initial data
    const root = ReactNoop.createRoot();
    await act(() => {
      root.render(
        <Suspense fallback={<Text text="Loading..." />}>
          <App />
        </Suspense>,
      );
    });
    assertLog(['Cache miss! [A]', 'Loading...']);
    expect(root).toMatchRenderedOutput('Loading...');

    await act(() => {
      resolveMostRecentTextCache('A');
    });
    assertLog(['A [v1]']);
    expect(root).toMatchRenderedOutput('A [v1]');

    // Refresh for new data.
    await act(() => {
      startTransition(() => refresh());
    });
    assertLog(['Cache miss! [A]', 'Loading...']);
    expect(root).toMatchRenderedOutput('A [v1]');

    await act(() => {
      resolveMostRecentTextCache('A');
    });
    // Note that the version has updated
    if (getCacheSignal) {
      assertLog(['A [v2]', 'Cache cleanup: A [v1]']);
    } else {
      assertLog(['A [v2]']);
    }
    expect(root).toMatchRenderedOutput('A [v2]');

    await act(() => {
      root.render('Bye');
    });
    expect(root).toMatchRenderedOutput('Bye');
  });

こういうのはやはりテストコードを読むといい。

今回の学び

近い将来Reactに fetch をCache APIでラップした実装が入る。Canaryなら今すぐ使える。
- これが入ればfetchに関してはUse APIと組み合わせてすぐにSuspenseなコンポーネントに使うことができるのですごく嬉しい。
CacheのinvalidationはCache boundaryやSuspenseなどの単位で行われる。
- Next.jsのServer-sideで用意されているような、任意のCacheキーを指定して選択的にboundaryを横断してinvalidationするようなことは現状できなさそう。
今回読んだCache APIにまつわる詳しいscrapboxの記事があった → Built-in Suspense Cache APIをどう使うか - tosuke
- TypeScriptで使うにあたっては @types/react で定義されていないものがちらほらあるが、Reactのコードでexportを見るとunstable なprefixでimportはできるっぽい。