allow for next logic in promise based middleware

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.
       })
     }

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() {