diff --git a/index.js b/index.js index a6724cf..ac5e0eb 100644 --- a/index.js +++ b/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) +} diff --git a/test/test.js b/test/test.js index b003662..77c985d 100644 --- a/test/test.js +++ b/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