Make it work with async hooks

HISTORY.md

@@ -1,3 +1,8 @@
+unreleased
+==========
+
+  * Prevent loss of async hooks context
+
 2.3.0 / 2015-05-26
 ==================
 

index.js

@@ -20,6 +20,7 @@ module.exports.isFinished = isFinished
  * @private
  */
 
+var asyncHooks = tryRequireAsyncHooks()
 var first = require('ee-first')
 
 /**
@@ -49,7 +50,7 @@ function onFinished (msg, listener) {
   }
 
   // attach the listener to the message
-  attachListener(msg, listener)
+  attachListener(msg, wrap(listener))
 
   return msg
 }
@@ -195,3 +196,33 @@ function patchAssignSocket (res, callback) {
     callback(socket)
   }
 }
+
+/**
+ * Try to require async_hooks
+ * @private
+ */
+
+function tryRequireAsyncHooks () {
+  try {
+    return require('async_hooks')
+  } catch (e) {
+    /* istanbul ignore next */
+    return {}
+  }
+}
+
+/**
+ * Wrap function with async resource
+ * @private
+ */
+
+function wrap (fn) {
+  if (!asyncHooks.AsyncResource) {
+    /* istanbul ignore next */
+    return fn
+  }
+
+  // AsyncResource.bind static method backported
+  var res = new asyncHooks.AsyncResource(fn.name || 'bound-anonymous-fn')
+  return res.runInAsyncScope.bind(res, fn, null)
+}

test/test.js

@@ -1,9 +1,14 @@
 
 var assert = require('assert')
+var asyncHooks = tryRequire('async_hooks')
 var http = require('http')
 var net = require('net')
 var onFinished = require('..')
 
+var describeAsyncHooks = typeof asyncHooks.AsyncLocalStorage === 'function'
+  ? describe
+  : describe.skip
+
 describe('onFinished(res, listener)', function () {
   it('should invoke listener given an unknown object', function (done) {
     onFinished({}, done)
@@ -32,15 +37,57 @@ describe('onFinished(res, listener)', function () {
       sendGet(server)
     })
 
-    it('should fire when called after finish', function (done) {
-      var server = http.createServer(function (req, res) {
-        onFinished(res, function () {
-          onFinished(res, done)
+    describe('when called after finish', function () {
+      it('should fire when called after finish', function (done) {
+        var server = http.createServer(function (req, res) {
+          onFinished(res, function () {
+            onFinished(res, done)
+          })
+          setTimeout(res.end.bind(res), 0)
         })
-        setTimeout(res.end.bind(res), 0)
+
+        sendGet(server)
       })
 
-      sendGet(server)
+      describeAsyncHooks('when async local storage', function () {
+        it('should presist store in callback', function (done) {
+          var asyncLocalStorage = new asyncHooks.AsyncLocalStorage()
+          var store = { foo: 'bar' }
+
+          var server = http.createServer(function (req, res) {
+            onFinished(res, function () {
+              asyncLocalStorage.run(store, function () {
+                onFinished(res, function () {
+                  assert.strictEqual(asyncLocalStorage.getStore().foo, 'bar')
+                  done()
+                })
+              })
+            })
+            setTimeout(res.end.bind(res), 0)
+          })
+
+          sendGet(server)
+        })
+      })
+    })
+
+    describeAsyncHooks('when async local storage', function () {
+      it('should presist store in callback', function (done) {
+        var asyncLocalStorage = new asyncHooks.AsyncLocalStorage()
+        var store = { foo: 'bar' }
+
+        var server = http.createServer(function (req, res) {
+          asyncLocalStorage.run(store, function () {
+            onFinished(res, function () {
+              assert.strictEqual(asyncLocalStorage.getStore().foo, 'bar')
+              done()
+            })
+          })
+          setTimeout(res.end.bind(res), 0)
+        })
+
+        sendGet(server)
+      })
     })
   })
 
@@ -385,16 +432,60 @@ describe('onFinished(req, listener)', function () {
       sendGet(server)
     })
 
-    it('should fire when called after finish', function (done) {
-      var server = http.createServer(function (req, res) {
-        onFinished(req, function () {
-          onFinished(req, done)
+    describe('when called after finish', function () {
+      it('should fire when called after finish', function (done) {
+        var server = http.createServer(function (req, res) {
+          onFinished(req, function () {
+            onFinished(req, done)
+          })
+          req.resume()
+          setTimeout(res.end.bind(res), 0)
         })
-        req.resume()
-        setTimeout(res.end.bind(res), 0)
+
+        sendGet(server)
       })
 
-      sendGet(server)
+      describeAsyncHooks('when async local storage', function () {
+        it('should presist store in callback', function (done) {
+          var asyncLocalStorage = new asyncHooks.AsyncLocalStorage()
+          var store = { foo: 'bar' }
+
+          var server = http.createServer(function (req, res) {
+            onFinished(req, function () {
+              asyncLocalStorage.run(store, function () {
+                onFinished(req, function () {
+                  assert.strictEqual(asyncLocalStorage.getStore().foo, 'bar')
+                  done()
+                })
+              })
+            })
+            req.resume()
+            setTimeout(res.end.bind(res), 0)
+          })
+
+          sendGet(server)
+        })
+      })
+    })
+
+    describeAsyncHooks('when async local storage', function () {
+      it('should presist store in callback', function (done) {
+        var asyncLocalStorage = new asyncHooks.AsyncLocalStorage()
+        var store = { foo: 'bar' }
+
+        var server = http.createServer(function (req, res) {
+          asyncLocalStorage.run(store, function () {
+            onFinished(req, function () {
+              assert.strictEqual(asyncLocalStorage.getStore().foo, 'bar')
+              done()
+            })
+          })
+          req.resume()
+          setTimeout(res.end.bind(res), 0)
+        })
+
+        sendGet(server)
+      })
     })
   })
 
@@ -1068,6 +1159,14 @@ function sendGet (server) {
   })
 }
 
+function tryRequire (name) {
+  try {
+    return require(name)
+  } catch (e) {
+    return {}
+  }
+}
+
 function writeRequest (socket, chunked) {
   socket.write('GET / HTTP/1.1\r\n')
   socket.write('Host: localhost\r\n')