instrumentation: fixed issue with bluebird promises losing transaction state

Author: Bryan Clement

lib/instrumentation/bluebird.js

@@ -0,0 +1,30 @@
+'use strict'
+
+var shimmer = require('../shimmer')
+
+module.exports = function initialize(agent, bluebird) {
+  var tracer = agent.tracer
+
+  shimmer.wrapMethod(
+    bluebird.prototype,
+    'bluebird.prototype',
+    '_then',
+    function wrapThen(original) {
+      return tracer.wrapFunctionNoSegment(original, null, wrapper)
+    }
+  )
+
+  function wrapper(args) {
+    var onResolve = args[0]
+    if (typeof onResolve === 'function') {
+      args[0] = tracer.bindFunction(onResolve)
+    }
+
+    var onReject = args[1]
+    if (typeof onReject === 'function') {
+      args[1] = tracer.bindFunction(onReject)
+    }
+
+    return args
+  }
+}

lib/instrumentations.js

@@ -4,6 +4,7 @@
 module.exports = function instrumentations() {
   return [
     'connect',
+    'bluebird',
     'express',
     'generic-pool',
     'hapi',

lib/transaction/tracer/index.js

@@ -213,7 +213,7 @@ function argSlice(args) {
   return array
 }
 
-function wrapFunctionNoSegment(original, name) {
+function wrapFunctionNoSegment(original, name, wrapper) {
   if (typeof original !== 'function') return original
 
   logger.trace('Wrapping function %s (no segment)', name || original.name || 'anonymous')
@@ -224,10 +224,15 @@ function wrapFunctionNoSegment(original, name) {
   function wrappedFunction() {
     if (!tracer.getTransaction()) return original.apply(this, arguments)
     var args = tracer.slice(arguments)
-    var last = args.length - 1
-    var cb = args[last]
-    if (typeof cb === 'function') {
-      args[last] = tracer.bindFunction(cb)
+
+    if (wrapper === undefined) {
+      var last = args.length - 1
+      var cb = args[last]
+      if (typeof cb === 'function') {
+        args[last] = tracer.bindFunction(cb)
+      }
+    } else {
+      args = wrapper(args)
     }
     return original.apply(this, args)
   }

test/integration/instrumentation/bluebird.tap.js

@@ -0,0 +1,113 @@
+'use strict'
+
+var test = require('tap').test
+var helper = require('../../lib/agent_helper')
+
+
+test('preserves transaction in resolve callback', function(t) {
+  var agent = setupAgent(t)
+  var Promise = require('bluebird')
+
+  var total = 2
+  var finished = 0
+
+  for (var i = 0; i < total; i++) {
+    helper.runInTransaction(agent, function transactionWrapper(transaction) {
+      Promise.resolve().then(function anonymous() {
+        checkTransaction(t, agent, transaction)
+        finished++
+        if (finished === total) t.end()
+      })
+    })
+  }
+})
+
+test('preserves transaction in reject callback', function(t) {
+  var agent = setupAgent(t)
+  var Promise = require('bluebird')
+
+  var total = 2
+  var finished = 0
+
+  for (var i = 0; i < total; i++) {
+    helper.runInTransaction(agent, function transactionWrapper(transaction) {
+      Promise.reject(new Error('some error')).then(
+        function() { /* success */ },
+        function anonymous() {
+          checkTransaction(t, agent, transaction)
+          finished++
+          if (finished === total) t.end()
+        }
+      )
+    })
+  }
+})
+
+test('preserves transaction in catch callback', function(t) {
+  var agent = setupAgent(t)
+  var Promise = require('bluebird')
+
+  var total = 2
+  var finished = 0
+
+  for (var i = 0; i < total; i++) {
+    helper.runInTransaction(agent, function transactionWrapper(transaction) {
+      Promise.reject(new Error('some error'))
+        .catch(function(error) {
+          checkTransaction(t, agent, transaction)
+          finished++
+          if (finished === total) t.end()
+        })
+    })
+  }
+})
+
+test('preserves transaction with chained promises', function(t) {
+  var agent = setupAgent(t)
+  var Promise = require('bluebird')
+
+  helper.runInTransaction(agent, function transactionWrapper(transaction) {
+    Promise.resolve(1).then(function() {
+      return 1
+    })
+    .then(function() {
+      return 2
+    })
+    .then(function() {
+      checkTransaction(t, agent, transaction)
+      t.end()
+    })
+  })
+})
+
+test('preserves transaction with chained promises and catch callback', function(t) {
+  var agent = setupAgent(t)
+  var Promise = require('bluebird')
+
+  helper.runInTransaction(agent, function transactionWrapper(transaction) {
+    Promise.resolve(1).then(function() {
+      return 1
+    })
+    .then(function() {
+      throw new Error('some error')
+    })
+    .catch(function() {
+      checkTransaction(t, agent, transaction)
+      t.end()
+    })
+  })
+})
+
+function setupAgent(t) {
+  var agent = helper.instrumentMockedAgent()
+  t.tearDown(function tearDown() {
+    helper.unloadAgent(agent)
+  })
+  return agent
+}
+
+function checkTransaction(t, agent, transaction) {
+  t.ok(agent.getTransaction() != null, 'there should be a transaction')
+  t.equal(agent.getTransaction().id, transaction.id)
+}
+

test/package.json

@@ -10,6 +10,7 @@
     "JSV": "~4.0.2",
     "architect": "*",
     "async": "*",
+    "bluebird": "*",
     "chai": "^1",
     "connect": "*",
     "cover": "*",