zustand 中間件介紹#
本文是對前文zustand 核心源碼解析的補充，以persist插件為例介紹zustand的中間件系統
zustand的核心代碼實現得非常簡潔，功能特性也相對較少。如果想要有更多的特性就需要自行實現或者通過引入中間件的方式使用現有的輪子。
zustand 中間件的本質#
如官方文檔middlware介紹，zustand的中間件實際上是一個高階函數，它的入參和create函數相同，都是createInitState類的函數，但是不同的是它的返回值仍然是一個createInitState類的函數，本質上是對createInitState做了一層包裹，注入特定的邏輯，實現對createInitState的改寫。
import { create } from "zustand";
const createInitState = (set) => ({
bees: false,
setBees: (input) => set((state) => void (state.bees = input)),
});
// no middleware
const useStore = create(createInitState);
// with middleware
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,
// 決定如何merge現有數據和保存在storage裡的數據
merge: (persistedState: unknown, currentState: S) => ({
...currentState,
...(persistedState as object),
}),
}
可以看到最開始的時候調用了createJSONStorage函數來生成 storage, 從我的理解，這個函數的作用是做一個膠水層，使得getItem支持async storage，核心在於
const str = (storage as StateStorage).getItem(name) ?? null;
// support async storage
if (str instanceof Promise) {
return str.then(parse);
}
return parse(str);
從createJSONStorage函數可以看出 persist 中間件最開始應該是基於localStorage進行的封裝，後來進行的拓展，數據的存儲都需要使用 string。可以看到存在這樣一個問題： 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) {
// prevent error if the storage is not defined (e.g. when server side rendering a page)
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 merge 到一起，就類似於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),
}),
// above are default configs
...baseOptions,
};
let hasHydrated = false;
const hydrationListeners = new Set<PersistListener<S>>();
const finishHydrationListeners = new Set<PersistListener<S>>();
let storage = options.storage;
// 沒有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
);
}
// set partialized item
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，每次更新後都種入storage
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
);
// a workaround to solve the issue of not storing rehydrated state in sync storage
// the set(state) value would be later overridden with initial state by create()
// to avoid this, we merge the state from localStorage into the initial state.
let stateFromStorage: S | undefined;
// rehydrate initial state with existing stored 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函數部分，主要流程就是從 storage 中讀取到初始的狀態，通過toThenable函數將非async storage的值也轉換成promisify的類型，統一函數的調用形式。先判斷獲取的值是否需要migrate，然後進行merge，再進行暴露的與hydrate有關的函數的調用。
const hydrate = () => {
if (!storage) return;
// On the first invocation of 'hydrate', state will not yet be defined (this is
// true for both the 'asynchronous' and 'synchronous' case). Pass 'configResult'
// as a backup to 'get()' so listeners and 'onRehydrateStorage' are called with
// the latest available state.
hasHydrated = false;
// 在沒有skipHydration的時候，調用hydrate的時候initState還沒有生成，get()結果是undefined，所以需要使用前置生成的configResult(還沒有和保存的值merge)
hydrationListeners.forEach((cb) => cb(get() ?? configResult));
const postRehydrationCallback =
options.onRehydrateStorage?.(get() ?? configResult) || undefined;
// bind is used to avoid `TypeError: Illegal invocation` error
return toThenable(storage.getItem.bind(storage))(options.name)
.then((deserializedStorageValue) => {
// 此步為了實現根據version進行舊數據的遷移
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) => {
// 這步進行merge
stateFromStorage = options.merge(
migratedState as S,
get() ?? configResult
);
set(stateFromStorage as S, true);
return setItem();
})
.then(() => {
// TODO: In the asynchronous case, it's possible that the state has changed
// since it was set in the prior callback. As such, it would be better to
// pass 'get()' to the 'postRehydrationCallback' to ensure the most up-to-date
// state is used. However, this could be a breaking change, so this isn't being
// done now.
postRehydrationCallback?.(stateFromStorage, undefined);
// It's possible that 'postRehydrationCallback' updated the state. To ensure
// that isn't overwritten when returning 'stateFromStorage' below
// (synchronous-case only), update 'stateFromStorage' to point to the latest
// state. In the asynchronous case, 'stateFromStorage' isn't used after this
// callback, so there's no harm in updating it to match the latest 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);
},
};
}
};