add support for bluebird (#372)

src/plugins/bluebird.js

@@ -0,0 +1,16 @@
+'use strict'
+
+const tx = require('./util/promise')
+
+module.exports = [
+  {
+    name: 'bluebird',
+    versions: ['2.0.2 - 3'], // 2.0.0 and 2.0.1 were removed from npm
+    patch (Promise, tracer, config) {
+      this.wrap(Promise.prototype, '_then', tx.createWrapThen(tracer, config))
+    },
+    unpatch (Promise) {
+      this.unwrap(Promise.prototype, '_then')
+    }
+  }
+]

src/plugins/index.js

@@ -3,6 +3,7 @@
 module.exports = {
   'amqp10': require('./amqp10'),
   'amqplib': require('./amqplib'),
+  'bluebird': require('./bluebird'),
   'elasticsearch': require('./elasticsearch'),
   'express': require('./express'),
   'graphql': require('./graphql'),

src/plugins/util/promise.js

@@ -0,0 +1,30 @@
+'use strict'
+
+module.exports = {
+  createWrapThen (tracer, config) {
+    return function wrapThen (then) {
+      return function thenWithTrace (onFulfilled, onRejected, onProgress) {
+        arguments[0] = wrapCallback(tracer, onFulfilled)
+        arguments[1] = wrapCallback(tracer, onRejected)
+
+        // not standard but sometimes supported
+        if (onProgress) {
+          arguments[2] = wrapCallback(tracer, onProgress)
+        }
+
+        return then.apply(this, arguments)
+      }
+    }
+  }
+}
+
+function wrapCallback (tracer, callback) {
+  if (typeof callback !== 'function') return callback
+
+  const scope = tracer.scopeManager().active()
+
+  return function () {
+    tracer.scopeManager().activate(scope ? scope.span() : null)
+    return callback.apply(this, arguments)
+  }
+}

src/scope/scope_manager.js

@@ -49,8 +49,10 @@ class ScopeManager {
     let execution = this._active
 
     while (execution !== null) {
-      if (execution.scope()) {
-        return execution.scope()
+      const scope = execution.scope()
+
+      if (scope) {
+        return scope.span() ? scope : null
       }
 
       execution = execution.parent()

test/plugins/bluebird.spec.js

@@ -0,0 +1,104 @@
+'use strict'
+
+const agent = require('./agent')
+const plugin = require('../../src/plugins/bluebird')
+
+wrapIt()
+
+describe('Plugin', () => {
+  let Promise
+  let tracer
+
+  describe('bluebird', () => {
+    withVersions(plugin, 'bluebird', version => {
+      beforeEach(() => {
+        tracer = require('../..')
+      })
+
+      afterEach(() => {
+        return agent.close()
+      })
+
+      describe('without configuration', () => {
+        beforeEach(() => {
+          return agent.load(plugin, 'bluebird')
+            .then(() => {
+              Promise = require(`../../versions/bluebird@${version}`).get()
+            })
+        })
+
+        it('should run the then() callback in context where then() was called', () => {
+          if (process.env.DD_CONTEXT_PROPAGATION === 'false') return
+
+          const span = {}
+          const promise = new Promise((resolve, reject) => {
+            setImmediate(() => {
+              tracer.scopeManager().activate({})
+              resolve()
+            })
+          })
+
+          tracer.scopeManager().activate(span)
+
+          return promise
+            .then(() => {
+              tracer.scopeManager().activate({})
+            })
+            .then(() => {
+              const scope = tracer.scopeManager().active()
+
+              expect(scope).to.not.be.null
+              expect(scope.span()).to.equal(span)
+            })
+        })
+
+        it('should run the catch() callback in context where catch() was called', () => {
+          if (process.env.DD_CONTEXT_PROPAGATION === 'false') return
+
+          const span = {}
+          const promise = new Promise((resolve, reject) => {
+            setImmediate(() => {
+              tracer.scopeManager().activate({})
+              reject(new Error())
+            })
+          })
+
+          tracer.scopeManager().activate(span)
+
+          return promise
+            .catch(err => {
+              tracer.scopeManager().activate({})
+              throw err
+            })
+            .catch(() => {
+              const scope = tracer.scopeManager().active()
+
+              expect(scope).to.not.be.null
+              expect(scope.span()).to.equal(span)
+            })
+        })
+
+        it('should allow to run without a scope if not available when calling then()', () => {
+          if (process.env.DD_CONTEXT_PROPAGATION === 'false') return
+
+          tracer.scopeManager().activate(null)
+
+          const promise = new Promise((resolve, reject) => {
+            setImmediate(() => {
+              tracer.scopeManager().activate({})
+              resolve()
+            })
+          })
+
+          return promise
+            .then(() => {
+              tracer.scopeManager().activate({})
+            })
+            .then(() => {
+              expect(tracer.scopeManager().active()).to.be.null
+            })
+        })
+      })
+    })
+  })
+})

test/scope/scope_manager.spec.js

@@ -241,4 +241,13 @@ describe('ScopeManager', () => {
       asyncHooks.before(1)
     }).not.to.throw()
   })
+
+  it('should allow deactivating a scope', () => {
+    const span = {}
+
+    scopeManager.activate(span)
+    scopeManager.activate(null)
+
+    expect(scopeManager.active()).to.be.null
+  })
 })