From 9b709b28d70995d9a7762c43af398d31a0417e6a Mon Sep 17 00:00:00 2001
From: rochdev <roch.devost@datadoghq.com>
Date: Wed, 9 Jan 2019 19:01:58 -0500
Subject: [PATCH] add log correlation of trace identifiers for bunyan

---
 src/plugins/bunyan.js       | 25 +++++++++++
 src/plugins/index.js        |  1 +
 test/plugins/bunyan.spec.js | 87 +++++++++++++++++++++++++++++++++++++
 3 files changed, 113 insertions(+)
 create mode 100644 src/plugins/bunyan.js
 create mode 100644 test/plugins/bunyan.spec.js

diff --git a/src/plugins/bunyan.js b/src/plugins/bunyan.js
new file mode 100644
index 00000000000..350fe6aa642
--- /dev/null
+++ b/src/plugins/bunyan.js
@@ -0,0 +1,25 @@
+'use strict'
+
+const tx = require('./util/log')
+
+function createWrapEmit (tracer, config) {
+  return function wrapEmit (emit) {
+    return function emitWithTrace (rec, noemit) {
+      tx.correlate(tracer, rec)
+
+      return emit.apply(this, arguments)
+    }
+  }
+}
+
+module.exports = {
+  name: 'bunyan',
+  versions: ['1 - 2'],
+  patch (Logger, tracer, config) {
+    if (!config.correlate) return
+    this.wrap(Logger.prototype, '_emit', createWrapEmit(tracer, config))
+  },
+  unpatch (Logger) {
+    this.unwrap(Logger.prototype, '_emit')
+  }
+}
diff --git a/src/plugins/index.js b/src/plugins/index.js
index 83436776b39..4f6e1154bfd 100644
--- a/src/plugins/index.js
+++ b/src/plugins/index.js
@@ -4,6 +4,7 @@ module.exports = {
   'amqp10': require('./amqp10'),
   'amqplib': require('./amqplib'),
   'bluebird': require('./bluebird'),
+  'bunyan': require('./bunyan'),
   'dns': require('./dns'),
   'elasticsearch': require('./elasticsearch'),
   'express': require('./express'),
diff --git a/test/plugins/bunyan.spec.js b/test/plugins/bunyan.spec.js
new file mode 100644
index 00000000000..5adcbeeaa03
--- /dev/null
+++ b/test/plugins/bunyan.spec.js
@@ -0,0 +1,87 @@
+'use strict'
+
+const Writable = require('stream').Writable
+const agent = require('./agent')
+const plugin = require('../../src/plugins/bunyan')
+
+wrapIt()
+
+describe('Plugin', () => {
+  let logger
+  let tracer
+  let stream
+  let span
+
+  function setup (version) {
+    const bunyan = require(`../../versions/bunyan@${version}`).get()
+
+    span = tracer.startSpan('test')
+
+    stream = new Writable()
+    stream._write = () => {}
+
+    sinon.spy(stream, 'write')
+
+    logger = bunyan.createLogger({ name: 'test', stream })
+  }
+
+  describe('bunyan', () => {
+    withVersions(plugin, 'bunyan', version => {
+      beforeEach(() => {
+        tracer = require('../..')
+      })
+
+      afterEach(() => {
+        return agent.close()
+      })
+
+      describe('without configuration', () => {
+        beforeEach(() => {
+          return agent.load(plugin, 'bunyan')
+            .then(() => {
+              setup(version)
+            })
+        })
+
+        it('should not alter the default behavior', () => {
+          tracer.scopeManager().activate(span)
+
+          logger.info('message')
+
+          expect(stream.write).to.have.been.called
+
+          const record = JSON.parse(stream.write.firstCall.args[0].toString())
+
+          expect(record).to.not.include({
+            'dd.trace_id': span.context().toTraceId(),
+            'dd.span_id': span.context().toSpanId()
+          })
+        })
+      })
+
+      describe('with configuration', () => {
+        beforeEach(() => {
+          return agent.load(plugin, 'bunyan', { correlate: true })
+            .then(() => {
+              setup(version)
+            })
+        })
+
+        it('should add the trace identifiers to logger instances', () => {
+          tracer.scopeManager().activate(span)
+
+          logger.info('message')
+
+          expect(stream.write).to.have.been.called
+
+          const record = JSON.parse(stream.write.firstCall.args[0].toString())
+
+          expect(record).to.include({
+            'dd.trace_id': span.context().toTraceId(),
+            'dd.span_id': span.context().toSpanId()
+          })
+        })
+      })
+    })
+  })
+})