Merge pull request #7 from matwojcik/customised-mdc-keys

Customised mdc keys
kamon-io · Mar 27, 2018 · 1d5612b · 1d5612b
2 parents 909296a + be5494a
commit 1d5612b
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 12 deletions.
diff --git a/src/main/resources/reference.conf b/src/main/resources/reference.conf
@@ -3,5 +3,7 @@ kamon {
     // if enabled kamon will set MDC context with TraceID and SpanID valuse
     // they could be referenced using kamonTraceID and kamonSpanID keys
     mdc-context-propagation = on
+    mdc-trace-id-key = kamonTraceId
+    mdc-span-id-key = kamonSpanID
   }
 }
diff --git a/src/main/scala/kamon/logback/instrumentation/AsyncAppenderInstrumentation.scala b/src/main/scala/kamon/logback/instrumentation/AsyncAppenderInstrumentation.scala
@@ -29,10 +29,13 @@ import scala.beans.BeanProperty
 
 object AsyncAppenderInstrumentation {
 
-  val MdcTraceKey : String = "kamonTraceID"
-  val MdcSpanKey : String = "kamonSpanID"
+  @volatile private var _mdcContextPropagation: Boolean = true
+  @volatile private var _mdcTraceKey: String = "kamonTraceID"
+  @volatile private var _mdcSpanKey: String = "kamonSpanID"
 
-  @volatile var mdcContextPropagation : Boolean = true
+  def mdcTraceKey: String = _mdcTraceKey
+  def mdcSpanKey: String = _mdcSpanKey
+  def mdcContextPropagation: Boolean = _mdcContextPropagation
 
   loadConfiguration(Kamon.config())
 
@@ -44,7 +47,9 @@ object AsyncAppenderInstrumentation {
 
   private def loadConfiguration(config: Config): Unit = synchronized {
     val logbackConfig = config.getConfig("kamon.logback")
-    mdcContextPropagation = logbackConfig.getBoolean("mdc-context-propagation")
+    _mdcContextPropagation = logbackConfig.getBoolean("mdc-context-propagation")
+    _mdcTraceKey = logbackConfig.getString("mdc-trace-id-key")
+    _mdcSpanKey = logbackConfig.getString("mdc-span-id-key")
   }
 }
 @Aspect
@@ -68,13 +73,13 @@ class AsyncAppenderInstrumentation {
     val context = Kamon.currentContext().get(Span.ContextKey)
 
     if (context.context().traceID != IdentityProvider.NoIdentifier && AsyncAppenderInstrumentation.mdcContextPropagation){
-      MDC.put(AsyncAppenderInstrumentation.MdcTraceKey, context.context().traceID.string)
-      MDC.put(AsyncAppenderInstrumentation.MdcSpanKey, context.context().spanID.string)
+      MDC.put(AsyncAppenderInstrumentation.mdcTraceKey, context.context().traceID.string)
+      MDC.put(AsyncAppenderInstrumentation.mdcSpanKey, context.context().spanID.string)
       try {
         pjp.proceed()
       } finally {
-        MDC.remove(AsyncAppenderInstrumentation.MdcTraceKey)
-        MDC.remove(AsyncAppenderInstrumentation.MdcSpanKey)
+        MDC.remove(AsyncAppenderInstrumentation.mdcTraceKey)
+        MDC.remove(AsyncAppenderInstrumentation.mdcSpanKey)
       }
     } else {
       pjp.proceed()

diff --git a/src/test/scala/kamon/logback/LogbackSpanConverterSpec.scala b/src/test/scala/kamon/logback/LogbackSpanConverterSpec.scala
@@ -53,7 +53,7 @@ class LogbackSpanConverterSpec extends WordSpec with Matchers with Eventually {
       }
 
       "MDC context" in {
-        val memoryAppender = buildMemoryAppender(configurator,s"%X{${AsyncAppenderInstrumentation.MdcTraceKey}} %X{${AsyncAppenderInstrumentation.MdcSpanKey}} %X{mdc_key}")
+        val memoryAppender = buildMemoryAppender(configurator,s"%X{${AsyncAppenderInstrumentation.mdcTraceKey}} %X{${AsyncAppenderInstrumentation.mdcSpanKey}} %X{mdc_key}")
 
         val span = Kamon.buildSpan("my-span").start()
         val traceID = span.context().traceID
@@ -66,8 +66,8 @@ class LogbackSpanConverterSpec extends WordSpec with Matchers with Eventually {
         }
 
         memoryAppender.getLastLine shouldBe traceID.string + " " + spanID.string + " mdc_value"
-        MDC.get(AsyncAppenderInstrumentation.MdcTraceKey) shouldBe null
-        MDC.get(AsyncAppenderInstrumentation.MdcSpanKey) shouldBe null
+        MDC.get(AsyncAppenderInstrumentation.mdcTraceKey) shouldBe null
+        MDC.get(AsyncAppenderInstrumentation.mdcSpanKey) shouldBe null
       }
 
       "disable MDC context" in {
@@ -78,7 +78,7 @@ class LogbackSpanConverterSpec extends WordSpec with Matchers with Eventually {
         )
 
 
-        val memoryAppender = buildMemoryAppender(configurator,s"%X{${AsyncAppenderInstrumentation.MdcTraceKey}}")
+        val memoryAppender = buildMemoryAppender(configurator,s"%X{${AsyncAppenderInstrumentation.mdcTraceKey}}")
 
         val span = Kamon.buildSpan("my-span").start()
         val contextWithSpan = Context.create(Span.ContextKey, span)