Make it work with async hooks

index.js

@@ -22,6 +22,12 @@ module.exports.isFinished = isFinished
 
 var first = require('ee-first')
 
+var asyncHooks
+try {
+  asyncHooks = require('async_hooks')
+} catch (ignored) {
+}
+
 /**
  * Variables.
  * @private
@@ -43,6 +49,7 @@ var defer = typeof setImmediate === 'function'
  */
 
 function onFinished (msg, listener) {
+  listener = preserveAsyncContext('OnFinishedListener', listener)
   if (isFinished(msg) !== false) {
     defer(listener, null, msg)
     return msg
@@ -195,3 +202,11 @@ function patchAssignSocket (res, callback) {
     callback(socket)
   }
 }
+
+function preserveAsyncContext (asyncResourceType, fn) {
+  if (!asyncHooks) {
+    return fn
+  }
+  var asyncResource = new asyncHooks.AsyncResource(asyncResourceType)
+  return asyncResource.runInAsyncScope.bind(asyncResource, fn, null)
+}

test/test.js

@@ -4,6 +4,13 @@ var http = require('http')
 var net = require('net')
 var onFinished = require('..')
 
+var describeWhenAsyncHooksAreSupported = describe.skip
+try {
+  var asyncHooks = require('async_hooks')
+  describeWhenAsyncHooksAreSupported = typeof asyncHooks.AsyncLocalStorage === 'function' ? describe : describe.skip
+} catch (ignored) {
+}
+
 describe('onFinished(res, listener)', function () {
   it('should invoke listener given an unknown object', function (done) {
     onFinished({}, done)
@@ -44,6 +51,27 @@ describe('onFinished(res, listener)', function () {
     })
   })
 
+  describeWhenAsyncHooksAreSupported('when used with async hooks', function () {
+    it('should maintain async context', function (done) {
+      var asyncLocalStorage = new asyncHooks.AsyncLocalStorage()
+      var store = {
+        contextMaintained: true
+      }
+
+      var server = http.createServer(function (req, res) {
+        asyncLocalStorage.run(store, function () {
+          onFinished(res, function () {
+            assert.strictEqual(asyncLocalStorage.getStore(), store)
+            done()
+          })
+        })
+        setTimeout(res.end.bind(res), 0)
+      })
+
+      sendGet(server)
+    })
+  })
+
   describe('when using keep-alive', function () {
     it('should fire for each response', function (done) {
       var called = false