zustand-middleware

zustand ミドルウェアの紹介#

この記事は前回のzustand コアソースコード解析の補足であり、persistプラグインを例にしてzustandのミドルウェアシステムを紹介します。

zustandのコアコードは非常にシンプルに実装されており、機能特性も比較的少ないです。より多くの機能を持たせたい場合は、自分で実装するか、ミドルウェアを導入して既存の機能を利用する必要があります。

zustand ミドルウェアの本質#

公式ドキュメントmiddlewareによると、zustandのミドルウェアは実際には高階関数であり、その引数はcreate関数と同じで、どちらもcreateInitStateクラスの関数ですが、異なるのはその戻り値もcreateInitStateクラスの関数であり、本質的にはcreateInitStateをラップし、特定のロジックを注入してcreateInitStateを改変するものです。

import { create } from "zustand";

const createInitState = (set) => ({
  bees: false,
  setBees: (input) => set((state) => void (state.bees = input)),
});
// ミドルウェアなし
const useStore = create(createInitState);

// ミドルウェアあり
const useStoreWithMiddleware = create(
  middleware1(middleware2(createInitState))
);

persist ミドルウェアソースコード詳細#

persist ミドルウェアの詳細な紹介はpersistを参照してください。

引数はcreateInitStateとoptionsであり、その中でoptions.nameは必須で、他はオプションです。optionsのデフォルト値は以下の通りです。

defaultOptions = {
    storage: createJSONStorage<S>(() => localStorage),
    // state内のどのデータを保存するかを決定
    partialize: (state: S) => state,
    version: 0,
    // 既存のデータとstorage内のデータをどのようにマージするかを決定
    merge: (persistedState: unknown, currentState: S) => ({
      ...currentState,
      ...(persistedState as object),
    }),
}

最初にcreateJSONStorage関数を呼び出してストレージを生成していることがわかります。私の理解では、この関数の役割はgetItemがasync storageをサポートするための接着層を作成することです。核心は以下の通りです。

const str = (storage as StateStorage).getItem(name) ?? null;
// async storageをサポート
if (str instanceof Promise) {
  return str.then(parse);
}
return parse(str);

createJSONStorage関数から、persist ミドルウェアは最初にlocalStorageを基にしたラッピングであることがわかります。その後の拡張では、データの保存には文字列を使用する必要があります。以下のような問題が存在します：setItemはasync storageをサポートしておらず、await操作を行っていません。呼び出し後、完了しているかどうかに関わらず、これはgetItemと一致しません。しかし、別の視点から見ると、実際にはsetItemの完了を保証する必要はありません。一般的にはsetItemの後にすぐにgetItemを実行することはないため、awaitすると一部のパフォーマンスが失われます。ただし、私の観点からはawaitした方が良いと思います。

createJSONStorageのソースコードは以下の通りです。

export interface StateStorage {
  getItem: (name: string) => string | null | Promise<string | null>;
  setItem: (name: string, value: string) => void | Promise<void>;
  removeItem: (name: string) => void | Promise<void>;
}

export function createJSONStorage<S>(
  getStorage: () => StateStorage,
  options?: JsonStorageOptions
): PersistStorage<S> | undefined {
  let storage: StateStorage | undefined;
  try {
    storage = getStorage();
  } catch (e) {
    // ストレージが定義されていない場合のエラーを防ぐ（例：サーバーサイドレンダリング時）
    return;
  }
  const persistStorage: PersistStorage<S> = {
    getItem: (name) => {
      const parse = (str: string | null) => {
        if (str === null) {
          return null;
        }
        return JSON.parse(str, options?.reviver) as StorageValue<S>;
      };
      const str = (storage as StateStorage).getItem(name) ?? null;
      // async storageをサポート
      if (str instanceof Promise) {
        return str.then(parse);
      }
      return parse(str);
    },
    setItem: (name, newValue) =>
      (storage as StateStorage).setItem(
        name,
        JSON.stringify(newValue, options?.replacer)
      ),
    removeItem: (name) => (storage as StateStorage).removeItem(name),
  };
  return persistStorage;
}

以下はコアソースコード（hydrate関数の内容は省略されています。次のステップで紹介します。）で、zustand の既存の API にapi.persistを追加して persist の API を公開するために使用されており、最終的な戻り値もcreateInitState（つまり config）の戻り値である initState です。

options.skipHydrationとhydrate関数が見られます。この関数の役割は、保存された state と既存の state をマージすることです。これはSSRの水と同様で、イベントを適切な DOM にマウントします。

const newImpl = (config, baseOptions) => (set, get, api) => {
  // ここでのSは実際には前述のcreateInitStateの戻り値の型です
  type S = ReturnType<typeof config>;
  let options = {
    storage: createJSONStorage<S>(() => localStorage),
    partialize: (state: S) => state,
    version: 0,
    merge: (persistedState: unknown, currentState: S) => ({
      ...currentState,
      ...(persistedState as object),
    }),
    // 上記はデフォルト設定
    ...baseOptions,
  };

  let hasHydrated = false;
  const hydrationListeners = new Set<PersistListener<S>>();
  const finishHydrationListeners = new Set<PersistListener<S>>();
  let storage = options.storage;

  // ストレージがない場合は保存せず、直接config(..args)を返す（つまりinitState）
  if (!storage) {
    return config(
      (...args) => {
        console.warn(
          `[zustand persist middleware] Unable to update item '${options.name}', the given storage is currently unavailable.`
        );
        set(...args);
      },
      get,
      api
    );
  }

  // 部分的なアイテムを設定
  const setItem = (): void | Promise<void> => {
    // partialize処理されたstateのみを保存
    const state = options.partialize({ ...get() });
    return (storage as PersistStorage<S>).setItem(options.name, {
      state,
      version: options.version,
    });
  };

  const savedSetState = api.setState;

  // 新しいsetStateを置き換え、毎回更新後にストレージに保存
  api.setState = (state, replace) => {
    savedSetState(state, replace);
    void setItem();
  };

  const configResult = config(
    /**
     *
    このステップはcreateState関数に注入されたsetStateとapi.setStateを等価にするためのものです
    こうして注入されるため、更新されたapi.setStateは注入されたset関数に影響を与えません
    const api = { setState, getState, subscribe, destroy }
    state = createState(setState, getState, api)
     */
    (...args) => {
      // ここでのset === savedSetState
      set(...args);
      void setItem();
    },
    get,
    api
  );

  // 非同期ストレージに再水和されたstateを保存しない問題を解決するためのワークアラウンド
  // set(state)の値は後でcreate()によって初期状態で上書きされる
  // これを避けるために、localStorageからのstateを初期状態にマージします。
  let stateFromStorage: S | undefined;

  // 既存の保存されたstateで初期状態を再水和
  const hydrate = () => {
    ...
  };

  (api as StoreApi<S> & StorePersist<S, S>).persist = {
    setOptions: (newOptions) => {
      options = {
        ...options,
        ...newOptions,
      };

      if (newOptions.storage) {
        storage = newOptions.storage;
      }
    },
    clearStorage: () => {
      storage?.removeItem(options.name);
    },
    getOptions: () => options,
    rehydrate: () => hydrate() as Promise<void>,
    hasHydrated: () => hasHydrated,
    onHydrate: (cb) => {
      hydrationListeners.add(cb);

      return () => {
        hydrationListeners.delete(cb);
      };
    },
    onFinishHydration: (cb) => {
      finishHydrationListeners.add(cb);

      return () => {
        finishHydrationListeners.delete(cb);
      };
    },
  };

  if (!options.skipHydration) {
    hydrate();
  }

  return stateFromStorage || configResult;
};

ここでhydrate関数の部分を紹介します。主な流れは、ストレージから初期状態を読み取り、toThenable関数を通じて非async storageの値もpromisifyの型に変換し、関数の呼び出し形式を統一します。まず、取得した値がmigrateする必要があるかどうかを判断し、次にmergeを行い、最後に公開されたhydrateに関連する関数を呼び出します。

const hydrate = () => {
  if (!storage) return;

  // 'hydrate'の最初の呼び出しでは、stateはまだ定義されていません（これは
  // '非同期'および'同期'の両方のケースに当てはまります）。バックアップとして
  // 'configResult'を'get()'に渡し、リスナーと'onRehydrateStorage'が
  // 最新の利用可能なstateで呼び出されるようにします。

  hasHydrated = false;
  // skipHydrationがない場合、hydrateを呼び出すとinitStateがまだ生成されていないため、get()の結果はundefinedです。したがって、前もって生成されたconfigResult（まだ保存された値とマージされていない）を使用する必要があります。
  hydrationListeners.forEach((cb) => cb(get() ?? configResult));

  const postRehydrationCallback =
    options.onRehydrateStorage?.(get() ?? configResult) || undefined;

  // bindは`TypeError: Illegal invocation`エラーを回避するために使用されます
  return toThenable(storage.getItem.bind(storage))(options.name)
    .then((deserializedStorageValue) => {
      // このステップは、バージョンに基づいて古いデータを移行するためのものです
      if (deserializedStorageValue) {
        if (
          typeof deserializedStorageValue.version === "number" &&
          deserializedStorageValue.version !== options.version
        ) {
          if (options.migrate) {
            return options.migrate(
              deserializedStorageValue.state,
              deserializedStorageValue.version
            );
          }
          console.error(
            `State loaded from storage couldn't be migrated since no migrate function was provided`
          );
        } else {
          return deserializedStorageValue.state;
        }
      }
    })
    .then((migratedState) => {
      // このステップでマージを行います
      stateFromStorage = options.merge(
        migratedState as S,
        get() ?? configResult
      );

      set(stateFromStorage as S, true);
      return setItem();
    })
    .then(() => {
      // TODO: 非同期の場合、以前のコールバックで設定された後にstateが変更される可能性があります。
      // そのため、最も最新のstateが使用されるように、'postRehydrationCallback'に'get()'を渡す方が良いでしょう。
      // ただし、これは破壊的な変更になる可能性があるため、今は行われていません。
      postRehydrationCallback?.(stateFromStorage, undefined);

      // 'postRehydrationCallback'がstateを更新した可能性があります。
      // これにより、以下で'stateFromStorage'を返すときに上書きされないようにするために、
      // 'stateFromStorage'を最新のstateを指すように更新します。
      // 非同期の場合、'stateFromStorage'はこのコールバックの後に使用されないため、
      // 最新のstateに一致させることに問題はありません。
      stateFromStorage = get();
      hasHydrated = true;
      finishHydrationListeners.forEach((cb) => cb(stateFromStorage as S));
    })
    .catch((e: Error) => {
      postRehydrationCallback?.(undefined, e);
    });
};

toThenable関数のソースコード

const toThenable =
  <Result, Input>(
    fn: (input: Input) => Result | Promise<Result> | Thenable<Result>
  ) =>
  (input: Input): Thenable<Result> => {
    try {
      const result = fn(input);
      if (result instanceof Promise) {
        return result as Thenable<Result>;
      }
      return {
        then(onFulfilled) {
          return toThenable(onFulfilled)(result as Result);
        },
        catch(_onRejected) {
          return this as Thenable<any>;
        },
      };
    } catch (e: any) {
      return {
        then(_onFulfilled) {
          return this as Thenable<any>;
        },
        catch(onRejected) {
          return toThenable(onRejected)(e);
        },
      };
    }
  };