fix using log correlation with a custom logger

LICENSE-3rdparty.csv

@@ -8,6 +8,7 @@ require,lodash.pick,MIT,Copyright JS Foundation and other contributors
 require,lodash.truncate,MIT,Copyright JS Foundation and other contributors
 require,lodash.uniq,MIT,Copyright JS Foundation and other contributors
 require,methods,MIT,Copyright 2013-2014 TJ Holowaychuk 2013-2014 TJ Holowaychuk
+require,module-details-from-path,MIT,Copyright 2016 Thomas Watson Steen
 require,msgpack-lite,MIT,Copyright 2015 Yusuke Kawasaki
 require,opentracing,MIT,Copyright 2016 Resonance Labs Inc
 require,parent-module,MIT,Copyright Sindre Sorhus

package.json

@@ -46,6 +46,7 @@
     "lodash.truncate": "^4.4.2",
     "lodash.uniq": "^4.5.0",
     "methods": "^1.1.2",
+    "module-details-from-path": "^1.0.3",
     "msgpack-lite": "^0.1.26",
     "opentracing": "0.14.1",
     "parent-module": "^0.1.0",

src/config.js

@@ -10,6 +10,7 @@ class Config {
 
     const enabled = coalesce(options.enabled, platform.env('DD_TRACE_ENABLED'), true)
     const debug = coalesce(options.debug, platform.env('DD_TRACE_DEBUG'), false)
+    const logInjection = coalesce(options.logInjection, platform.env('DD_LOGS_INJECTION'), false)
     const env = coalesce(options.env, platform.env('DD_ENV'))
     const url = coalesce(options.url, platform.env('DD_TRACE_AGENT_URL'), null)
     const protocol = 'http'
@@ -26,6 +27,7 @@ class Config {
 
     this.enabled = String(enabled) === 'true'
     this.debug = String(debug) === 'true'
+    this.logInjection = String(logInjection) === 'true'
     this.env = env
     this.url = url ? new URL(url) : new URL(`${protocol}://${hostname}:${port}`)
     this.tags = Object.assign({}, options.tags)

src/instrumenter.js

@@ -2,6 +2,7 @@
 
 const semver = require('semver')
 const hook = require('require-in-the-middle')
+const parse = require('module-details-from-path')
 const path = require('path')
 const shimmer = require('shimmer')
 const uniq = require('lodash.uniq')
@@ -129,8 +130,7 @@ class Instrumenter {
             .filter(instrumentation => moduleName === filename(instrumentation))
             .filter(instrumentation => matchVersion(moduleVersion, instrumentation.versions))
             .forEach(instrumentation => {
-              this._instrumented.set(instrumentation, moduleExports)
-              instrumentation.patch.call(this, moduleExports, this._tracer._tracer, this._plugins.get(plugin).config)
+              this._patch(instrumentation, moduleExports, this._plugins.get(plugin).config)
             })
         } catch (e) {
           log.error(e)
@@ -147,7 +147,10 @@ class Instrumenter {
   }
 
   _set (plugin, meta) {
-    this._plugins.set(plugin, Object.assign({ config: {} }, meta))
+    meta = Object.assign({ config: {} }, meta)
+
+    this._plugins.set(plugin, meta)
+    this._load(plugin, meta)
   }
 
   _validate (plugin, moduleBaseDir, moduleVersion) {
@@ -177,13 +180,88 @@ class Instrumenter {
     this._plugins.delete(plugin)
   }
 
+  _patch (instrumentation, moduleExports, config) {
+    let instrumented = this._instrumented.get(instrumentation)
+
+    if (!instrumented) {
+      this._instrumented.set(instrumentation, instrumented = new Set())
+    }
+
+    if (!instrumented.has(moduleExports)) {
+      instrumented.add(moduleExports)
+      instrumentation.patch.call(this, moduleExports, this._tracer._tracer, config)
+    }
+  }
+
   _unpatch (instrumentation) {
-    try {
-      instrumentation.unpatch.call(this, this._instrumented.get(instrumentation))
-    } catch (e) {
-      log.error(e)
+    const instrumented = this._instrumented.get(instrumentation)
+
+    if (instrumented) {
+      instrumented.forEach(moduleExports => {
+        try {
+          instrumentation.unpatch.call(this, moduleExports)
+        } catch (e) {
+          log.error(e)
+        }
+      })
     }
   }
+
+  _load (plugin, meta) {
+    if (this._enabled) {
+      const instrumentations = [].concat(plugin)
+
+      try {
+        instrumentations
+          .forEach(instrumentation => {
+            getModules(instrumentation).forEach(nodule => {
+              this._patch(instrumentation, nodule, meta.config)
+            })
+          })
+      } catch (e) {
+        log.error(e)
+        this._fail(plugin)
+        log.debug(`Error while trying to patch ${meta.name}. The plugin has been disabled.`)
+      }
+    }
+  }
+}
+
+function getModules (instrumentation) {
+  const modules = []
+  const ids = Object.keys(require.cache)
+
+  let pkg
+
+  for (let i = 0, l = ids.length; i < l; i++) {
+    const id = ids[i].replace(pathSepExpr, '/')
+
+    if (!id.includes(`/node_modules/${instrumentation.name}/`)) continue
+
+    if (instrumentation.file) {
+      if (!id.endsWith(`/node_modules/${filename(instrumentation)}`)) continue
+
+      const basedir = getBasedir(ids[i])
+
+      pkg = require(`${basedir}/package.json`)
+    } else {
+      const basedir = getBasedir(ids[i])
+
+      pkg = require(`${basedir}/package.json`)
+
+      if (!id.endsWith(`/node_modules/${instrumentation.name}/${pkg.main}`)) continue
+    }
+
+    if (!matchVersion(pkg.version, instrumentation.versions)) continue
+
+    modules.push(require.cache[ids[i]].exports)
+  }
+
+  return modules
+}
+
+function getBasedir (id) {
+  return parse(id).basedir.replace(pathSepExpr, '/')
 }
 
 function matchVersion (version, ranges) {

src/opentracing/tracer.js

@@ -27,6 +27,7 @@ class DatadogTracer extends Tracer {
     this._url = config.url
     this._env = config.env
     this._tags = config.tags
+    this._logInjection = config.logInjection
     this._prioritySampler = new PrioritySampler(config.env)
     this._writer = new Writer(this._prioritySampler, config.url, config.bufferSize)
     this._recorder = new Recorder(this._writer, config.flushInterval)

src/plugins/bunyan.js

@@ -16,7 +16,7 @@ module.exports = {
   name: 'bunyan',
   versions: ['>=1'],
   patch (Logger, tracer, config) {
-    if (!config.correlate) return
+    if (!tracer._logInjection) return
     this.wrap(Logger.prototype, '_emit', createWrapEmit(tracer, config))
   },
   unpatch (Logger) {

src/plugins/pino.js

@@ -41,8 +41,7 @@ module.exports = [
     name: 'pino',
     versions: ['>=5'],
     patch (pino, tracer, config) {
-      if (!config.correlate) return
-
+      if (!tracer._logInjection) return
       this.wrap(Object.getPrototypeOf(pino()), pino.symbols.writeSym, createWrapWrite(tracer, config))
     },
     unpatch (pino) {
@@ -54,8 +53,7 @@ module.exports = [
     versions: ['4'],
     file: 'lib/tools.js',
     patch (tools, tracer, config) {
-      if (!config.correlate) return
-
+      if (!tracer._logInjection) return
       this.wrap(tools, 'genLog', createWrapGenLog(tracer, config))
     },
     unpatch (tools) {
@@ -66,8 +64,7 @@ module.exports = [
     name: 'pino',
     versions: ['2 - 3'],
     patch (pino, tracer, config) {
-      if (!config.correlate) return
-
+      if (!tracer._logInjection) return
       this.wrap(Object.getPrototypeOf(pino()), 'asJson', createWrapWrite(tracer, config))
     },
     unpatch (pino) {

src/plugins/winston.js

@@ -47,7 +47,7 @@ module.exports = [
     file: 'lib/winston/logger.js',
     versions: ['>=3'],
     patch (Logger, tracer, config) {
-      if (!config.correlate) return
+      if (!tracer._logInjection) return
       this.wrap(Logger.prototype, 'write', createWrapWrite(tracer, config))
     },
     unpatch (Logger) {
@@ -59,7 +59,7 @@ module.exports = [
     file: 'lib/winston/logger.js',
     versions: ['1 - 2'],
     patch (logger, tracer, config) {
-      if (!config.correlate) return
+      if (!tracer._logInjection) return
       this.wrap(logger.Logger.prototype, 'log', createWrapLog(tracer, config))
     },
     unpatch (logger) {

test/instrumenter.spec.js

@@ -1,6 +1,7 @@
 'use strict'
 
 const proxyquire = require('proxyquire').noCallThru()
+const path = require('path')
 
 describe('Instrumenter', () => {
   let Instrumenter
@@ -67,7 +68,13 @@ describe('Instrumenter', () => {
   })
 
   afterEach(() => {
-    delete require.cache[require.resolve('mysql-mock')]
+    const basedir = path.resolve(path.join(__dirname, 'node_modules'))
+
+    Object.keys(require.cache)
+      .filter(name => name.indexOf(basedir) !== -1)
+      .forEach(name => {
+        delete require.cache[name]
+      })
   })
 
   describe('with integrations enabled', () => {
@@ -142,6 +149,25 @@ describe('Instrumenter', () => {
         expect(integrations.mysql[0].patch).to.not.have.been.called
         expect(integrations.mysql[1].patch).to.not.have.been.called
       })
+
+      it('should attempt to patch already loaded modules', () => {
+        const express = require('express-mock')
+
+        instrumenter.use('express-mock')
+
+        expect(integrations.express.patch).to.have.been.called
+        expect(integrations.express.patch).to.have.been.calledWith(express, 'tracer', {})
+      })
+
+      it('should not patch twice already loaded modules', () => {
+        require('express-mock')
+
+        instrumenter.use('express-mock')
+
+        require('express-mock')
+
+        expect(integrations.express.patch).to.have.been.calledOnce
+      })
     })
 
     describe('patch', () => {

test/plugins/bunyan.spec.js

@@ -29,6 +29,7 @@ describe('Plugin', () => {
     withVersions(plugin, 'bunyan', version => {
       beforeEach(() => {
         tracer = require('../..')
+        return agent.load(plugin, 'bunyan')
       })
 
       afterEach(() => {
@@ -37,10 +38,7 @@ describe('Plugin', () => {
 
       describe('without configuration', () => {
         beforeEach(() => {
-          return agent.load(plugin, 'bunyan')
-            .then(() => {
-              setup(version)
-            })
+          setup(version)
         })
 
         it('should not alter the default behavior', () => {
@@ -61,10 +59,8 @@ describe('Plugin', () => {
 
       describe('with configuration', () => {
         beforeEach(() => {
-          return agent.load(plugin, 'bunyan', { correlate: true })
-            .then(() => {
-              setup(version)
-            })
+          tracer._tracer._logInjection = true
+          setup(version)
         })
 
         it('should add the trace identifiers to logger instances', () => {

test/plugins/pino.spec.js

@@ -29,6 +29,7 @@ describe('Plugin', () => {
     withVersions(plugin, 'pino', version => {
       beforeEach(() => {
         tracer = require('../..')
+        return agent.load(plugin, 'pino')
       })
 
       afterEach(() => {
@@ -37,10 +38,7 @@ describe('Plugin', () => {
 
       describe('without configuration', () => {
         beforeEach(() => {
-          return agent.load(plugin, 'pino')
-            .then(() => {
-              setup(version)
-            })
+          setup(version)
         })
 
         it('should not alter the default behavior', () => {
@@ -61,10 +59,8 @@ describe('Plugin', () => {
 
       describe('with configuration', () => {
         beforeEach(() => {
-          return agent.load(plugin, 'pino', { correlate: true })
-            .then(() => {
-              setup(version)
-            })
+          tracer._tracer._logInjection = true
+          setup(version)
         })
 
         it('should add the trace identifiers to logger instances', () => {

test/plugins/winston.spec.js

@@ -37,6 +37,7 @@ describe('Plugin', () => {
     withVersions(plugin, 'winston', version => {
       beforeEach(() => {
         tracer = require('../..')
+        return agent.load(plugin, 'winston')
       })
 
       afterEach(() => {
@@ -45,10 +46,7 @@ describe('Plugin', () => {
 
       describe('without configuration', () => {
         beforeEach(() => {
-          return agent.load(plugin, 'winston')
-            .then(() => {
-              setup(version)
-            })
+          setup(version)
         })
 
         it('should not alter the default behavior', () => {
@@ -71,10 +69,8 @@ describe('Plugin', () => {
 
       describe('with configuration', () => {
         beforeEach(() => {
-          return agent.load(plugin, 'winston', { correlate: true })
-            .then(() => {
-              setup(version)
-            })
+          tracer._tracer._logInjection = true
+          setup(version)
         })
 
         it('should add the trace identifiers to the default logger', () => {