+ core: allow accessing tags on TraceContext and Segment (#386)

3 files changed, 33 insertions, 0 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala b/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala
index c3a1571a..c1e23d91 100644
--- a/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala
+++ b/kamon-core/src/main/scala/kamon/trace/MetricsOnlyContext.scala
@@ -49,6 +49,7 @@ private[kamon] class MetricsOnlyContext(traceName: String,
       log.warning("Can't rename trace from [{}] to [{}] because the trace is already closed.", name, newName)
 
   def name: String = _name
+  def tags: Map[String, String] = _tags.toMap
   def isEmpty: Boolean = false
   def status: Status = _status
   def addMetadata(key: String, value: String): Unit = {}
@@ -132,6 +133,7 @@ private[kamon] class MetricsOnlyContext(traceName: String,
     @volatile private var _status: Status = Status.Open
 
     def name: String = _segmentName
+    def tags: Map[String, String] = _tags.toMap
     def isEmpty: Boolean = false
     def status: Status = _status
     def addMetadata(key: String, value: String): Unit = {}
diff --git a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
index 935fe0ba..c5704caa 100644
--- a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala
@@ -28,6 +28,7 @@ import scala.concurrent.Future
 trait TraceContext {
   def name: String
   def token: String
+  def tags: Map[String, String]
   def isEmpty: Boolean
   def nonEmpty: Boolean = !isEmpty
   def isClosed: Boolean = !(Status.Open == status)
@@ -86,6 +87,7 @@ trait Segment {
   def name: String
   def category: String
   def library: String
+  def tags: Map[String, String]
   def isEmpty: Boolean
   def nonEmpty: Boolean = !isEmpty
   def isClosed: Boolean = !(Status.Open == status)
@@ -101,6 +103,7 @@ trait Segment {
 case object EmptyTraceContext extends TraceContext {
   def name: String = "empty-trace"
   def token: String = ""
+  def tags: Map[String, String] = Map.empty
   def isEmpty: Boolean = true
   def status: Status = Closed
   def finish(): Unit = {}
@@ -117,6 +120,7 @@ case object EmptyTraceContext extends TraceContext {
     val name: String = "empty-segment"
     val category: String = "empty-category"
     val library: String = "empty-library"
+    def tags: Map[String, String] = Map.empty
     def isEmpty: Boolean = true
     def status: Status = Closed
     def finish(): Unit = {}
diff --git a/kamon-core/src/test/scala/kamon/trace/TraceContextManipulationSpec.scala b/kamon-core/src/test/scala/kamon/trace/TraceContextManipulationSpec.scala
index 2756eb30..58b51a73 100644
--- a/kamon-core/src/test/scala/kamon/trace/TraceContextManipulationSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/TraceContextManipulationSpec.scala
@@ -63,6 +63,20 @@ class TraceContextManipulationSpec extends BaseKamonSpec("trace-metrics-spec") {
       createdContext.name shouldBe "renamed-trace"
     }
 
+    "allow tagging and untagging a trace" in {
+      val createdContext = Tracer.withContext(newContext("trace-before-rename")) {
+        Tracer.currentContext.addTag("trace-tag", "tag-1")
+        Tracer.currentContext
+      }
+
+      Tracer.currentContext shouldBe empty
+      createdContext.tags shouldBe Map("trace-tag" -> "tag-1")
+
+      createdContext.removeTag("trace-tag", "tag-1")
+
+      createdContext.tags shouldBe Map.empty
+    }
+
     "allow creating a segment within a trace" in {
       val createdContext = Tracer.withContext(newContext("trace-with-segments")) {
         Tracer.currentContext.startSegment("segment-1", "segment-1-category", "segment-library")
@@ -82,5 +96,18 @@ class TraceContextManipulationSpec extends BaseKamonSpec("trace-metrics-spec") {
         segment.name should be("new-segment-name")
       }
     }
+
+    "allow tagging and untagging a segment" in {
+      Tracer.withContext(newContext("trace-with-renamed-segment")) {
+        val segment = Tracer.currentContext.startSegment("segment-name", "segment-label", "segment-library")
+        segment.tags should be(Map.empty)
+
+        segment.addTag("segment-tag", "tag-1")
+        segment.tags should be(Map("segment-tag" -> "tag-1"))
+
+        segment.removeTag("segment-tag", "tag-1")
+        segment.tags should be(Map.empty)
+      }
+    }
   }
 }
\ No newline at end of file