From 4d2ab543e63517a32fd7a4d25fe29217cca7683e Mon Sep 17 00:00:00 2001 From: Bryan Clement Date: Mon, 12 Mar 2018 20:52:20 -0700 Subject: [PATCH] allow for next logic in promise based middleware --- lib/shim/webframework-shim.js | 56 ++++++++++++++++++++---- test/unit/shim/webframework-shim.test.js | 40 ++++++++++++++++- 2 files changed, 87 insertions(+), 9 deletions(-) diff --git a/lib/shim/webframework-shim.js b/lib/shim/webframework-shim.js index 39ccaabc7a..36e7fbc3dc 100644 --- a/lib/shim/webframework-shim.js +++ b/lib/shim/webframework-shim.js @@ -818,7 +818,7 @@ function _recordMiddleware(shim, middleware, spec) { // if context leaves this middleware. var nextWrapper = null if (shim.isFunction(spec.next)) { - nextWrapper = _makeNextBinder(route, txInfo, spec.next, isErrorWare) + nextWrapper = _makeNextBinder(route, txInfo, spec.next, isErrorWare, false) } else { var nextIdx = shim.normalizeIndex(args.length, spec.next) if (nextIdx !== null && args[nextIdx] instanceof Function) { @@ -828,7 +828,8 @@ function _recordMiddleware(shim, middleware, spec) { function wrapNext(s, f, n, _args, wrap) { wrap(_args, nextIdx) }, - isErrorWare + isErrorWare, + false ) } } @@ -895,16 +896,37 @@ function _recordMiddleware(shim, middleware, spec) { recorder = _makeMiddlewareRecorder(shim, metricName + stackPath) } - function pushSegment(shim, fn, _name, segment) { + // The next callback style can still apply to promise based + // middleware (e.g. koa). In this case we would like to remove the + // path for the current executing middleware, then readd it once the + // next callback is done (either asynchronously or after the + // returned promise is resolved). + var nextWrapper = function pushSegment(shim, _fn, _name, segment) { txInfo.segmentStack.push(segment) } + if (shim.isFunction(spec.next)) { + nextWrapper = _makeNextBinder(route, txInfo, spec.next, isErrorWare, true) + } else { + var nextIdx = shim.normalizeIndex(args.length, spec.next) + if (nextIdx !== null && args[nextIdx] instanceof Function) { + nextWrapper = _makeNextBinder( + route, + txInfo, + function wrapNext(s, f, n, _args, wrap) { + wrap(_args, nextIdx) + }, + isErrorWare, + true + ) + } + } // Finally, return the segment descriptor. return { name: segmentName, parent: txInfo.segmentStack[txInfo.segmentStack.length - 1], promise: spec.promise, - callback: pushSegment, + callback: nextWrapper, recorder: recorder, parameters: params, after: function afterExec(shim, _fn, _name, err, result) { @@ -934,7 +956,7 @@ function _makeGetReq(shim, req) { } } -function _makeNextBinder(route, txInfo, wrapNext, isErrorWare) { +function _makeNextBinder(route, txInfo, wrapNext, isErrorWare, isPromise) { return function bindNext(shim, fn, _name, segment, args) { if (!segment) { return @@ -956,9 +978,27 @@ function _makeNextBinder(route, txInfo, wrapNext, isErrorWare) { segment.transaction.nameState.popPath(route) } - txInfo.segmentStack.pop() - segment.end() - return original.apply(this, arguments) + // The next call does not signify the end of the segment + // calling next in the promise case. Keep the segment on the + // stack and wait for its promise to be resolved to end it. + if (!isPromise) { + txInfo.segmentStack.pop() + segment.end() + } + var ret = original.apply(this, arguments) + + if (isPromise && shim.isPromise(ret)) { + // After the next call has resolved, we should reinstate the + // segment responsible for calling next in case there is + // more work to do in that scope. + return ret.then(function onNextFinish(v) { + segment.transaction.nameState.appendPath(route) + txInfo.segmentStack.push(segment) + return v + }) + } + + return ret }, shim.getSegment() || segment) // Bind to parent. }) } diff --git a/test/unit/shim/webframework-shim.test.js b/test/unit/shim/webframework-shim.test.js index 94b03b28fe..ceede4a2a3 100644 --- a/test/unit/shim/webframework-shim.test.js +++ b/test/unit/shim/webframework-shim.test.js @@ -724,12 +724,21 @@ describe('WebFrameworkShim', function() { beforeEach(function() { unwrappedTimeout = shim.unwrap(setTimeout) - middleware = function(_req, err) { + middleware = function(_req, err, next) { segment = shim.getSegment() return new Promise(function(resolve, reject) { unwrappedTimeout(function() { try { expect(txInfo.transaction.nameState.getPath()).to.equal('/foo/bar') + if (next) { + return next().then(function() { + expect(txInfo.transaction.nameState.getPath()).to.equal('/foo/bar') + resolve() + }, function(e) { + expect(txInfo.transaction.nameState.getPath()).to.equal('/') + resolve() + }) + } if (err) { throw err } else { @@ -744,6 +753,7 @@ describe('WebFrameworkShim', function() { wrapped = shim.recordMiddleware(middleware, { route: '/foo/bar', + next: shim.LAST, promise: true }) }) @@ -785,6 +795,34 @@ describe('WebFrameworkShim', function() { }) }) }) + + it('should pop the name of the handler off when next is called', function() { + return helper.runInTransaction(agent, function(tx) { + tx.nameState.appendPath('/') + txInfo.transaction = tx + return wrapped(req, null, function next() { + expect(tx.nameState.getPath()).to.equal('/') + return new Promise(function(resolve) { + expect(agent.tracer.getTransaction()).to.equal(tx) + resolve() + }) + }) + }) + }) + + it('should have the right name when the next handler errors', function() { + return helper.runInTransaction(agent, function(tx) { + tx.nameState.appendPath('/') + txInfo.transaction = tx + return wrapped(req, null, function next() { + expect(tx.nameState.getPath()).to.equal('/') + return new Promise(function(resolve, reject) { + expect(agent.tracer.getTransaction()).to.equal(tx) + reject() + }) + }) + }) + }) }) describe('when wrapping errorware', function() {