GC友好的memoize

122 阅读1分钟

计算机没有黑魔法,性能优化无非时间和空间二者转换。

常做性能优化的同学,肯定对lodash的memoize很熟悉了。但是memoize有几个天生自带的问题:

  1. 多个复杂类型的入参不好计算缓存key
  2. 缓存强引用,越是大数据,越吃内存

针对上述问题,搞一个GC友好的memoize定制。主要思路如下:

  1. 利用WeakMap来维护一个高效的对象->ID映射关系
  2. 多参数生成缓存key改为ID拼接
  3. 若入参被GC,则释放缓存的计算结果

注意:因为是浅比较,所以!

{} !== {}
Symbol("1") !== Symbol("1")

贴一下代码实现:

import { memoize } from "lodash@4.17.21";

class CacheKey<T extends object> {
  private store: WeakMap<T, string>;
  private idx: number;
  private keyPrefix: string;

  constructor(keyPrefix?: string) {
    this.store = new WeakMap();
    this.idx = 0;
    this.keyPrefix = keyPrefix ?? "key";
  }

  public getKey(obj: T) {
    if (this.store.has(obj)) return this.store.get(obj);
    const nextKey = `${this.keyPrefix}-${++this.idx}`;
    this.store.set(obj, nextKey);
    return nextKey;
  }
}

const keyGenerator = new CacheKey();
const gcRegisterSymbol = Symbol("GCRegister");

const isGcFriendly = typeof FinalizationRegistry !== "undefined";

/**
 * 1. 支持复杂类型的多个参数作为缓存key
 * 2. gc友好,当入参被gc后,缓存的结果也会被gc
 * 3. 复杂数据类型,采用浅比较
 * @param fn 原始函数
 * @returns Memorized版本
 */
export const memorize = (fn: (...args: any[]) => any) => {
  const register = isGcFriendly ? (
    () => new FinalizationRegistry((key) => {
      ret.cache.delete(key);
    })
  )() : null;
  const ret = memoize(fn, (...args: any) => {
    const key = args.map((arg: any) => {
      const type = typeof arg;
      if (type === "number") return `_0_${arg}`;
      if (type === "string") return arg;
      if (type === null) return `_\$null`;
      if (type === "undefined") return `_\$undefined`;
      return keyGenerator.getKey(arg);
    }).join();
    if (register) {
      args.forEach((arg: any) => {
        if (arg !== null && typeof arg === 'object') {
          register.register(arg, key);
        }
      });
    }
    return key;
  });
  return Object.assign(ret, { [gcRegisterSymbol]: register });
}

以下是测试代码: