diff --git a/lib/instrumentation/modules/elasticsearch.js b/lib/instrumentation/modules/elasticsearch.js index 1cf37fa28a..705e76c6d2 100644 --- a/lib/instrumentation/modules/elasticsearch.js +++ b/lib/instrumentation/modules/elasticsearch.js @@ -55,14 +55,34 @@ module.exports = function (elasticsearch, agent, { enabled }) { } return original.apply(this, args) } else { - return original.apply(this, arguments) - .then(function (originalP) { + const originalPromise = original.apply(this, arguments) + + const descriptors = Object.getOwnPropertyDescriptors(originalPromise) + delete descriptors.domain + + const inspectedPromise = originalPromise + .then(function (value) { span.end() - return originalP + return value }, function (err) { span.end() throw err }) + + Object.defineProperties(inspectedPromise, descriptors) + + // we have to properly end the span when user aborts the request + shimmer.wrap(inspectedPromise, 'abort', function wrapAbort (originalAbort) { + return function wrappedAbort () { + if (span.ended) return + agent.logger.debug('intercepted call to elasticsearch.Transport.request.abort %o', { id: id, method: method, path: path }) + const originalReturn = originalAbort.apply(this, args) + span.end() + return originalReturn + } + }) + + return inspectedPromise } } else { agent.logger.debug('could not instrument elasticsearch request %o', { id: id }) diff --git a/test/instrumentation/modules/elasticsearch.js b/test/instrumentation/modules/elasticsearch.js index 64e683f77f..9b93d05af2 100644 --- a/test/instrumentation/modules/elasticsearch.js +++ b/test/instrumentation/modules/elasticsearch.js @@ -63,6 +63,23 @@ test('client.search with callback', function userLandCode (t) { }) }) +test('client.search with abort', function userLandCode (t) { + resetAgent(3, done(t, 'POST', '/_search', '{"q":"pants"}', true)) + + agent.startTransaction('foo') + + var client = new elasticsearch.Client({ host: host }) + var query = { q: 'pants' } + + var req = client.search(query) + + setImmediate(() => { + req.abort() + agent.endTransaction() + agent.flush() + }) +}) + if (semver.satisfies(pkg.version, '>= 10')) { test('client.searchTemplate with callback', function userLandCode (t) { var body = { @@ -167,7 +184,7 @@ test('client.count with callback', function userLandCode (t) { }) var queryRegexp = /_((search|msearch)(\/template)?|count)$/ -function done (t, method, path, query) { +function done (t, method, path, query, abort = false) { return function (data, cb) { t.strictEqual(data.transactions.length, 1) t.strictEqual(data.spans.length, 2) @@ -220,14 +237,22 @@ function done (t, method, path, query) { } t.ok(span1.timestamp > span2.timestamp, 'http span should start after elasticsearch span') - t.ok(span1.timestamp + span1.duration * 1000 < span2.timestamp + span2.duration * 1000, 'http span should end before elasticsearch span') + if (abort) { + t.ok(span1.timestamp + span1.duration * 1000 > span2.timestamp + span2.duration * 1000, 'http span should end after elasticsearch span when req is aborted') + } else { + t.ok(span1.timestamp + span1.duration * 1000 < span2.timestamp + span2.duration * 1000, 'http span should end before elasticsearch span') + } t.end() } } -function resetAgent (cb) { +function resetAgent (expected, cb) { + if (typeof expected === 'function') { + cb = expected + expected = 3 + } agent._instrumentation.currentTransaction = null - agent._transport = mockClient(3, cb) + agent._transport = mockClient(expected, cb) agent.captureError = function (err) { throw err } }