Implement process.addAsyncListener polyfill

bootstrap.js

@@ -1,44 +1,15 @@
 'use strict';
+if (process.addAsyncListener) {
+  throw new Error("Don't require polyfill unless needed");
+}
+process.addAsyncListener = addAsyncListener;
+var listeners = [];
 
-var cls      = require('continuation-local-storage')
-  , shimmer  = require('shimmer')
+var shimmer  = require('shimmer')
   , wrap     = shimmer.wrap
   , massWrap = shimmer.massWrap
   ;
 
-var slice = [].slice;
-function each(obj, callback) {
-  var keys = Object.keys(obj);
-  for (var i = 0, l = keys.length; i < l; ++i) {
-    var key = keys[i];
-    callback(key, obj[key]);
-  }
-}
-
-function wrapCallback(callback) {
-  // Get the currently active contexts in all the namespaces.
-  var contexts = {};
-  each(process.namespaces, function (name, namespace) {
-    contexts[name] = namespace.active;
-  });
-
-  // Return a callback that enters all the saved namespaces when called.
-  return function () {
-    var namespaces = process.namespaces;
-    each(contexts, function (name, context) {
-      namespaces[name].enter(context);
-    });
-    try {
-      return callback.apply(this, arguments);
-    }
-    finally {
-      each(contexts, function (name, context) {
-        namespaces[name].exit(context);
-      });
-    }
-  };
-}
-
 var net = require('net');
 wrap(net.Server.prototype, "_listen2", function (original) {
   return function () {
@@ -61,39 +32,6 @@ wrap(net.Socket.prototype, "connect", function (original) {
   };
 });
 
-// Shim activator for functions that have callback last
-function activator(fn) {
-  return function () {
-    var args = slice.call(arguments);
-    var callback = args[args.length - 1];
-
-    // If there is no callback, there will be no continuation to trap.
-    if (typeof callback !== "function") {
-      return fn.apply(this, arguments);
-    }
-
-    // Wrap the callback so that the continuation keeps the current contexts.
-    args[args.length - 1] = wrapCallback(callback);
-    return fn.apply(this, args);
-  };
-}
-
-// Shim activator for functions that have callback first
-function activatorFirst(fn) {
-  return function () {
-    var args = slice.call(arguments);
-    var callback = args[0];
-
-    // If there is no callback, there will be no continuation to trap.
-    if (typeof callback !== "function") {
-      return fn.apply(this, arguments);
-    }
-
-    // Wrap the callback so that the continuation keeps the current contexts.
-    args[0] = wrapCallback(callback);
-    return fn.apply(this, args);
-  };
-}
 
 var processors = ['nextTick'];
 if (process._nextDomainTick) processors.push('_nextDomainTick');
@@ -183,5 +121,137 @@ if (fs.lchmod) wrap(fs, 'lchmod', activator);
 // only wrap ftruncate in versions of node that have it
 if (fs.ftruncate) wrap(fs, 'ftruncate', activator);
 
-// PUBLIC API STARTS HERE: there isn't much of one
-module.exports = cls;
+// Wrap zlib streams
+var zProto = Object.getPrototypeOf(require('zlib').Deflate.prototype);
+wrap(zProto, "_transform", activator);
+
+// Wrap Crypto
+var crypto;
+try { crypto = require('crypto'); }
+catch (err) { }
+if (crypto) {
+  massWrap(crypto, [
+    "pbkdf2",
+    "randomBytes",
+    "pseudoRandomBytes",
+  ], activator);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+// Polyfilled version of process.addAsyncListener
+function addAsyncListener(onAsync, callbackObject) {
+  listeners.push({
+    onAsync: onAsync,
+    callbackObject: callbackObject
+  });
+}
+
+// Shim activator for functions that have callback last
+function activator(fn) {
+  return function () {
+    var index = arguments.length - 1;
+    if (typeof arguments[index] === "function") {
+      arguments[index] = wrapCallback(arguments[index]);
+    }
+    return fn.apply(this, arguments);
+  }
+}
+
+// Shim activator for functions that have callback first
+function activatorFirst(fn) {
+  return function () {
+    if (typeof arguments[0] === "function") {
+      arguments[0] = wrapCallback(arguments[0]);
+    }
+    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 runSetup(list, length) {
+  var data = new Array(length);
+  for (var i = 0; i < length; ++i) {
+    var listener = list[i];
+    data[i] = listener.onAsync();
+  }
+  return data;
+}
+
+function runBefore(data, list, length) {
+  for (var i = 0; i < length; ++i) {
+    var obj = list[i].callbackObject;
+    if (obj && obj.before) obj.before(data[i]);
+  }
+}
+
+function runError(data, list, length) {
+  for (i = 0; i < length; ++i) {
+    obj = list[i].callbackObject;
+    if (obj && obj.after) obj.after(data[i]);
+  }
+}
+
+function runAfter(data, list, length) {
+  var i, obj;
+  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 catchyWrap(original, list, length) {
+  var data = runSetup(list, length);
+  return function () {
+    runBefore(data, list, length);
+    try {
+      return original.apply(this, arguments);
+    }
+    catch (err) {
+      runError(data, list, length);
+    }
+    finally {
+      runAfter(data, list, length);
+    }
+  }
+}
+
+function normalWrap(original, list, length) {
+  var data = runSetup(list, length);
+  return function () {
+    runBefore(data, list, length);
+    try {
+      return original.apply(this, arguments);
+    }
+    finally {
+      runAfter(data, list, length);
+    }
+  }
+}
+
+function noWrap(original, list, length) {
+  for (var i = 0; i < length; ++i) {
+    list[i].onAsync();
+  }
+  return original;
+}
+

test/async-context.tap.js

@@ -1,6 +1,5 @@
 'use strict';
-
-require('../bootstrap.js');
+if (!process.addAsyncListener) require('../bootstrap.js');
 
 var tap             = require('tap')
   , test            = tap.test

test/crypto.tap.js

@@ -0,0 +1,69 @@
+'use strict';
+if (!process.addAsyncListener) require('../bootstrap.js');
+
+var tap             = require('tap')
+  , test            = tap.test
+  , createNamespace = require('continuation-local-storage').createNamespace
+  ;
+
+var crypto;
+try { crypto = require('crypto'); }
+catch (err) {}
+if (crypto) {
+
+  test("continuation-local state with crypto.randomBytes", function (t) {
+    t.plan(1);
+
+    var namespace = createNamespace('namespace');
+    namespace.set('test', 0xabad1dea);
+
+    t.test("deflate", function (t) {
+      namespace.run(function () {
+        namespace.set('test', 42);
+        crypto.randomBytes(100, function (err, bytes) {
+          if (err) throw err;
+          t.equal(namespace.get('test'), 42, "mutated state was preserved");
+          t.end();
+        });
+      });
+    });
+  });
+
+  test("continuation-local state with crypto.pseudoRandomBytes", function (t) {
+    t.plan(1);
+
+    var namespace = createNamespace('namespace');
+    namespace.set('test', 0xabad1dea);
+
+    t.test("deflate", function (t) {
+      namespace.run(function () {
+        namespace.set('test', 42);
+        crypto.pseudoRandomBytes(100, function (err, bytes) {
+          if (err) throw err;
+          t.equal(namespace.get('test'), 42, "mutated state was preserved");
+          t.end();
+        });
+      });
+    });
+  });
+
+  test("continuation-local state with crypto.pbkdf2", function (t) {
+    t.plan(1);
+
+    var namespace = createNamespace('namespace');
+    namespace.set('test', 0xabad1dea);
+
+    t.test("deflate", function (t) {
+      namespace.run(function () {
+        namespace.set('test', 42);
+        crypto.pbkdf2("s3cr3tz", "451243", 10, 40, function (err, key) {
+          if (err) throw err;
+          t.equal(namespace.get('test'), 42, "mutated state was preserved");
+          t.end();
+        });
+      });
+    });
+  });
+
+}
+

test/error-handling.tap.js

@@ -1,8 +1,9 @@
 'use strict';
+if (!process.addAsyncListener) require('../bootstrap.js');
 
 var domain = require('domain')
   , test   = require('tap').test
-  , cls    = require('../bootstrap.js')
+  , cls    = require('continuation-local-storage')
   ;
 
 test("continuation-local storage glue with a throw in the continuation chain",

test/makecallback-dns.tap.js

@@ -1,6 +1,5 @@
 'use strict';
-
-require('../bootstrap.js');
+if (!process.addAsyncListener) require('../bootstrap.js');
 
 var dns             = require('dns')
   , tap             = require('tap')

test/makecallback-fs.tap.js

@@ -1,6 +1,5 @@
 'use strict';
-
-require('../bootstrap.js');
+if (!process.addAsyncListener) require('../bootstrap.js');
 
 var fs              = require('fs')
   , path            = require('path')

test/monkeypatching.tap.js

@@ -1,7 +1,9 @@
 'use strict';
-
 var test = require('tap').test;
 
+if (!process.addAsyncListener) {
+
+
 test("overwriting startup.processNextTick", function (t) {
   t.plan(2);
 
@@ -62,3 +64,6 @@ test("overwriting setImmediate", function (t) {
    */
 
 });
+
+
+}

test/net-events.tap.js

@@ -1,6 +1,5 @@
 'use strict';
-
-require('../bootstrap.js');
+if (!process.addAsyncListener) require('../bootstrap.js');
 
 var net             = require('net')
   , tap             = require('tap')

test/timers.tap.js

@@ -1,6 +1,5 @@
 'use strict';
-
-require('../bootstrap.js');
+if (!process.addAsyncListener) require('../bootstrap.js');
 
 var tap             = require('tap')
   , test            = tap.test

test/zlib.tap.js

@@ -0,0 +1,27 @@
+'use strict';
+if (!process.addAsyncListener) require('../bootstrap.js');
+
+var tap             = require('tap')
+  , test            = tap.test
+  , createNamespace = require('continuation-local-storage').createNamespace
+  ;
+
+var zlib = require('zlib');
+
+test("continuation-local state with zlib", function (t) {
+  t.plan(1);
+
+  var namespace = createNamespace('namespace');
+  namespace.set('test', 0xabad1dea);
+
+  t.test("deflate", function (t) {
+    namespace.run(function () {
+      namespace.set('test', 42);
+      zlib.deflate(new Buffer("Goodbye World"), function (err, deflated) {
+        if (err) throw err;
+        t.equal(namespace.get('test'), 42, "mutated state was preserved");
+        t.end();
+      });
+    });
+  });
+});