add support for simple broadcast String keys

6 files changed, 59 insertions, 22 deletions

diff --git a/kamon-core-tests/src/test/resources/reference.conf b/kamon-core-tests/src/test/resources/reference.conf
index 862cfa32..e69de29b 100644
--- a/kamon-core-tests/src/test/resources/reference.conf
+++ b/kamon-core-tests/src/test/resources/reference.conf
@@ -1,12 +0,0 @@
-kamon {
-
-  context.codecs {
-    http-headers-keys {
-      string-broadcast-key = "kamon.testkit.SimpleStringCodec$Headers"
-    }
-
-    binary-keys {
-      string-broadcast-key = "kamon.testkit.SimpleStringCodec$Binary"
-    }
-  }
-}
\ No newline at end of file
diff --git a/kamon-core/src/main/resources/reference.conf b/kamon-core/src/main/resources/reference.conf
index e140db83..a779ad02 100644
--- a/kamon-core/src/main/resources/reference.conf
+++ b/kamon-core/src/main/resources/reference.conf
@@ -149,6 +149,8 @@ kamon {
       # Size of the encoding buffer for the Binary Codec.
       binary-buffer-size = 256
 
+      string-keys = [ ]
+
       # Codecs to be used when propagating a Context through a HTTP Headers transport.
       http-headers-keys {
         span = "kamon.trace.SpanCodec$B3"
diff --git a/kamon-core/src/main/scala/kamon/context/Codecs.scala b/kamon-core/src/main/scala/kamon/context/Codecs.scala
index b50e991d..8ba70157 100644
--- a/kamon-core/src/main/scala/kamon/context/Codecs.scala
+++ b/kamon-core/src/main/scala/kamon/context/Codecs.scala
@@ -40,10 +40,15 @@ class Codecs(initialConfig: Config) {
     binary
 
   def reconfigure(config: Config): Unit = {
+    import scala.collection.JavaConverters._
     try {
       val codecsConfig = config.getConfig("kamon.context.codecs")
-      httpHeaders = new Codecs.HttpHeaders(readEntryCodecs("http-headers-keys", codecsConfig))
-      binary = new Codecs.Binary(codecsConfig.getBytes("binary-buffer-size"), readEntryCodecs("binary-keys", codecsConfig))
+      val stringKeys = codecsConfig.getStringList("string-keys").asScala
+      val knownHttpHeaderCodecs = readEntryCodecs[TextMap]("http-headers-keys", codecsConfig) ++ stringHeaderCodecs(stringKeys)
+      val knownBinaryCodecs = readEntryCodecs[ByteBuffer]("binary-keys", codecsConfig) ++ stringBinaryCodecs(stringKeys)
+
+      httpHeaders = new Codecs.HttpHeaders(knownHttpHeaderCodecs)
+      binary = new Codecs.Binary(codecsConfig.getBytes("binary-buffer-size"), knownBinaryCodecs)
     } catch {
       case t: Throwable => log.error("Failed to initialize Context Codecs", t)
     }
@@ -66,6 +71,14 @@ class Codecs(initialConfig: Config) {
 
     entries.result()
   }
+
+  private def stringHeaderCodecs(keys: Seq[String]): Map[String, Codecs.ForEntry[TextMap]] = {
+    keys.map(key => (key, new Codecs.StringHeadersCodec(key))).toMap
+  }
+
+  private def stringBinaryCodecs(keys: Seq[String]): Map[String, Codecs.ForEntry[ByteBuffer]] = {
+    keys.map(key => (key, new Codecs.StringBinaryCodec(key))).toMap
+  }
 }
 
 object Codecs {
@@ -210,6 +223,43 @@ object Codecs {
       override def initialValue(): Array[Byte] = Array.ofDim[Byte](size.toInt)
     }
   }
+
+  private class StringHeadersCodec(key: String) extends Codecs.ForEntry[TextMap] {
+    private val dataKey = "X-KamonContext-" + key
+    private val contextKey = Key.broadcast[Option[String]](key, None)
+
+    override def encode(context: Context): TextMap = {
+      val textMap = TextMap.Default()
+      context.get(contextKey).foreach { value =>
+        textMap.put(dataKey, value)
+      }
+
+      textMap
+    }
+
+    override def decode(carrier: TextMap, context: Context): Context = {
+      carrier.get(dataKey) match {
+        case value @ Some(_) => context.withKey(contextKey, value)
+        case None            => context
+      }
+    }
+  }
+
+  private class StringBinaryCodec(key: String) extends Codecs.ForEntry[ByteBuffer] {
+    val emptyBuffer: ByteBuffer = ByteBuffer.allocate(0)
+    private val contextKey = Key.broadcast[Option[String]](key, None)
+
+    override def encode(context: Context): ByteBuffer = {
+      context.get(contextKey) match {
+        case Some(value)  => ByteBuffer.wrap(value.getBytes)
+        case None         => emptyBuffer
+      }
+    }
+
+    override def decode(carrier: ByteBuffer, context: Context): Context = {
+      context.withKey(contextKey, Some(new String(carrier.array())))
+    }
+  }
 }
 
 
diff --git a/kamon-core/src/main/scala/kamon/context/Context.scala b/kamon-core/src/main/scala/kamon/context/Context.scala
index 340a7803..18e77f5f 100644
--- a/kamon-core/src/main/scala/kamon/context/Context.scala
+++ b/kamon-core/src/main/scala/kamon/context/Context.scala
@@ -88,6 +88,9 @@ object Key {
   def broadcast[T](name: String, emptyValue: T): Key[T] =
     new Default[T](name, emptyValue, true)
 
+  def broadcastString(name: String): Key[Option[String]] =
+    new Default[Option[String]](name, None, true)
+
 
   private class Default[T](val name: String, val emptyValue: T, val broadcast: Boolean) extends Key[T] {
     override def hashCode(): Int =
diff --git a/kamon-testkit/src/main/resources/reference.conf b/kamon-testkit/src/main/resources/reference.conf
index bd16a022..44abb535 100644
--- a/kamon-testkit/src/main/resources/reference.conf
+++ b/kamon-testkit/src/main/resources/reference.conf
@@ -1,11 +1,5 @@
 kamon {
   context.codecs {
-    http-headers-keys {
-      string-broadcast-key = "kamon.testkit.SimpleStringCodec$Headers"
-    }
-
-    binary-keys {
-      string-broadcast-key = "kamon.testkit.SimpleStringCodec$Binary"
-    }
+    string-keys = [ "string-broadcast-key" ]
   }
 }
\ No newline at end of file
diff --git a/kamon-testkit/src/main/scala/kamon/testkit/ContextTesting.scala b/kamon-testkit/src/main/scala/kamon/testkit/ContextTesting.scala
index 3330fe68..9f17dff0 100644
--- a/kamon-testkit/src/main/scala/kamon/testkit/ContextTesting.scala
+++ b/kamon-testkit/src/main/scala/kamon/testkit/ContextTesting.scala
@@ -19,7 +19,7 @@ import kamon.context.{Context, Key}
 
 trait ContextTesting {
   val StringKey = Key.local[Option[String]]("string-key", None)
-  val StringBroadcastKey = Key.broadcast[Option[String]]("string-broadcast-key", None)
+  val StringBroadcastKey = Key.broadcastString("string-broadcast-key")
 
   def contextWithLocal(value: String): Context =
     Context.create(StringKey, Some(value))