fix scope manager being instantiated multiple times

src/plugins/mysql.js

@@ -6,8 +6,9 @@ function createWrapQuery (tracer, config) {
   return function wrapQuery (query) {
     return function queryWithTrace (sql, values, cb) {
       const parentScope = tracer.scopeManager().active()
+      const parent = parentScope && parentScope.span()
       const span = tracer.startSpan('mysql.query', {
-        childOf: parentScope && parentScope.span(),
+        childOf: parent,
         tags: {
           [Tags.SPAN_KIND]: Tags.SPAN_KIND_RPC_CLIENT,
           'service.name': config.service || 'mysql',
@@ -28,7 +29,7 @@ function createWrapQuery (tracer, config) {
       span.setTag('resource.name', sequence.sql)
 
       if (sequence._callback) {
-        sequence._callback = wrapCallback(tracer, span, sequence._callback)
+        sequence._callback = wrapCallback(tracer, span, parent, sequence._callback)
       } else {
         sequence.on('end', () => {
           span.finish()
@@ -40,7 +41,7 @@ function createWrapQuery (tracer, config) {
   }
 }
 
-function wrapCallback (tracer, span, done) {
+function wrapCallback (tracer, span, parent, done) {
   return (err, res) => {
     if (err) {
       span.addTags({
@@ -52,6 +53,10 @@ function wrapCallback (tracer, span, done) {
 
     span.finish()
 
+    if (parent) {
+      tracer.scopeManager().activate(parent)
+    }
+
     done(err, res)
   }
 }

src/plugins/mysql2.js

@@ -6,8 +6,9 @@ function createWrapQuery (tracer, config) {
   return function wrapQuery (query) {
     return function queryWithTrace (sql, values, cb) {
       const parentScope = tracer.scopeManager().active()
+      const parent = parentScope && parentScope.span()
       const span = tracer.startSpan('mysql.query', {
-        childOf: parentScope && parentScope.span(),
+        childOf: parent,
         tags: {
           [Tags.SPAN_KIND]: Tags.SPAN_KIND_RPC_CLIENT,
           'service.name': config.service || 'mysql',
@@ -28,7 +29,7 @@ function createWrapQuery (tracer, config) {
       span.setTag('resource.name', sequence.sql)
 
       if (sequence.onResult) {
-        sequence.onResult = wrapCallback(span, sequence.onResult)
+        sequence.onResult = wrapCallback(tracer, span, parent, sequence.onResult)
       } else {
         sequence.on('end', () => {
           span.finish()
@@ -40,7 +41,7 @@ function createWrapQuery (tracer, config) {
   }
 }
 
-function wrapCallback (span, done) {
+function wrapCallback (tracer, span, parent, done) {
   return (err, res) => {
     if (err) {
       span.addTags({
@@ -52,6 +53,10 @@ function wrapCallback (span, done) {
 
     span.finish()
 
+    if (parent) {
+      tracer.scopeManager().activate(parent)
+    }
+
     done(err, res)
   }
 }

src/scope/scope_manager.js

@@ -5,7 +5,7 @@ const Scope = require('./scope')
 const Context = require('./context')
 const ContextExecution = require('./context_execution')
 
-const singleton = null
+let singleton = null
 
 /**
  * The Datadog Scope Manager. This is used for context propagation.
@@ -18,6 +18,8 @@ class ScopeManager {
       return singleton
     }
 
+    singleton = this
+
     const id = -1
     const execution = new ContextExecution()
 

test/plugins/agent.js

@@ -49,7 +49,6 @@ module.exports = {
 
         server.on('close', () => {
           tracer._instrumenter.unpatch()
-          tracer.scopeManager()._disable()
           tracer = null
         })
 

test/plugins/mysql.spec.js

@@ -48,8 +48,7 @@ describe('Plugin', () => {
 
         connection.query('SELECT 1 + 1 AS solution', () => {
           const active = tracer.scopeManager().active()
-          scope.close()
-          expect(active).to.equal(scope)
+          expect(active.span()).to.equal(scope.span())
           done()
         })
       })

test/plugins/mysql2.spec.js

@@ -48,8 +48,7 @@ describe('Plugin', () => {
 
         connection.query('SELECT 1 + 1 AS solution', () => {
           const active = tracer.scopeManager().active()
-          scope.close()
-          expect(active).to.equal(scope)
+          expect(active.span()).to.equal(scope.span())
           done()
         })
       })

test/scope/scope_manager.spec.js

@@ -31,8 +31,8 @@ describe('ScopeManager', () => {
     scopeManager = new ScopeManager()
   })
 
-  afterEach(() => {
-    scopeManager._disable()
+  it('should be a singleton', () => {
+    expect(new ScopeManager()).to.equal(scopeManager)
   })
 
   it('should enable its hooks', () => {

test/setup.js

@@ -14,6 +14,9 @@ const elasticsearch = require('elasticsearch')
 const amqplib = require('amqplib/callback_api')
 const platform = require('../src/platform')
 const node = require('../src/platform/node')
+const ScopeManager = require('../src/scope/scope_manager')
+
+const scopeManager = new ScopeManager()
 
 const retryOptions = {
   retries: 60,
@@ -33,6 +36,10 @@ global.wrapIt = wrapIt
 
 platform.use(node)
 
+after(() => {
+  scopeManager._disable()
+})
+
 waitForServices()
   .then(run)
   .catch(err => {
@@ -185,6 +192,44 @@ function waitForRabbitMQ () {
   })
 }
 
+function withoutScope (fn) {
+  return function () {
+    let active
+
+    while ((active = scopeManager.active())) {
+      active.close()
+    }
+
+    return fn.apply(this, arguments)
+  }
+}
+
 function wrapIt () {
-  // placeholder
+  const it = global.it
+
+  global.it = function (title, fn) {
+    if (!fn) {
+      return it.apply(this, arguments)
+    }
+
+    if (fn.length > 0) {
+      return it.call(this, title, function (done) {
+        arguments[0] = withoutScope(done)
+
+        return fn.apply(this, arguments)
+      })
+    } else {
+      return it.call(this, title, function () {
+        const result = fn.apply(this, arguments)
+
+        if (result && result.then) {
+          return result
+            .then(withoutScope(res => res))
+            .catch(withoutScope(err => Promise.reject(err)))
+        }
+
+        return result
+      })
+    }
+  }
 }

test/tracer.spec.js

@@ -29,10 +29,6 @@ describe('Tracer', () => {
     tracer = new Tracer(config)
   })
 
-  afterEach(() => {
-    tracer.scopeManager()._disable()
-  })
-
   describe('trace', () => {
     it('should run the callback with the new span', done => {
       tracer.trace('name', current => {