From 8e7f64134c6e2ef48926bf63943fa3338636a62c Mon Sep 17 00:00:00 2001
From: dxspasojevic <dxspasojevic@agiledigital.com.au>
Date: Sat, 18 Jun 2016 12:37:01 +1000
Subject: = kamon-core: add context name and token to MDC

---
 .../scala/kamon/trace/logging/MdcKeysSupport.scala | 16 ++++++--
 .../kamon/trace/logging/MdcKeysSupportSpec.scala   | 44 ++++++++++++++++++++++
 2 files changed, 57 insertions(+), 3 deletions(-)
 create mode 100644 kamon-core/src/test/scala/kamon/trace/logging/MdcKeysSupportSpec.scala

(limited to 'kamon-core/src')

diff --git a/kamon-core/src/main/scala/kamon/trace/logging/MdcKeysSupport.scala b/kamon-core/src/main/scala/kamon/trace/logging/MdcKeysSupport.scala
index abe56cd3..60a02ef8 100644
--- a/kamon-core/src/main/scala/kamon/trace/logging/MdcKeysSupport.scala
+++ b/kamon-core/src/main/scala/kamon/trace/logging/MdcKeysSupport.scala
@@ -17,13 +17,18 @@
 package kamon.trace.logging
 
 import kamon.trace.TraceLocal.AvailableToMdc
-import kamon.trace.{ Tracer, EmptyTraceContext, MetricsOnlyContext, TraceContext }
+import kamon.trace.{ EmptyTraceContext, MetricsOnlyContext, TraceContext, Tracer }
 import kamon.util.Supplier
-
 import org.slf4j.MDC
 
 trait MdcKeysSupport {
 
+  val traceTokenKey = "traceToken"
+
+  val traceNameKey = "traceName"
+
+  private val defaultKeys = Seq(traceTokenKey, traceNameKey)
+
   def withMdc[A](thunk: ⇒ A): A = {
     val keys = copyToMdc(Tracer.currentContext)
     try thunk finally keys.foreach(key ⇒ MDC.remove(key))
@@ -34,7 +39,12 @@ trait MdcKeysSupport {
 
   private[this] def copyToMdc(traceContext: TraceContext): Iterable[String] = traceContext match {
     case ctx: MetricsOnlyContext ⇒
-      ctx.traceLocalStorage.underlyingStorage.collect {
+
+      // Add the default key value pairs for the trace token and trace name.
+      MDC.put(traceTokenKey, ctx.token)
+      MDC.put(traceNameKey, ctx.name)
+
+      defaultKeys ++ ctx.traceLocalStorage.underlyingStorage.collect {
         case (available: AvailableToMdc, value) ⇒ Map(available.mdcKey -> String.valueOf(value))
       }.map { value ⇒ value.map { case (k, v) ⇒ MDC.put(k, v); k } }.flatten
 
diff --git a/kamon-core/src/test/scala/kamon/trace/logging/MdcKeysSupportSpec.scala b/kamon-core/src/test/scala/kamon/trace/logging/MdcKeysSupportSpec.scala
new file mode 100644
index 00000000..6193f028
--- /dev/null
+++ b/kamon-core/src/test/scala/kamon/trace/logging/MdcKeysSupportSpec.scala
@@ -0,0 +1,44 @@
+package kamon.trace.logging
+
+import kamon.testkit.BaseKamonSpec
+import kamon.trace.{ EmptyTraceContext, Tracer }
+import org.slf4j.MDC
+
+class MdcKeysSupportSpec extends BaseKamonSpec("mdc-keys-support-spec") {
+
+  "Running code with MDC support" should {
+    "add nothing to the MDC" when {
+      "the trace context is empty" in {
+        // Given an empty trace context.
+        Tracer.withContext(EmptyTraceContext) {
+          // When some code is executed with MDC support.
+          MdcKeysSupport.withMdc {
+            // Then the MDC should not contain the trace token.
+            Option(MDC.get(MdcKeysSupport.traceTokenKey)) should be(None)
+            // Or name
+            Option(MDC.get(MdcKeysSupport.traceNameKey)) should be(None)
+          }
+        }
+      }
+    }
+    "add the trace token and name to the context" when {
+      "the trace context is not empty" in {
+        // Given a trace context.
+        Tracer.withNewContext("name", Some("token")) {
+          // When some code is executed with MDC support.
+          MdcKeysSupport.withMdc {
+            // Then the MDC should contain the trace token.
+            Option(MDC.get(MdcKeysSupport.traceTokenKey)) should be(Some("token"))
+            // And name
+            Option(MDC.get(MdcKeysSupport.traceNameKey)) should be(Some("name"))
+          }
+
+          // Then after code is executed the MDC should have been cleared.
+          Option(MDC.get(MdcKeysSupport.traceTokenKey)) should be(None)
+          Option(MDC.get(MdcKeysSupport.traceNameKey)) should be(None)
+        }
+      }
+    }
+  }
+
+}
-- 
cgit v1.2.3