update bluebird instrumentation to account for getNewLibraryCopy

packages/datadog-plugin-bluebird/src/index.js

@@ -2,7 +2,57 @@
 
 const tx = require('../../dd-trace/src/plugins/util/promise')
 
+const DD_LIB_COPIES = '_datadog_library_copies'
+
+function createGetNewLibraryCopyWrap (tracer, config, originalLib, shim) {
+  return function wrapGetNewLibraryCopy (getNewLibraryCopy) {
+    return function getNewLibraryCopyWithTrace () {
+      const libraryCopy = getNewLibraryCopy.apply(this, arguments)
+      shim.wrap(libraryCopy.prototype, '_then', tx.createWrapThen(tracer, config))
+      addToLibraryCopies(originalLib, libraryCopy)
+      return libraryCopy
+    }
+  }
+}
+
+function addToLibraryCopies (originalLib, libraryCopy) {
+  let libraryCopies = originalLib[DD_LIB_COPIES]
+
+  if (!libraryCopies) {
+    libraryCopies = new Set()
+
+    Object.defineProperty(originalLib, DD_LIB_COPIES, {
+      writable: true,
+      configurable: true,
+      value: libraryCopies
+    })
+  }
+
+  libraryCopies.add(libraryCopy)
+}
+
+function unwrapLibraryCopies (originalLib, shim) {
+  const libraryCopies = originalLib[DD_LIB_COPIES]
+
+  if (libraryCopies) {
+    libraryCopies.forEach(libraryCopy => shim.unwrap(libraryCopy.prototype, '_then'))
+    libraryCopies.clear()
+    delete originalLib[DD_LIB_COPIES]
+  }
+}
+
 module.exports = [
+  {
+    name: 'bluebird',
+    versions: ['^2.11.0', '^3.4.1'],
+    patch (Promise, tracer, config) {
+      this.wrap(Promise, 'getNewLibraryCopy', createGetNewLibraryCopyWrap(tracer, config, Promise, this))
+    },
+    unpatch (Promise) {
+      this.unwrap(Promise, 'getNewLibraryCopy')
+      unwrapLibraryCopies(Promise, this)
+    }
+  },
   {
     name: 'bluebird',
     versions: ['>=2.0.2'], // 2.0.0 and 2.0.1 were removed from npm

packages/datadog-plugin-bluebird/test/index.spec.js

@@ -3,3 +3,5 @@
 const assertPromise = require('../../dd-trace/test/plugins/promise')
 
 assertPromise('bluebird')
+
+assertPromise('bluebird', bluebird => bluebird.getNewLibraryCopy(), '^2.11.0 || ^3.4.1')

packages/datadog-plugin-promise-js/src/index.js

@@ -7,10 +7,14 @@ module.exports = [
     name: 'promise-js',
     versions: ['>=0.0.3'],
     patch (Promise, tracer, config) {
-      this.wrap(Promise.prototype, 'then', tx.createWrapThen(tracer, config))
+      if (Promise !== global.Promise) {
+        this.wrap(Promise.prototype, 'then', tx.createWrapThen(tracer, config))
+      }
     },
     unpatch (Promise) {
-      this.unwrap(Promise.prototype, 'then')
+      if (Promise !== global.Promise) {
+        this.unwrap(Promise.prototype, 'then')
+      }
     }
   }
 ]

packages/dd-trace/test/plugins/promise.js

@@ -1,8 +1,9 @@
 'use strict'
 
-module.exports = (name, factory) => {
+module.exports = (name, factory, versionRange) => {
   const agent = require('./agent')
   const plugin = require(`../../../datadog-plugin-${name}/src`)
+  const semver = require('semver')
 
   wrapIt()
 
@@ -12,14 +13,12 @@ module.exports = (name, factory) => {
 
     describe(name, () => {
       withVersions(plugin, name, version => {
+        if (versionRange && !semver.intersects(version, versionRange)) return
+
         beforeEach(() => {
           tracer = require('../..')
         })
 
-        afterEach(() => {
-          return agent.close()
-        })
-
         describe('without configuration', () => {
           beforeEach(() => {
             return agent.load(plugin, name)
@@ -31,6 +30,10 @@ module.exports = (name, factory) => {
             Promise = factory ? factory(moduleExports) : moduleExports
           })
 
+          afterEach(() => {
+            return agent.close()
+          })
+
           it('should run the then() callbacks in the context where then() was called', () => {
             if (process.env.DD_CONTEXT_PROPAGATION === 'false') return
 
@@ -101,6 +104,53 @@ module.exports = (name, factory) => {
             })
           })
         })
+
+        describe('unpatching', () => {
+          beforeEach(() => {
+            const moduleExports = require(`../../../../versions/${name}@${version}`).get()
+
+            Promise = factory ? factory(moduleExports) : moduleExports
+          })
+
+          afterEach(() => {
+            return agent.close()
+          })
+
+          it('should behave the same before and after unpatching', async () => {
+            const span = {}
+
+            const testPatching = async function (Promise, tracer) {
+              const promise = new Promise((resolve, reject) => {
+                setImmediate(() => {
+                  tracer.scope().activate({}, () => {
+                    reject(new Error())
+                  })
+                })
+              })
+
+              await tracer.scope().activate(span, () => {
+                return promise
+                  .catch(err => {
+                    throw err
+                  })
+                  .catch(() => {
+                    return tracer.scope().active() !== span
+                  })
+              })
+            }
+
+            await agent.load()
+
+            const beforePatching = await testPatching(Promise, tracer)
+
+            tracer.use(name)
+            tracer._instrumenter.disable()
+
+            const afterUnpatching = await testPatching(Promise, tracer)
+
+            expect(beforePatching).to.equal(afterUnpatching)
+          })
+        })
       })
     })
   })