DataLoader解读及踩坑

951 阅读5分钟

DataLoader 是一个通用的数据加载函数,可用于数据获取的一部分。通过对获取行为的批量处理及对数据缓存达到简化操作及优化性能。

  • 通过事件循环机制,达到对循环内容的延迟及批量处理
  • 通过Map及Promise对象的配合实现对已有数据的缓存

执行流程及解读

创建实例

// 业务代码
const getUserIdLoader = new DataLoader(getUserId)
async function getUserId(uids) => {
  return uids.map(i => ({ uid: i }))
}
  • DataLoader 构造函数源码逻辑(简化代码)
    • _batchLoadFn 就是上面传入的 getUserId 函数,它被保存起里并在特定时机被执行
    • _cacheMap 本质就是 Map 对象,内部以 { key: value } 形式缓存着执行过的数据
function DataLoader(batchLoadFn) {
	// 传入的待执行函数
  this._batchLoadFn = batchLoadFn;
  // 调度函数(批量处理的核心内容)
  this._batchScheduleFn = getValidBatchScheduleFn(options);
  // 生成一个 map{} 缓存对象
  this._cacheMap = getValidCacheMap(options);
  // 内部 newBatch 缓存对象
  this._batch = null;
}
  • getValidBatchScheduleFn() 调度函数生成器
    • 当未传递 options.batchScheduleFn 时,使用内置的调度任务生成器
    • enqueuePostPromiseJob 本身会有一个递减的兼容处理
      • process.nextTick -> setImmediate -> setTimeout
    • 所以其本质就是把传入的 batchLoadFn 执行时机给推迟了。放到本次调用栈结束,下次 event loop 开始前执行。在这里的效果就是会在所有的 promise 任务执行后再执行 fn 的内容
function getValidBatchScheduleFn(options) {
  var batchScheduleFn = options && options.batchScheduleFn;

  if (batchScheduleFn === undefined) {
    return enqueuePostPromiseJob;
  }

  return batchScheduleFn;
}

var enqueuePostPromiseJob = Promise.resolve().then(function () {
  // fn -> dispatchBatch(loader, newBatch) 后续说明
  process.nextTick(fn);
})

批量执行

// 业务代码
// getUserIdLoader.load(i) -> DataLoader.prototype.load(key)
const pormiseList = userIds.map(i => getUserIdLoader.load(i))
const result = await Promise.all(pormiseList)
  • load(key) 加载函数
    • 通过 getCurrentBatch()(后面解读) 获取 batch 对象
    • 生成一个pending状态的promsie对象,并把它的 resolve/reject保存到batch.callbacks队列中。以便在后续的dispatchBatch()去更改状态
    • 并把当前的key塞入batch.keys队列中,用于后续执行dispatchBatch()时的批量值
    • 并把key-promsie存入this._cacheMap中,当后续的x.load(key)执行相同的key就不再生成新的promise而是从缓冲中读取
DataLoader.prototype.load = function load(key) {
  // *batch -> { hasDispatched: false, keys: [], callbacks: [] }
  var batch = getCurrentBatch(this);
  // *this._cacheMap -> new Map()
  var cacheMap = this._cacheMap;
  var cacheKey = key;

	// *存在缓存的情况
  if (cacheMap) {
    // *cachedPromise 是按照单个 key 进行缓存,而不是 [keys] 的
    // *这个 cachedPromise 就是下面的 var promise = new Promise... 的
    var cachedPromise = cacheMap.get(cacheKey);

    if (cachedPromise) {
      // cacheHits 以缓存的 cachedPromise 列表
      var cacheHits = batch.cacheHits || (batch.cacheHits = []);
      return new Promise(function (resolve) {
        cacheHits.push(function () {
          resolve(cachedPromise);
        });
      });
    }
  }

  /**
   * *例如:传入 [1, 2, 1, 1, 2] 的数据
   * *batch.keys -> 只有 [1, 2]
   * *后续的都会命中缓存,并获取之前的值
   */
  batch.keys.push(key);
  var promise = new Promise(function (resolve, reject) {
    // !这里 x.load() -> promise 状态是 pending 的,在后续的 dispatchBatch() 中才真正执行并更改状态
    batch.callbacks.push({
      resolve: resolve,
      reject: reject
    });
  });
  
	/**
   * !对 promise 进行缓存,等价于缓存的 promise.resolve/reject 的状态
   * !也可以说是缓存的 promise.resolve(value) 中的 value 值
   */
  if (cacheMap) {
    cacheMap.set(cacheKey, promise);
  }

  return promise;
}
  • getCurrentBatch() 生成批量信息缓存对象
    • 每次执行 x.load(key) 都执行该方法。但在一个循环内,只有第一个 .load() 会创建新的 newBatch。
    • 同时把dispatchBatch()函数,塞入process.nextTick()中推迟执行
      • dispatchBatch 内会执行,最初配置的 getUserId(keys) 函数
    • 注:上述动作只在本次循环内的第一个load()会执行
/**
 * !每次执行 .load() 都执行该方法
 * *在一个循环内,只有第一个 .load() 会创建新的 newBatch。
 * *之后都是使用这个 existingBatch -> newBatch。
 * *所以在第一次 .load() 时,会往 nextTick 内塞入 dispatchBatch(loader, newBatch)
 */
function getCurrentBatch(loader) {
  var existingBatch = loader._batch;

  // *当 loader._batch 有值,且 hasDispatched = false 是进入
  if (
    existingBatch !== null &&
    !existingBatch.hasDispatched
  ) {
    // *只有每次触发 .load() 的第一个走外面,之后都走这里
    return existingBatch;
  }

  var newBatch = {
    hasDispatched: false,
    keys: [],
    callbacks: []
  };
  // *保存当前的 newBatch,后续就不再生成新的了
  loader._batch = newBatch;

  // !loader._batchScheduleFn -> Promise.resolve().then(function () { process.nextTick(function () { dispatchBatch(loader, newBatch) }) })
  loader._batchScheduleFn(function () {
    dispatchBatch(loader, newBatch);
  });

  return newBatch;
}
  • dispatchBatch() 执行批量处理动作,才真正执行 loadFn
    • const pormiseList = userIds.map(i => getUserInfoLoader.load(i))执行完后,才会执行dispatchBatch()函数
    • 将搜集的batch.keys传入_batchLoadFn(keys)从而达到批量处理的效果
    • 执行resolveCacheHits()去更改缓存中promise的状态
    • 之后拿到values并执行batch.callbacks[i].resolve(value)按照下标位置放入对应的值,完成整个逻辑的执行
    • 注:
      • batch.keys.length===0说明没有新的key进入,那么就不用重新执行_batchLoadFn()可从缓存中获取,实现了性能优化
      • 坑点:当_batchLoadFn()执行后的结果values内容的顺序和batch.callbacks顺序不对应时,那么这个key所缓存的value可能并不是正确的
function dispatchBatch(loader, batch) {
  // 更改标识位状态
  batch.hasDispatched = true;

  if (batch.keys.length === 0) {
    // *没有新key前,都从缓冲的获取
    resolveCacheHits(batch);
    return;
  }

  var batchPromise;

  try {
    // *_batchLoadFn -> new DataLoader(async (keys)=>{})
    batchPromise = loader._batchLoadFn(batch.keys);
  } catch (e) {}

  // *values 为 new DataLoader(async (keys)=>{}) 返回的是 []
  batchPromise.then(function (values) {
    // !先把缓存包裹的 promise 更改状态,在执行真正的 promise。让所有的 promise 状态都变更以完成执行逻辑
    // *遍历 batch.cacheHits 内的缓存的 cachedPromise -> promise
    resolveCacheHits(batch);

    // *遍历 batch.callbacks 执行新的 promise
    for (var i = 0; i < batch.callbacks.length; i++) {
      // !其实真正意义上缓存的就是这个 value
      var value = values[i];

      if (value instanceof Error) {
        batch.callbacks[i].reject(value);
      } else {
        batch.callbacks[i].resolve(value);
      }
    }
  });
}
  • resolveCacheHits() 更改缓存包裹 promise 的状态
// *遍历 batch.cacheHits 及是执行 resolve(cachedPromise)
function resolveCacheHits(batch) {
  if (batch.cacheHits) {
    for (var i = 0; i < batch.cacheHits.length; i++) {
      batch.cacheHits[i]();
    }
  }
}

爬坑事例

前面说到当_batchLoadFn()执行后的结果values内容的顺序和batch.callbacks顺序不对应时。而dispatchBatch()又是按照下标放入val值,那么会造成这个key所缓存的value可能并不是正确的 如:在sql查询中,当我们按照3,2,1的 id 顺序传入并进行查询时。最终其实会按 id 默认升序,按照1,2,3的顺序返回。这样它们的缓存对应关系就变成这样了key:3 -> val:1

  • 解决方式一:可以通过控制 keys 的顺序和 values 的顺序一致解决
  • 解决方式二:配置option.cache=false不进行缓存
const getUserInfoLoader = new DataLoader<number, any>(async (uids: number[]) => {
  // 这里反转 uids 数组,模拟 keys 顺序和 values 顺序不一致的情况
  return uids.reverse()
})

// 入参:keys = [1,1,1,2,3,3] 是 3个1,1个2,2个3
// 期望结果:安顺序也返回的是 [1,1,1,2,3,3]
// 实际结果:[3,3,3,2,1,1] 却是 3个3,1个2,2个1
const result = await Promise.all([1,1,1,2,3,3].map(i => getUserInfoLoader.load(i)))
  • 坑点执行伪代码模拟
// 入参:keys = [1,1,1,2,3,3]
// 执行伪代码输出
key=1
create new promise // 取名叫 p1
cached key=1 and promise // 缓存 promise
key=1 // key 重复使用缓存
use cachedPromise // 是上面的 p1
key=1
use cachedPromise // 是上面的 p1
key=2
create new promise // 取名叫 p2
cached key=2 and promise
key=3
create new promise // 取名叫 p3
cached key=3 and promise
key=3
use cachedPromise // 是上面的 p3

// 进入 dispatchBatch() 阶段
cacheHits = [ p1, p1, p3 ] // 缓存的 promise 队列中
callbacks = [ p1, p2, p3 ] // 待更改状态的 promise
values = [3, 2, 1] // 返回的 values 数据顺序和如参 keys 不一致
// 循环 batch.callbacks 并更改状态
// 它是按照 callbacks 下标更改
batch.callbacks[1].resolve(3) // p1 改为 resolve 并塞入 val;也就改变 cacheHits 内 p1 的状态了
batch.callbacks[2].resolve(2) // p2 改为 resolve
batch.callbacks[3].resolve(1) // p3 改为 resolve

// 最终导致了缓存数据异常
key=1 -> value=3
key=2 -> value=2
key=3 -> value=1