Add tags for traces and closes #327

7 files changed, 58 insertions, 27 deletions

diff --git a/kamon-akka-remote/src/main/scala/kamon/akka/instrumentation/RemotingInstrumentation.scala b/kamon-akka-remote/src/main/scala/kamon/akka/instrumentation/RemotingInstrumentation.scala
index a0940c10..ee88dc64 100644
--- a/kamon-akka-remote/src/main/scala/kamon/akka/instrumentation/RemotingInstrumentation.scala
+++ b/kamon-akka-remote/src/main/scala/kamon/akka/instrumentation/RemotingInstrumentation.scala
@@ -89,6 +89,7 @@ class RemotingInstrumentation {
       val ctx = tracer.newContext(
         remoteTraceContext.getTraceName,
         Option(remoteTraceContext.getTraceToken),
+        tags = Map.empty,
         new MilliTimestamp(remoteTraceContext.getStartMilliTime).toRelativeNanoTimestamp,
         remoteTraceContext.getIsOpen,
         isLocal = false)
diff --git a/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala b/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala
index c044719c..c45a49d7 100644
--- a/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala
+++ b/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala
@@ -25,7 +25,7 @@ import kamon.util.{ NanoInterval, RelativeNanoTimestamp }
 
 import scala.annotation.tailrec
 
-private[kamon] class MetricsOnlyContext(traceName: String, val token: String, izOpen: Boolean, val levelOfDetail: LevelOfDetail,
+private[kamon] class MetricsOnlyContext(traceName: String, val token: String, tags: Map[String, String], izOpen: Boolean, val levelOfDetail: LevelOfDetail,
   val startTimestamp: RelativeNanoTimestamp, log: LoggingAdapter)
     extends TraceContext {
 
@@ -53,7 +53,7 @@ private[kamon] class MetricsOnlyContext(traceName: String, val token: String, iz
     _elapsedTime = traceElapsedTime
 
     if (Kamon.metrics.shouldTrack(name, TraceMetrics.category))
-      Kamon.metrics.entity(TraceMetrics, name).elapsedTime.record(traceElapsedTime.nanos)
+      Kamon.metrics.entity(TraceMetrics, name, tags).elapsedTime.record(traceElapsedTime.nanos)
     drainFinishedSegments()
   }
 
diff --git a/kamon-core/src/main/scala/kamon/trace/TracerModule.scala b/kamon-core/src/main/scala/kamon/trace/TracerModule.scala
index 06286cae..3ff32751 100644
--- a/kamon-core/src/main/scala/kamon/trace/TracerModule.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TracerModule.scala
@@ -1,6 +1,6 @@
 /*
  * =========================================================================================
- * Copyright © 2013 the kamon project <http://kamon.io/>
+ * Copyright © 2013-2016 the kamon project <http://kamon.io/>
  *
  * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
  * except in compliance with the License. You may obtain a copy of the License at
@@ -16,17 +16,21 @@
 
 package kamon.trace
 
+import java.util
+
 import akka.actor._
-import akka.event.{ LoggingAdapter, Logging }
+import akka.event.{ Logging, LoggingAdapter }
 import com.typesafe.config.Config
 import kamon.Kamon
 import kamon.metric.MetricsModule
 import kamon.util._
+import scala.collection.JavaConverters._
 
 trait TracerModule {
   def newContext(name: String): TraceContext
   def newContext(name: String, token: Option[String]): TraceContext
-  def newContext(name: String, token: Option[String], timestamp: RelativeNanoTimestamp, isOpen: Boolean, isLocal: Boolean): TraceContext
+  def newContext(name: String, token: Option[String], tags: Map[String, String]): TraceContext
+  def newContext(name: String, token: Option[String], tags: Map[String, String], timestamp: RelativeNanoTimestamp, isOpen: Boolean, isLocal: Boolean): TraceContext
 
   def subscribe(subscriber: ActorRef): Unit
   def unsubscribe(subscriber: ActorRef): Unit
@@ -57,8 +61,8 @@ object Tracer {
   def withContext[T](context: TraceContext, code: Supplier[T]): T =
     withContext(context)(code.get)
 
-  def withNewContext[T](traceName: String, traceToken: Option[String], autoFinish: Boolean)(code: ⇒ T): T = {
-    withContext(Kamon.tracer.newContext(traceName, traceToken)) {
+  def withNewContext[T](traceName: String, traceToken: Option[String], tags: Map[String, String], autoFinish: Boolean)(code: ⇒ T): T = {
+    withContext(Kamon.tracer.newContext(traceName, traceToken, tags)) {
       val codeResult = code
       if (autoFinish)
         currentContext.finish()
@@ -68,26 +72,47 @@ object Tracer {
   }
 
   def withNewContext[T](traceName: String)(code: ⇒ T): T =
-    withNewContext(traceName, None, false)(code)
+    withNewContext(traceName, None)(code)
+
+  def withNewContext[T](traceName: String, tags: Map[String, String])(code: ⇒ T): T =
+    withNewContext(traceName, None, tags)(code)
 
   def withNewContext[T](traceName: String, traceToken: Option[String])(code: ⇒ T): T =
-    withNewContext(traceName, traceToken, false)(code)
+    withNewContext(traceName, traceToken, Map.empty[String, String])(code)
+
+  def withNewContext[T](traceName: String, traceToken: Option[String], tags: Map[String, String])(code: ⇒ T): T =
+    withNewContext(traceName, traceToken, tags, autoFinish = false)(code)
 
   def withNewContext[T](traceName: String, autoFinish: Boolean)(code: ⇒ T): T =
-    withNewContext(traceName, None, autoFinish)(code)
+    withNewContext(traceName, None, Map.empty[String, String], autoFinish)(code)
+
+  def withNewContext[T](traceName: String, tags: Map[String, String], autoFinish: Boolean)(code: ⇒ T): T =
+    withNewContext(traceName, None, tags, autoFinish)(code)
 
   // Java variants.
   def withNewContext[T](traceName: String, traceToken: Option[String], autoFinish: Boolean, code: Supplier[T]): T =
-    withNewContext(traceName, traceToken, autoFinish)(code.get)
+    withNewContext(traceName, traceToken, Map.empty[String, String], autoFinish)(code.get)
+
+  def withNewContext[T](traceName: String, traceToken: Option[String], tags: util.Map[String, String], autoFinish: Boolean, code: Supplier[T]): T =
+    withNewContext(traceName, traceToken, tags.asScala.toMap, autoFinish)(code.get)
 
   def withNewContext[T](traceName: String, code: Supplier[T]): T =
-    withNewContext(traceName, None, false)(code.get)
+    withNewContext(traceName, None, Map.empty[String, String], autoFinish = false)(code.get)
+
+  def withNewContext[T](traceName: String, tags: util.Map[String, String], code: Supplier[T]): T =
+    withNewContext(traceName, None, tags.asScala.toMap, autoFinish = false)(code.get)
 
   def withNewContext[T](traceName: String, traceToken: Option[String], code: Supplier[T]): T =
-    withNewContext(traceName, traceToken, false)(code.get)
+    withNewContext(traceName, traceToken, Map.empty[String, String], autoFinish = false)(code.get)
+
+  def withNewContext[T](traceName: String, traceToken: Option[String], tags: util.Map[String, String], code: Supplier[T]): T =
+    withNewContext(traceName, traceToken, tags.asScala.toMap, autoFinish = false)(code.get)
 
   def withNewContext[T](traceName: String, autoFinish: Boolean, code: Supplier[T]): T =
-    withNewContext(traceName, None, autoFinish)(code.get)
+    withNewContext(traceName, None, Map.empty[String, String], autoFinish)(code.get)
+
+  def withNewContext[T](traceName: String, tags: util.Map[String, String], autoFinish: Boolean, code: Supplier[T]): T =
+    withNewContext(traceName, None, tags.asScala.toMap, autoFinish)(code.get)
 }
 
 private[kamon] class TracerModuleImpl(metricsExtension: MetricsModule, config: Config) extends TracerModule {
@@ -96,7 +121,7 @@ private[kamon] class TracerModuleImpl(metricsExtension: MetricsModule, config: C
   private val _subscriptions = new LazyActorRef
   private val _incubator = new LazyActorRef
   private val _dynamic = new akka.actor.ReflectiveDynamicAccess(getClass.getClassLoader)
-  private val _tokenGenerator = _dynamic.createInstanceFor[Function0[String]](_settings.tokenGeneratorFQN, Nil).get // let's bubble up any problems.
+  private val _tokenGenerator = _dynamic.createInstanceFor[() ⇒ String](_settings.tokenGeneratorFQN, Nil).get // let's bubble up any problems.
 
   private def newToken: String = _tokenGenerator()
 
@@ -106,14 +131,17 @@ private[kamon] class TracerModuleImpl(metricsExtension: MetricsModule, config: C
   def newContext(name: String, token: Option[String]): TraceContext =
     createTraceContext(name, token)
 
-  def newContext(name: String, token: Option[String], timestamp: RelativeNanoTimestamp, isOpen: Boolean, isLocal: Boolean): TraceContext =
-    createTraceContext(name, token, timestamp, isOpen, isLocal)
+  def newContext(name: String, token: Option[String], tags: Map[String, String]): TraceContext =
+    createTraceContext(name, token, tags)
+
+  def newContext(name: String, token: Option[String], tags: Map[String, String], timestamp: RelativeNanoTimestamp, isOpen: Boolean, isLocal: Boolean): TraceContext =
+    createTraceContext(name, token, tags, timestamp, isOpen, isLocal)
 
-  private def createTraceContext(traceName: String, token: Option[String], startTimestamp: RelativeNanoTimestamp = RelativeNanoTimestamp.now,
+  private def createTraceContext(traceName: String, token: Option[String], tags: Map[String, String] = Map.empty, startTimestamp: RelativeNanoTimestamp = RelativeNanoTimestamp.now,
     isOpen: Boolean = true, isLocal: Boolean = true): TraceContext = {
 
     def newMetricsOnlyContext(token: String): TraceContext =
-      new MetricsOnlyContext(traceName, token, isOpen, _settings.levelOfDetail, startTimestamp, _logger)
+      new MetricsOnlyContext(traceName, token, tags, isOpen, _settings.levelOfDetail, startTimestamp, _logger)
 
     val traceToken = token.getOrElse(newToken)
 
@@ -123,7 +151,7 @@ private[kamon] class TracerModuleImpl(metricsExtension: MetricsModule, config: C
       case _ if !isLocal || !_settings.sampler.shouldTrace ⇒
         newMetricsOnlyContext(traceToken)
       case _ ⇒
-        new TracingContext(traceName, traceToken, true, _settings.levelOfDetail, isLocal, startTimestamp, _logger, dispatchTracingContext)
+        new TracingContext(traceName, traceToken, tags, izOpen = true, _settings.levelOfDetail, isLocal, startTimestamp, _logger, dispatchTracingContext)
     }
   }
 
diff --git a/kamon-core/src/main/scala/kamon/trace/TracingContext.scala b/kamon-core/src/main/scala/kamon/trace/TracingContext.scala
index 9269a99e..1c735fbe 100644
--- a/kamon-core/src/main/scala/kamon/trace/TracingContext.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TracingContext.scala
@@ -25,9 +25,9 @@ import kamon.metric.MetricsModule
 
 import scala.collection.concurrent.TrieMap
 
-private[trace] class TracingContext(traceName: String, token: String, izOpen: Boolean, levelOfDetail: LevelOfDetail,
+private[trace] class TracingContext(traceName: String, token: String, tags: Map[String, String], izOpen: Boolean, levelOfDetail: LevelOfDetail,
   isLocal: Boolean, startTimeztamp: RelativeNanoTimestamp, log: LoggingAdapter, traceInfoSink: TracingContext ⇒ Unit)
-    extends MetricsOnlyContext(traceName, token, izOpen, levelOfDetail, startTimeztamp, log) {
+    extends MetricsOnlyContext(traceName, token, tags, izOpen, levelOfDetail, startTimeztamp, log) {
 
   private val _openSegments = new AtomicInteger(0)
   private val _startTimestamp = NanoTimestamp.now
diff --git a/kamon-core/src/test/scala/kamon/testkit/BaseKamonSpec.scala b/kamon-core/src/test/scala/kamon/testkit/BaseKamonSpec.scala
index f23d974c..497cb77d 100644
--- a/kamon-core/src/test/scala/kamon/testkit/BaseKamonSpec.scala
+++ b/kamon-core/src/test/scala/kamon/testkit/BaseKamonSpec.scala
@@ -20,8 +20,8 @@ import akka.testkit.{ ImplicitSender, TestKitBase }
 import akka.actor.ActorSystem
 import com.typesafe.config.{ Config, ConfigFactory }
 import kamon.Kamon
-import kamon.metric.{ Entity, SubscriptionsDispatcher, EntitySnapshot }
-import kamon.trace.TraceContext
+import kamon.metric.{ Entity, EntitySnapshot, SubscriptionsDispatcher }
+import kamon.trace.{ TraceContext, Tracer }
 import kamon.util.LazyActorRef
 import org.scalatest.{ BeforeAndAfterAll, Matchers, WordSpecLike }
 
@@ -41,6 +41,9 @@ abstract class BaseKamonSpec(actorSystemName: String) extends TestKitBase with W
   def newContext(name: String, token: String): TraceContext =
     Kamon.tracer.newContext(name, Option(token))
 
+  def newContext(name: String, token: String, tags: Map[String, String]): TraceContext =
+    Kamon.tracer.newContext(name, Option(token), tags)
+
   def takeSnapshotOf(name: String, category: String): EntitySnapshot = {
     val recorder = Kamon.metrics.find(name, category).get
     recorder.collect(collectionContext)
diff --git a/kamon-core/src/test/scala/kamon/trace/SimpleTraceSpec.scala b/kamon-core/src/test/scala/kamon/trace/SimpleTraceSpec.scala
index d6886d4c..95deeef9 100644
--- a/kamon-core/src/test/scala/kamon/trace/SimpleTraceSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/SimpleTraceSpec.scala
@@ -1,6 +1,6 @@
 /*
  * =========================================================================================
- * Copyright © 2013-2014 the kamon project <http://kamon.io/>
+ * Copyright © 2013-2016 the kamon project <http://kamon.io/>
  *
  * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
  * except in compliance with the License. You may obtain a copy of the License at
@@ -16,9 +16,9 @@
 
 package kamon.trace
 
-import com.typesafe.config.ConfigFactory
 import kamon.Kamon
 import kamon.testkit.BaseKamonSpec
+
 import scala.concurrent.duration._
 
 class SimpleTraceSpec extends BaseKamonSpec("simple-trace-spec") {
diff --git a/kamon-core/src/test/scala/kamon/trace/TraceContextManipulationSpec.scala b/kamon-core/src/test/scala/kamon/trace/TraceContextManipulationSpec.scala
index ecef9271..6a454149 100644
--- a/kamon-core/src/test/scala/kamon/trace/TraceContextManipulationSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/TraceContextManipulationSpec.scala
@@ -16,7 +16,6 @@
 
 package kamon.trace
 
-import com.typesafe.config.ConfigFactory
 import kamon.testkit.BaseKamonSpec
 
 class TraceContextManipulationSpec extends BaseKamonSpec("trace-metrics-spec") {