Simple promise-based lock implementation

promise-lock.js

'use strict';

module.exports = {
  createSingleLock: createSingleLock,
  createKeyedLock: createKeyedLock
};

// Maintains a list of locks referenced by keys
function createKeyedLock() {
  const locks = {};

  return function execute(id, fn) {
    let entry;
    if (locks[id]) {
      entry = locks[id];
    } else {
      entry = locks[id] = {
        pending: 0,
        lock: createSingleLock()
      };
    }

    entry.pending += 1;
    const decrementPending = () => {
      entry.pending -= 1;
      if (entry.pending === 0) {
        locks[id] = undefined;
      }
    };

    return entry.lock(fn)
      .then((r) => {
        decrementPending();
        return r;
      }, (err) => {
        decrementPending();
        throw err;
      })
  };
}

function createSingleLock() {
  const queue = [];
  let locked = false;

  return function execute(fn) {
    return acquire()
      .then(fn)
      .then((r) => {
        release();
        return r;
      }, (err) => {
        release();
        throw err;
      })
  };

  function acquire() {
    if (locked) {
      return new Promise((resolve) => queue.push(resolve));
    } else {
      locked = true;
      return Promise.resolve();
    }
  }

  function release() {
    const next = queue.shift();
    if (next) {
      next();
    } else {
      locked = false;
    }
  }
}

promise-lock.test.js

'use strict';

const Bluebird = require('bluebird');
const promiseLock = require('./index');

describe('#createSingleLock', () => {
  it('Should only allow one call at a time to a given function', () => {

    const lock = promiseLock.createSingleLock();

    const startTimes = {};
    const endTimes = {};

    // f is our function that we expect to only be executed one at a time.
    const f = (x) => {
      return lock(() => {
        startTimes[x] = Date.now();
        return Bluebird.delay(25)
          .then(function () {
            endTimes[x] = Date.now();
            return x;
          });
      });
    };

    const startTime = Date.now();
    return Bluebird.all([
      f('a'),
      f('b'),
      f('c')
    ]).then(function (results) {
      expect(results).toEqual(['a', 'b', 'c']); // Make sure each function returns the right value

      expect(endTimes['a']).toBeGreaterThan(startTimes['a']); // Sanity check
      expect(startTimes['b']).toBeGreaterThanOrEqual(endTimes['a']); // B needs to start after A has ended
      expect(startTimes['c']).toBeGreaterThanOrEqual(endTimes['b']); // C needs to start after B has ended

      // If we actually performed this in sequence, then the total time should be at
      //  least (time for one function) * (time per function) = 25*3
      expect(Date.now() - startTime).toBeGreaterThan(75);
    });
  });

  it('Should handle one of the calls failing', () => {

    const lock = promiseLock.createSingleLock();

    const f = (x) => {
      return lock(() => {
        return Bluebird.delay(25).then(() => {
          if (x > 1) {
            throw new Error('x was greater than 1');
          } else {
            return x;
          }
        });
      });
    };

    const p1 = Bluebird.resolve(f(2)) // This we expect to fail
      .reflect()
      .then((res) => {
        const err = res.reason();
        expect(err.message).toEqual('x was greater than 1');
      });
    const p2 = f(1); // This should succeed

    return Bluebird.all([p1, p2]);
  });
});

describe('#createKeyedLock', () => {
  it('Should allow concurrent execution for different keys', function() {
    const startTimes = {};
    const endTimes = {};

    const lock = promiseLock.createKeyedLock();

    const f = function(id, x) {
      return lock(id, function() {
        startTimes[x] = Date.now();
        return Bluebird.delay(25).then(function() {
          endTimes[x] = Date.now();
          return x;
        });
      });
    };

    return Bluebird.all([
      f('id1', 'A'),
      f('id1', 'B'),
      f('id2', 'C'),
      f('id2', 'D')
    ]).then(function(results) {
      expect(results).toEqual(['A', 'B', 'C', 'D']); // Make sure each function returns the right value

      expect(endTimes.A).toBeGreaterThan(startTimes.A); // Sanity check
      expect(startTimes.B).toBeGreaterThanOrEqual(endTimes.A); // B needs to start after A has ended because they share the same ID
      expect(startTimes.C).toBeLessThan(startTimes.B); // C should have started before B started because it's a different ID
      expect(startTimes.D).toBeGreaterThanOrEqual(startTimes.C); // D needs to start after C has ended
    });
  });
});

SWEET!