working on ID generation and SpanContext encoding/decoding

1 files changed, 91 insertions, 90 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/SpanContextCodec.scala b/kamon-core/src/main/scala/kamon/trace/SpanContextCodec.scala
index 23eb40db..11d6de2c 100644
--- a/kamon-core/src/main/scala/kamon/trace/SpanContextCodec.scala
+++ b/kamon-core/src/main/scala/kamon/trace/SpanContextCodec.scala
@@ -13,152 +13,134 @@
  * =========================================================================================
  */
 
-
 package kamon.trace
 
+import java.lang.StringBuilder
 import java.net.{URLDecoder, URLEncoder}
 import java.nio.ByteBuffer
-import java.util.concurrent.ThreadLocalRandom
-
 import kamon.trace.SpanContext.{Baggage, SamplingDecision, Source}
-
-import scala.collection.JavaConverters._
-import kamon.util.HexCodec
-
+import scala.collection.mutable
 
 trait SpanContextCodec[T] {
-  def inject(spanContext: SpanContext, carrier: T): Unit
+  def inject(spanContext: SpanContext, carrier: T): T
+  def inject(spanContext: SpanContext): T
   def extract(carrier: T): Option[SpanContext]
 }
 
-trait TextMap {
-  def get(key: String): Option[String]
-  def put(key: String, value: String): Unit
-  def values: Iterator[(String, String)]
-}
-
-
 object SpanContextCodec {
 
-  trait Format[C]
+  sealed trait Format[C]
   object Format {
     case object TextMap extends Format[TextMap]
     case object HttpHeaders extends Format[TextMap]
     case object Binary extends Format[ByteBuffer]
   }
 
-//  val ExtendedB3: SpanContextCodec[TextMap] = new TextMapSpanCodec(
-//    traceIDKey  = "X-B3-TraceId",
-//    parentIDKey = "X-B3-ParentSpanId",
-//    spanIDKey   = "X-B3-SpanId",
-//    sampledKey  =  "X-B3-Sampled",
-//    baggageKey = "X-Kamon-Baggage-",
-//    baggageValueEncoder = urlEncode,
-//    baggageValueDecoder = urlDecode
-//  )
-
-  private def urlEncode(s: String): String = URLEncoder.encode(s, "UTF-8")
-  private def urlDecode(s: String): String = URLDecoder.decode(s, "UTF-8")
-
-  private class ExtendedB3(identityProvider: IdentityProvider) extends SpanContextCodec[TextMap] {
+  class ExtendedB3(identityProvider: IdentityProvider) extends SpanContextCodec[TextMap] {
     import ExtendedB3.Headers
 
+    override def inject(spanContext: SpanContext, carrier: TextMap): TextMap = {
+      if(spanContext != SpanContext.EmptySpanContext) {
+        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))
 
-    override def inject(spanContext: SpanContext, carrier: TextMap): Unit = {
-      carrier.put(Headers.TraceIdentifier, baggageValueEncoder(spanContext.traceID.string))
-      carrier.put(parentIDKey, baggageValueEncoder(spanContext.parentID.string))
-      carrier.put(spanIDKey, baggageValueEncoder(spanContext.spanID.string))
-
-      spanContext.baggage.getAll().foreach {
-        case (key, value) => carrier.put(baggageKey + key, baggageValueEncoder(value))
+        spanContext.baggage.get(Headers.Flags).foreach { flags =>
+          carrier.put(Headers.Flags, flags)
+        }
       }
+
+      carrier
+    }
+
+    override def inject(spanContext: SpanContext): TextMap = {
+      val mutableTextMap = TextMap.Default()
+      inject(spanContext, mutableTextMap)
+      mutableTextMap
     }
 
     override def extract(carrier: TextMap): Option[SpanContext] = {
       val traceID = carrier.get(Headers.TraceIdentifier)
-        .map(identityProvider.traceIdentifierGenerator().from)
+        .map(id => identityProvider.traceIdentifierGenerator().from(urlDecode(id)))
         .getOrElse(IdentityProvider.NoIdentifier)
 
       val spanID = carrier.get(Headers.SpanIdentifier)
-        .map(identityProvider.spanIdentifierGenerator().from)
+        .map(id => identityProvider.spanIdentifierGenerator().from(urlDecode(id)))
         .getOrElse(IdentityProvider.NoIdentifier)
 
       if(traceID != IdentityProvider.NoIdentifier && spanID != IdentityProvider.NoIdentifier) {
         val parentID = carrier.get(Headers.ParentSpanIdentifier)
-          .map(identityProvider.spanIdentifierGenerator().from)
+          .map(id => identityProvider.spanIdentifierGenerator().from(urlDecode(id)))
           .getOrElse(IdentityProvider.NoIdentifier)
 
-        val samplingDecision = carrier.get(Headers.Flags).orElse(carrier.get(Headers.Sampled)) match {
+        val baggage = decodeBaggage(carrier.get(Headers.Baggage))
+        val flags = carrier.get(Headers.Flags)
+
+        flags.foreach { flags =>
+          baggage.add(Headers.Flags, flags)
+        }
+
+        val samplingDecision = flags.orElse(carrier.get(Headers.Sampled)) match {
           case Some(sampled) if sampled == "1" => SamplingDecision.Sample
           case Some(sampled) if sampled == "0" => SamplingDecision.DoNotSample
           case _ => SamplingDecision.Unknown
         }
 
-
-
-
-        Some(SpanContext(traceID, spanID, parentID, samplingDecision, ???, Source.Remote))
+        Some(SpanContext(traceID, spanID, parentID, samplingDecision, baggage, Source.Remote))
 
       } else None
-
-      val minimalSpanContext =
-        for {
-          traceID   <- carrier.get(traceIDKey).map(identityProvider.traceIdentifierGenerator().from)
-          spanID    <- carrier.get(spanIDKey).map(identityProvider.spanIdentifierGenerator().from)
-        } yield {
-
-
-        }
-
-
-
-//      var traceID: String = null
-//      var parentID: String = null
-//      var spanID: String = null
-//      var sampled: String = null
-//      var baggage: Map[String, String] = Map.empty
-//
-//      carrier.iterator().asScala.foreach { entry =>
-//        if(entry.getKey.equals(traceIDKey))
-//          traceID = baggageValueDecoder(entry.getValue)
-//        else if(entry.getKey.equals(parentIDKey))
-//          parentID = baggageValueDecoder(entry.getValue)
-//        else if(entry.getKey.equals(spanIDKey))
-//          spanID = baggageValueDecoder(entry.getValue)
-//        else if(entry.getKey.equals(sampledKey))
-//          sampled = entry.getValue
-//        else if(entry.getKey.startsWith(baggagePrefix))
-//          baggage = baggage + (entry.getKey.substring(baggagePrefix.length) -> baggageValueDecoder(entry.getValue))
-//      }
-//
-//      if(traceID != null && spanID != null) {
-//        val actualParent = if(parentID == null) 0L else decodeLong(parentID)
-//        val isSampled = if(sampled == null) sampler.decide(ThreadLocalRandom.current().nextLong()) else sampled.equals("1")
-//
-//        new SpanContext(decodeLong(traceID), decodeLong(spanID), actualParent, isSampled, baggage)
-//      } else null
-//
-//      None
     }
 
     private def encodeBaggage(baggage: Baggage): String = {
       if(baggage.getAll().nonEmpty) {
-
+        val encodedBaggage = new StringBuilder()
         baggage.getAll().foreach {
-          case (key, value) =>
+          case (key, value) if key != Headers.Flags =>
+            if(encodedBaggage.length() > 0)
+              encodedBaggage.append(';')
+
+            encodedBaggage
+              .append(urlEncode(key))
+              .append('=')
+              .append(urlEncode(value))
         }
+
+        encodedBaggage.toString()
       } else ""
     }
 
-    private def decodeLong(input: String): Long =
-      HexCodec.lowerHexToUnsignedLong(input)
+    private def decodeBaggage(encodedBaggage: Option[String]): Baggage = {
+      val baggage = Baggage()
+      encodedBaggage.foreach { baggageString =>
+        baggageString.split(";").foreach { group =>
+          val pair = group.split("=")
+          if(pair.length >= 2 && pair(0).nonEmpty) {
+            baggage.add(urlDecode(pair(0)), urlDecode(pair(1)))
+          }
+        }
+      }
+
+      baggage
+    }
+
+    private def encodeSamplingDecision(samplingDecision: SamplingDecision): String = samplingDecision match {
+      case SamplingDecision.Sample      => "1"
+      case SamplingDecision.DoNotSample => "0"
+      case SamplingDecision.Unknown     => ""
+    }
 
-    private def encodeLong(input: Long): String =
-      HexCodec.toLowerHex(input)
+    private def urlEncode(s: String): String = URLEncoder.encode(s, "UTF-8")
+    private def urlDecode(s: String): String = URLDecoder.decode(s, "UTF-8")
 
   }
 
   object ExtendedB3 {
+
+    def apply(identityProvider: IdentityProvider): ExtendedB3 =
+      new ExtendedB3(identityProvider)
+
     object Headers {
       val TraceIdentifier = "X-B3-TraceId"
       val ParentSpanIdentifier = "X-B3-ParentSpanId"
@@ -169,3 +151,22 @@ object SpanContextCodec {
     }
   }
 }
+
+trait TextMap {
+  def get(key: String): Option[String]
+  def put(key: String, value: String): Unit
+  def values: Iterator[(String, String)]
+}
+
+object TextMap {
+  class Default extends TextMap {
+    private val storage = mutable.Map.empty[String, String]
+    override def get(key: String): Option[String] = storage.get(key)
+    override def put(key: String, value: String): Unit = storage.put(key, value)
+    override def values: Iterator[(String, String)] = storage.toIterator
+  }
+
+  object Default {
+    def apply(): Default = new Default()
+  }
+}