如何在 Zustand 中创建自定义中间件？ - 面试题

在 Zustand 中创建自定义中间件非常灵活，可以用来实现各种功能，如日志记录、状态验证、性能监控等。

基本自定义中间件结构：

javascript
const customMiddleware = (config) => (set, get, api) => {
  // 在原始 store 之前执行的逻辑
  const originalSet = set;
  
  // 包装 set 函数
  const wrappedSet = (partial, replace) => {
    // 在状态更新前执行逻辑
    console.log('State will update:', partial);
    
    // 调用原始 set
    const result = originalSet(partial, replace);
    
    // 在状态更新后执行逻辑
    console.log('State updated:', get());
    
    return result;
  };
  
  // 创建 store
  const store = config(wrappedSet, get, api);
  
  // 返回增强后的 store
  return store;
};

示例 1：日志中间件

javascript
const loggerMiddleware = (config) => (set, get, api) => {
  const originalSet = set;
  
  const wrappedSet = (partial, replace) => {
    const previousState = get();
    const result = originalSet(partial, replace);
    const nextState = get();
    
    console.log('Previous state:', previousState);
    console.log('Action:', partial);
    console.log('Next state:', nextState);
    
    return result;
  };
  
  return config(wrappedSet, get, api);
};

// 使用
const useStore = create(
  loggerMiddleware((set) => ({
    count: 0,
    increment: () => set((state) => ({ count: state.count + 1 }))
  }))
);

示例 2：状态验证中间件

javascript
const validationMiddleware = (schema) => (config) => (set, get, api) => {
  const originalSet = set;
  
  const wrappedSet = (partial, replace) => {
    // 验证状态更新
    const newState = typeof partial === 'function' 
      ? partial(get()) 
      : partial;
    
    const validation = schema.safeParse({ ...get(), ...newState });
    
    if (!validation.success) {
      console.error('State validation failed:', validation.error);
      throw new Error('Invalid state update');
    }
    
    return originalSet(partial, replace);
  };
  
  return config(wrappedSet, get, api);
};

// 使用
import { z } from 'zod';

const storeSchema = z.object({
  count: z.number().min(0),
  user: z.object({
    id: z.string(),
    name: z.string().min(1)
  }).nullable()
});

const useStore = create(
  validationMiddleware(storeSchema)((set) => ({
    count: 0,
    user: null,
    increment: () => set((state) => ({ count: state.count + 1 })),
    decrement: () => set((state) => ({ count: Math.max(0, state.count - 1) }))
  }))
);

示例 3：性能监控中间件

javascript
const performanceMiddleware = (config) => (set, get, api) => {
  const originalSet = set;
  const renderCounts = {};
  
  const wrappedSet = (partial, replace) => {
    const startTime = performance.now();
    const result = originalSet(partial, replace);
    const endTime = performance.now();
    
    const duration = endTime - startTime;
    if (duration > 10) {
      console.warn(`Slow state update: ${duration.toFixed(2)}ms`, partial);
    }
    
    return result;
  };
  
  const store = config(wrappedSet, get, api);
  
  // 跟踪组件渲染次数
  const originalSubscribe = api.subscribe;
  api.subscribe = (listener, selector) => {
    const wrappedListener = (state, previousState) => {
      const key = selector ? selector.toString() : 'full-store';
      renderCounts[key] = (renderCounts[key] || 0) + 1;
      
      if (renderCounts[key] % 10 === 0) {
        console.log(`Render count for ${key}:`, renderCounts[key]);
      }
      
      listener(state, previousState);
    };
    
    return originalSubscribe(wrappedListener, selector);
  };
  
  return store;
};

// 使用
const useStore = create(
  performanceMiddleware((set) => ({
    count: 0,
    increment: () => set((state) => ({ count: state.count + 1 }))
  }))
);

示例 4：撤销/重做中间件

javascript
const undoRedoMiddleware = (config) => (set, get, api) => {
  let history = [];
  let future = [];
  const MAX_HISTORY = 50;
  
  const originalSet = set;
  
  const wrappedSet = (partial, replace) => {
    const previousState = get();
    const result = originalSet(partial, replace);
    const nextState = get();
    
    // 保存到历史记录
    history.push(previousState);
    if (history.length > MAX_HISTORY) {
      history.shift();
    }
    
    // 清空未来记录
    future = [];
    
    return result;
  };
  
  const store = config(wrappedSet, get, api);
  
  // 添加撤销功能
  store.undo = () => {
    if (history.length === 0) return;
    
    const previousState = history.pop();
    future.push(get());
    originalSet(previousState, true);
  };
  
  // 添加重做功能
  store.redo = () => {
    if (future.length === 0) return;
    
    const nextState = future.pop();
    history.push(get());
    originalSet(nextState, true);
  };
  
  // 清空历史
  store.clearHistory = () => {
    history = [];
    future = [];
  };
  
  return store;
};

// 使用
const useStore = create(
  undoRedoMiddleware((set, get) => ({
    count: 0,
    increment: () => set((state) => ({ count: state.count + 1 })),
    decrement: () => set((state) => ({ count: state.count - 1 }))
  }))
);

// 在组件中使用
function Counter() {
  const { count, increment, decrement } = useStore();
  const undo = useStore((state) => state.undo);
  const redo = useStore((state) => state.redo);
  
  return (
    <div>
      <h1>Count: {count}</h1>
      <button onClick={increment}>Increment</button>
      <button onClick={decrement}>Decrement</button>
      <button onClick={undo}>Undo</button>
      <button onClick={redo}>Redo</button>
    </div>
  );
}

关键点：

自定义中间件是一个高阶函数，接收 config 并返回新的配置函数
可以包装 set、get 和 api 来增强功能
中间件的执行顺序很重要，通常外层中间件先执行
可以在中间件中添加额外的功能，如日志、验证、性能监控等
中间件可以返回增强后的 store，添加新的方法或属性