cls.js

module.exports = require('./multi-cls.js')();

multi-cls.js

if (!process.addAsyncListener) require('./polyfill.js');

module.exports = function () {

  // Mini CLS library that consumes the new hook API
  var ctx = Object.create(null);
  var stack = [ctx];

  process.addAsyncListener(function () {
    return ctx;
  }, {
    before: add,
    after: remove
  });

  return {
    get: function (key) { return ctx[key]; },
    set: function (key, value) { return ctx[key] = value; },
    scope: scope
  };

  function add(newContext) {
    ctx = newContext || Object.create(ctx);
    stack.push(ctx);
  }

  function remove(oldContext) {
    while (stack.length > 1) {
      if (stack.pop() === oldContext) break;
    }
    ctx = stack[stack.length - 1];
  }

  function scope(fn) {
    var context = add();
    try { fn(); }
    finally { remove(context); }
  }

};

polyfill.js

process.addAsyncListener = addAsyncListener;
var listeners = [];

// Wrap some easy event sources to test things out
setTimeout = wrapFirst(setTimeout);
setImmediate = wrapFirst(setImmediate);
setInterval = wrapFirst(setInterval);
process.nextTick = wrapFirst(process.nextTick);
process.nextTick = wrapFirst(process.nextTick);
var fs = require('fs');
fs.stat = wrapLast(fs.stat);

function addAsyncListener(onAsync, callbackObject) {
  listeners.push({
    onAsync: onAsync,
    callbackObject: callbackObject
  });
}

function wrapFirst(fn) {
  return function () {
    if (typeof arguments[0] === "function") {
      arguments[0] = wrapCallback(arguments[0]);
    }
    return fn.apply(this, arguments);
  };
}

function wrapLast(fn) {
  return function () {
    var index = arguments.length - 1;
    if (typeof arguments[index] === "function") {
      arguments[index] = wrapCallback(arguments[index]);
    }
    return fn.apply(this, arguments);
  }
}

function wrapCallback(original) {
  var list = Array.prototype.slice.call(listeners);
  var length = list.length;
  var hasAny = false, hasErr = false;
  for (var i = 0; i < length; ++i) {
    var obj = list[i].callbackObject;
    if (obj) {
      hasAny = true;
      if (obj.error) hasErr = true;
    }
  }
  return hasAny ? hasErr ? catchyWrap(original, list, length)
                         : normalWrap(original, list, length)
                : noWrap(original, list, length);
}

function catchyWrap(original, list, length) {
  var data = new Array(length);
  for (var i = 0; i < length; ++i) {
    var listener = list[i];
    data[i] = listener.onAsync();
  }
  return function () {
    var i, obj;
    for (var i = 0; i < length; ++i) {
      obj = list[i].callbackObject;
      if (obj && obj.before) obj.before(data[i]);
    }
    try {
      return original.apply(this, arguments);
    }
    catch (err) {
      for (i = 0; i < length; ++i) {
        obj = list[i].callbackObject;
        if (obj && obj.after) obj.after(data[i]);
      }
    }
    finally {
      for (i = 0; i < length; ++i) {
        obj = list[i].callbackObject;
        if (obj && obj.after) obj.after(data[i]);
      }
      for (i = 0; i < length; ++i) {
        obj = list[i].callbackObject;
        if (obj && obj.done) obj.done(data[i]);
      }
    }
  }
}

function normalWrap(original, list, length) {
  var data = new Array(length);
  for (var i = 0; i < length; ++i) {
    var listener = list[i];
    data[i] = listener.onAsync();
  }
  return function () {
    var i, obj;
    for (var i = 0; i < length; ++i) {
      obj = list[i].callbackObject;
      if (obj && obj.before) obj.before(data[i]);
    }
    try {
      return original.apply(this, arguments);
    }
    finally {
      for (i = 0; i < length; ++i) {
        obj = list[i].callbackObject;
        if (obj && obj.after) obj.after(data[i]);
      }
      for (i = 0; i < length; ++i) {
        obj = list[i].callbackObject;
        if (obj && obj.done) obj.done(data[i]);
      }
    }
  }
}

function noWrap(original, list, length) {
  for (var i = 0; i < length; ++i) {
    list[i].onAsync();
  }
  return original;
}

test.js

var cls = require('./cls.js');

// Mini test that uses the mini CLS library
var fs = require('fs');
var assert = require('assert');
cls.set("value", 0);
console.log("before", cls.get("value"))
cls.scope(function () {
  var left = 0;
  cls.set("value", 42);
  console.log("inside before", cls.get("value"))
  setTimeout(after("setTimeout"), 100);
  setImmediate(after("setImmediate"), 100);
  process.nextTick(after("process.nextTick"));
  fs.stat(__filename, after("fs.stat"));
  console.log("inside after", cls.get("value"))
  process.on('exit', function () {
    if (left) throw new Error(left + " callbacks failed to fire");
  });

  function after(name) {
    left++;
    return function () {
      left--;
      console.log("after " + name + " tick", cls.get("value"));
      assert.equal(cls.get("value"), 42);
    }
  }
});
console.log("afterfn", cls.get("value"))
cls.set("value", 1);