Refactor Lock to be counting (for rate limiting)

src/utils/CountingLock.js

@@ -0,0 +1,58 @@
+'use strict';
+require('regenerator-runtime/runtime');
+require('source-map-support').install();
+
+const _ = require('lodash');
+
+class CountingLock {
+  constructor(width=1) {
+    this.executing = 0;
+    this.waiters = [];
+    this.width = width;
+  }
+
+  lockForPath(f) {
+    const waiter = {f};
+    const promise = new Promise((resolve, reject) => {
+      waiter.resolve = resolve;
+      waiter.reject = reject;
+    });
+    this.waiters.push(waiter);
+    this._execute();
+    return promise;
+  }
+
+  async _execute() {
+    if (this.executing >= this.width) {
+      return;
+    }
+    if (this.waiters.length > 0) {
+      this.executing = this.executing + 1;
+      const {f, resolve, reject} = this.waiters.shift();
+      let value;
+      let error;
+      try {
+        value = await f();
+      } catch (e) {
+        error = e;
+      }
+
+      try {
+        if (error) {
+          reject(error);
+        } else {
+          resolve(value);
+        }
+      } finally {
+        this.executing = this.executing - 1;
+        if (this.waiters.length > 0) {
+          _.defer(() => {
+            this._execute();
+          });
+        }
+      }
+    }
+  }
+}
+
+module.exports = CountingLock;

src/utils/Lock.js

@@ -1,52 +1,11 @@
 'use strict';
 require('regenerator-runtime/runtime');
 require('source-map-support').install();
+const CountingLock = require('./CountingLock');
 
-const _ = require('lodash');
-
-class Lock {
+class Lock extends CountingLock {
   constructor() {
-    this.executing = false;
-    this.waiters = [];
-  }
-
-  lockForPath(f) {
-    const waiter = {f};
-    const promise = new Promise((resolve, reject) => {
-      waiter.resolve = resolve;
-      waiter.reject = reject;
-    });
-    this.waiters.push(waiter);
-    this._execute();
-    return promise;
-  }
-
-  async _execute() {
-    if (this.executing) {
-      return;
-    }
-    this.executing = true;
-    if (this.waiters.length > 0) {
-      const {f, resolve, reject} = this.waiters.shift();
-      let value;
-      let error;
-      try {
-        value = await f();
-      } catch (e) {
-        error = e;
-      }
-      if (error) {
-        reject(error);
-      } else {
-        resolve(value);
-      }
-    }
-    this.executing = false;
-    if (this.waiters.length > 0) {
-      _.defer(() => {
-        this._execute();
-      });
-    }
+    super(1);
   }
 }
 

test/utils/testCountingLock.js

@@ -0,0 +1,111 @@
+const sinon = require('sinon');
+const should = require('should');
+require('should-sinon');
+const Promise = require('bluebird');
+const _ = require('lodash');
+const CountingLock = require('../../src/utils/CountingLock');
+
+const sleep = ms => new Promise(resolve => setTimeout(resolve, ms));
+
+const width = 5;
+
+describe('CountingLock', () => {
+  let lock;
+
+  beforeEach(() => {
+    lock = new CountingLock(width);
+  });
+
+  afterEach(() => {
+    lock = null;
+  });
+
+  it('Locked code doesn\'t execute concurrently beyond width', async () => {
+    const nConcurrentPromises = 15;
+    const promiseRunning = _.map(_.range(nConcurrentPromises), () => false);
+    const promiseAtIndex = async index => {
+      promiseRunning[index] = true;
+
+      let nRunningPromises = 0;
+      for (const p of promiseRunning) {
+        if (p) {
+          nRunningPromises++;
+        }
+      }
+      nRunningPromises.should.be.lessThanOrEqual(width);
+
+      await sleep(0.2);
+      promiseRunning[index] = false;
+    };
+    await Promise.all(_.map(_.range(nConcurrentPromises), i => lock.lockForPath(() => (promiseAtIndex(i)))));
+  });
+
+  it('More than one thing can run at a time', async () => {
+    const nConcurrentPromises = 15;
+    const promiseRunning = _.map(_.range(nConcurrentPromises), () => false);
+    const runningPromisesSamples = [];
+    const promiseAtIndex = async index => {
+      promiseRunning[index] = true;
+
+      let nRunningPromises = 0;
+      for (const p of promiseRunning) {
+        if (p) {
+          nRunningPromises++;
+        }
+      }
+      runningPromisesSamples.push(nRunningPromises);
+
+      await sleep(0.2);
+      promiseRunning[index] = false;
+    };
+    await Promise.all(_.map(_.range(nConcurrentPromises), i => lock.lockForPath(() => (promiseAtIndex(i)))));
+    let largest = 0;
+    for (const sample of runningPromisesSamples) {
+      if (sample > largest) {
+        largest = sample;
+      }
+    }
+    largest.should.be.equal(width);
+  });
+
+  it('Exceptions propagated to caller', async () => {
+    let error;
+    try {
+      await lock.lockForPath(async () => {
+        throw new Error();
+      });
+    } catch (e) {
+      error = e;
+    }
+    should.exist(error);
+  });
+
+  it('Exceptions don\'t impact other callers', async () => {
+    const promises = [];
+    promises[0] = lock.lockForPath(async () => {
+      return 0;
+    });
+    let myError = new Error();
+    promises[1] = lock.lockForPath(async () => {
+      throw myError;
+    });
+    promises[2] = lock.lockForPath(async () => {
+      return 2;
+    })
+    const results = [];
+    const errors = [];
+    let idx = 0;
+    for (const promise of promises) {
+      results[idx] = null;
+      errors[idx] = null;
+      try {
+        results[idx] = await promise;
+      } catch (e) {
+        errors[idx] = e;
+      }
+      idx++;
+    }
+    results.should.be.eql([0, null, 2]);
+    errors.should.be.eql([null, myError, null]);
+  });
+});

test/utils/testLock.js

@@ -33,45 +33,4 @@ describe('Lock', () => {
     };
     await Promise.all(_.map(_.range(nConcurrentPromises), i => lock.lockForPath(() => (promiseAtIndex(i)))));
   });
-
-  it('Exceptions propagated to caller', async () => {
-    let error;
-    try {
-      await lock.lockForPath(async () => {
-        throw new Error();
-      });
-    } catch (e) {
-      error = e;
-    }
-    should.exist(error);
-  });
-
-  it('Exceptions don\'t impact other callers', async () => {
-    const promises = [];
-    promises[0] = lock.lockForPath(async () => {
-      return 0;
-    });
-    let myError = new Error();
-    promises[1] = lock.lockForPath(async () => {
-        throw myError;
-    });
-    promises[2] = lock.lockForPath(async () => {
-      return 2;
-    })
-    const results = [];
-    const errors = [];
-    let idx = 0;
-    for (const promise of promises) {
-      results[idx] = null;
-      errors[idx] = null;
-      try {
-        results[idx] = await promise;
-      } catch (e) {
-        errors[idx] = e;
-      }
-      idx++;
-    }
-    results.should.be.eql([0, null, 2]);
-    errors.should.be.eql([null, myError, null]);
-  });
 });