ES6でPromiseのレートリミット実装

DynamoDBのようなスループットを気にする必要があるAPIにアクセスする場合、API実行に実行回数制限をかけたくなる場合があります。npmに沢山ライブラリが公開されていますが、Promise関数にレートリミットを付与したクロージャーを返す実装を書いて使うことにしました。

例えばあるPromiseを返す関数の実行を秒間20回に制限したい場合、次のように書くと設定値を超える回数の実行をsleepで調整します。
rateLimit関数では結果的に引数の関数を実行するクロージャーを返しているので、関数の使い方は変わりません。

const testFunctionRateLimit = rateLimit(testFunction, 20, 1000);

実装

ソースコード:

// ES2015以外のPromiseライブラリでも動作可
// import Promise from 'bluebird';

const delay = (ms) => {
  return new Promise((resolve) => {
      setTimeout(resolve, ms);
  });
};

const rateLimit = (func, rate, inMs) => {
  let executedHistories = [];

  const wrapped = (...args) => {
    const now = Date.now();
    executedHistories.push(now);
    executedHistories = executedHistories.filter((executedAt) => {
      return (now - executedAt) < inMs;
    });

    if (rate < executedHistories.length) {
      const sleepTime = inMs - (now - Math.min(...executedHistories));
      return delay(sleepTime).then(() => {
          executedHistories.pop();
          return wrapped(...args);
      });
    } else {
      return func(...args);
    }
  }
  return wrapped;
};

動作確認

ソースコード:

// 動作確認用の関数を定義
const startTime = Date.now();
const testFunction = (arg1, arg2, arg3) => {
  return Promise.resolve().then(() => {
      const elapsed = Date.now() - startTime;
      console.log('arg1:', arg1, ', arg2:', arg2, ', arg3:', arg3, 'elapsed:', elapsed);
  });
};

// testFunctionの実行を1000msに20回に制限
const testFunctionRateLimit = rateLimit(testFunction, 20, 1000);

// 直列に100回実行
let promise = Promise.resolve(null);
for (let i = 0; i < 100; i++) {
  promise = promise.then(() => {
    return testFunctionRateLimit(i, i + 1, i + 2);
  });
}

promise.then(() => {
  console.log('done');
});

出力:

arg1: 0 , arg2: 1 , arg3: 2 elapsed: 2
arg1: 1 , arg2: 2 , arg3: 3 elapsed: 3
...
arg1: 18 , arg2: 19 , arg3: 20 elapsed: 5
arg1: 19 , arg2: 20 , arg3: 21 elapsed: 5
arg1: 20 , arg2: 21 , arg3: 22 elapsed: 1004
arg1: 21 , arg2: 22 , arg3: 23 elapsed: 1005
...
arg1: 38 , arg2: 39 , arg3: 40 elapsed: 1006
arg1: 39 , arg2: 40 , arg3: 41 elapsed: 1006
arg1: 40 , arg2: 41 , arg3: 42 elapsed: 2006
arg1: 41 , arg2: 42 , arg3: 43 elapsed: 2006
...
arg1: 58 , arg2: 59 , arg3: 60 elapsed: 2007
arg1: 59 , arg2: 60 , arg3: 61 elapsed: 2007
arg1: 60 , arg2: 61 , arg3: 62 elapsed: 3003
arg1: 61 , arg2: 62 , arg3: 63 elapsed: 3003
...
arg1: 78 , arg2: 79 , arg3: 80 elapsed: 3006
arg1: 79 , arg2: 80 , arg3: 81 elapsed: 3006
arg1: 80 , arg2: 81 , arg3: 82 elapsed: 4002
arg1: 81 , arg2: 82 , arg3: 83 elapsed: 4002
...
done