tests for Span building and the ExtendedB3 codec, related bugfixes

4 files changed, 48 insertions, 33 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/IdentityProvider.scala b/kamon-core/src/main/scala/kamon/trace/IdentityProvider.scala
index 0e7d3e7d..3f44629e 100644
--- a/kamon-core/src/main/scala/kamon/trace/IdentityProvider.scala
+++ b/kamon-core/src/main/scala/kamon/trace/IdentityProvider.scala
@@ -13,7 +13,14 @@ trait IdentityProvider {
 }
 
 object IdentityProvider {
-  case class Identifier(string: String, bytes: Array[Byte])
+  case class Identifier(string: String, bytes: Array[Byte]) {
+
+    override def equals(obj: Any): Boolean = {
+      if(obj != null && obj.isInstanceOf[Identifier])
+        obj.asInstanceOf[Identifier].string == string
+      else false
+    }
+  }
 
   val NoIdentifier = Identifier("", new Array[Byte](0))
 
diff --git a/kamon-core/src/main/scala/kamon/trace/Span.scala b/kamon-core/src/main/scala/kamon/trace/Span.scala
index ce09a36e..01cfbfc3 100644
--- a/kamon-core/src/main/scala/kamon/trace/Span.scala
+++ b/kamon-core/src/main/scala/kamon/trace/Span.scala
@@ -97,20 +97,20 @@ object Span {
     *
     * @param spanContext
     * @param initialOperationName
-    * @param initialTags
+    * @param initialSpanTags
     * @param startTimestampMicros
     * @param reporterRegistry
     */
-  final class Real(spanContext: SpanContext, initialOperationName: String, initialTags: Map[String, Span.TagValue],
-    startTimestampMicros: Long, reporterRegistry: ReporterRegistryImpl, tracer: Tracer) extends Span {
+  final class Real(spanContext: SpanContext, initialOperationName: String, initialSpanTags: Map[String, Span.TagValue],
+    initialMetricTags: Map[String, String], startTimestampMicros: Long, reporterRegistry: ReporterRegistryImpl, tracer: Tracer) extends Span {
 
     private var collectMetrics: Boolean = true
     private var open: Boolean = true
     private val sampled: Boolean = spanContext.samplingDecision == SamplingDecision.Sample
     private var operationName: String = initialOperationName
 
-    private var spanTags: Map[String, Span.TagValue] = initialTags
-    private var customMetricTags = Map.empty[String, String]
+    private var spanTags: Map[String, Span.TagValue] = initialSpanTags
+    private var customMetricTags = initialMetricTags
     private var annotations = List.empty[Span.Annotation]
 
     def annotate(annotation: Annotation): Span = synchronized {
@@ -201,9 +201,9 @@ object Span {
   }
 
   object Real {
-    def apply(spanContext: SpanContext, initialOperationName: String, initialTags: Map[String, Span.TagValue],
-        startTimestampMicros: Long, reporterRegistry: ReporterRegistryImpl, tracer: Tracer): Real =
-      new Real(spanContext, initialOperationName, initialTags, startTimestampMicros, reporterRegistry, tracer)
+    def apply(spanContext: SpanContext, initialOperationName: String, initialSpanTags: Map[String, Span.TagValue],
+        initialMetricTags: Map[String, String], startTimestampMicros: Long, reporterRegistry: ReporterRegistryImpl, tracer: Tracer): Real =
+      new Real(spanContext, initialOperationName, initialSpanTags, initialMetricTags, startTimestampMicros, reporterRegistry, tracer)
   }
 
   sealed trait TagValue
diff --git a/kamon-core/src/main/scala/kamon/trace/SpanContextCodec.scala b/kamon-core/src/main/scala/kamon/trace/SpanContextCodec.scala
index 11d6de2c..43b5e8e4 100644
--- a/kamon-core/src/main/scala/kamon/trace/SpanContextCodec.scala
+++ b/kamon-core/src/main/scala/kamon/trace/SpanContextCodec.scala
@@ -44,9 +44,12 @@ object SpanContextCodec {
         carrier.put(Headers.TraceIdentifier, urlEncode(spanContext.traceID.string))
         carrier.put(Headers.SpanIdentifier, urlEncode(spanContext.spanID.string))
         carrier.put(Headers.ParentSpanIdentifier, urlEncode(spanContext.parentID.string))
-        carrier.put(Headers.Sampled, encodeSamplingDecision(spanContext.samplingDecision))
         carrier.put(Headers.Baggage, encodeBaggage(spanContext.baggage))
 
+        encodeSamplingDecision(spanContext.samplingDecision).foreach { samplingDecision =>
+          carrier.put(Headers.Sampled, samplingDecision)
+        }
+
         spanContext.baggage.get(Headers.Flags).foreach { flags =>
           carrier.put(Headers.Flags, flags)
         }
@@ -55,11 +58,8 @@ object SpanContextCodec {
       carrier
     }
 
-    override def inject(spanContext: SpanContext): TextMap = {
-      val mutableTextMap = TextMap.Default()
-      inject(spanContext, mutableTextMap)
-      mutableTextMap
-    }
+    override def inject(spanContext: SpanContext): TextMap =
+      inject(spanContext, TextMap.Default())
 
     override def extract(carrier: TextMap): Option[SpanContext] = {
       val traceID = carrier.get(Headers.TraceIdentifier)
@@ -97,14 +97,16 @@ object SpanContextCodec {
       if(baggage.getAll().nonEmpty) {
         val encodedBaggage = new StringBuilder()
         baggage.getAll().foreach {
-          case (key, value) if key != Headers.Flags =>
-            if(encodedBaggage.length() > 0)
-              encodedBaggage.append(';')
-
-            encodedBaggage
-              .append(urlEncode(key))
-              .append('=')
-              .append(urlEncode(value))
+          case (key, value) =>
+            if(key != Headers.Flags) {
+              if (encodedBaggage.length() > 0)
+                encodedBaggage.append(';')
+
+              encodedBaggage
+                .append(urlEncode(key))
+                .append('=')
+                .append(urlEncode(value))
+            }
         }
 
         encodedBaggage.toString()
@@ -125,10 +127,10 @@ object SpanContextCodec {
       baggage
     }
 
-    private def encodeSamplingDecision(samplingDecision: SamplingDecision): String = samplingDecision match {
-      case SamplingDecision.Sample      => "1"
-      case SamplingDecision.DoNotSample => "0"
-      case SamplingDecision.Unknown     => ""
+    private def encodeSamplingDecision(samplingDecision: SamplingDecision): Option[String] = samplingDecision match {
+      case SamplingDecision.Sample      => Some("1")
+      case SamplingDecision.DoNotSample => Some("0")
+      case SamplingDecision.Unknown     => None
     }
 
     private def urlEncode(s: String): String = URLEncoder.encode(s, "UTF-8")
diff --git a/kamon-core/src/main/scala/kamon/trace/Tracer.scala b/kamon-core/src/main/scala/kamon/trace/Tracer.scala
index 47b633ac..714f0215 100644
--- a/kamon-core/src/main/scala/kamon/trace/Tracer.scala
+++ b/kamon-core/src/main/scala/kamon/trace/Tracer.scala
@@ -148,7 +148,8 @@ object Tracer {
   final class SpanBuilder(operationName: String, tracer: Tracer.Default, reporterRegistry: ReporterRegistryImpl) {
     private var parentContext: SpanContext = _
     private var startTimestamp = 0L
-    private var initialTags = Map.empty[String, Span.TagValue]
+    private var initialSpanTags = Map.empty[String, Span.TagValue]
+    private var initialMetricTags = Map.empty[String, String]
     private var useActiveSpanAsParent = true
 
     def asChildOf(parentContext: SpanContext): SpanBuilder = {
@@ -159,19 +160,24 @@ object Tracer {
     def asChildOf(parentSpan: Span): SpanBuilder =
       asChildOf(parentSpan.context())
 
+    def withMetricTag(key: String, value: String): SpanBuilder = {
+      this.initialMetricTags = this.initialMetricTags + (key -> value)
+      this
+    }
+
     def withSpanTag(key: String, value: String): SpanBuilder = {
-      this.initialTags = this.initialTags + (key -> TagValue.String(value))
+      this.initialSpanTags = this.initialSpanTags + (key -> TagValue.String(value))
       this
     }
 
     def withSpanTag(key: String, value: Long): SpanBuilder = {
-      this.initialTags = this.initialTags + (key -> TagValue.Number(value))
+      this.initialSpanTags = this.initialSpanTags + (key -> TagValue.Number(value))
       this
     }
 
     def withSpanTag(key: String, value: Boolean): SpanBuilder = {
       val tagValue = if (value) TagValue.True else TagValue.False
-      this.initialTags = this.initialTags + (key -> tagValue)
+      this.initialSpanTags = this.initialSpanTags + (key -> tagValue)
       this
     }
 
@@ -195,7 +201,7 @@ object Tracer {
       val samplingDecision: SamplingDecision = parentSpanContext
         .map(_.samplingDecision)
         .filter(_ != SamplingDecision.Unknown)
-        .getOrElse(tracer.sampler.decide(operationName, initialTags))
+        .getOrElse(tracer.sampler.decide(operationName, initialSpanTags))
 
       val spanContext = parentSpanContext match {
         case Some(parent) => joinParentContext(parent, samplingDecision)
@@ -203,7 +209,7 @@ object Tracer {
       }
 
       tracer.tracerMetrics.createdSpans.increment()
-      Span.Real(spanContext, operationName, initialTags, startTimestampMicros, reporterRegistry, tracer)
+      Span.Real(spanContext, operationName, initialSpanTags, initialMetricTags, startTimestampMicros, reporterRegistry, tracer)
     }
 
     private def joinParentContext(parent: SpanContext, samplingDecision: SamplingDecision): SpanContext =