+ Ensure that the TraceLocalStorage can be used from Java and close #196

3 files changed, 34 insertions, 16 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala b/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
index 0a0a120a..a99fb776 100644
--- a/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
@@ -17,17 +17,14 @@
 package kamon.trace
 
 import kamon.trace.TraceLocal.TraceLocalKey
-
+import kamon.util.Supplier
 import scala.collection.concurrent.TrieMap
 
 object TraceLocal {
 
-  trait TraceLocalKey {
-    type ValueType
-  }
+  trait TraceLocalKey[T]
 
-  trait AvailableToMdc extends TraceLocalKey {
-    override type ValueType = String
+  trait AvailableToMdc extends TraceLocalKey[String] {
     def mdcKey: String
   }
 
@@ -40,24 +37,32 @@ object TraceLocal {
 
   case class HttpContext(agent: String, uri: String, xforwarded: String)
 
-  object HttpContextKey extends TraceLocal.TraceLocalKey { type ValueType = HttpContext }
+  object HttpContextKey extends TraceLocal.TraceLocalKey[HttpContext]
 
-  def store(key: TraceLocalKey)(value: key.ValueType): Unit = Tracer.currentContext match {
+  def store[T](key: TraceLocalKey[T])(value: Any): Unit = Tracer.currentContext match {
     case ctx: MetricsOnlyContext ⇒ ctx.traceLocalStorage.store(key)(value)
     case EmptyTraceContext       ⇒ // Can't store in the empty context.
   }
 
-  def retrieve(key: TraceLocalKey): Option[key.ValueType] = Tracer.currentContext match {
+  def retrieve[T](key: TraceLocalKey[T]): Option[T] = Tracer.currentContext match {
     case ctx: MetricsOnlyContext ⇒ ctx.traceLocalStorage.retrieve(key)
     case EmptyTraceContext       ⇒ None // Can't retrieve anything from the empty context.
   }
 
+  // Java variant
+  @throws(classOf[NoSuchElementException])
+  def get[T](key: TraceLocalKey[T]): T = retrieve(key).get
+
+  def getOrElse[T](key: TraceLocalKey[T], code: Supplier[T]): T = retrieve(key).getOrElse(code.get)
+
   def storeForMdc(key: String, value: String): Unit = store(AvailableToMdc.fromKey(key))(value)
+
+  def newTraceLocalKey[T]: TraceLocalKey[T] = new TraceLocalKey[T] {}
 }
 
 class TraceLocalStorage {
-  val underlyingStorage = TrieMap[TraceLocal.TraceLocalKey, Any]()
+  val underlyingStorage = TrieMap[TraceLocalKey[_], Any]()
 
-  def store(key: TraceLocalKey)(value: key.ValueType): Unit = underlyingStorage.put(key, value)
-  def retrieve(key: TraceLocalKey): Option[key.ValueType] = underlyingStorage.get(key).map(_.asInstanceOf[key.ValueType])
+  def store[T](key: TraceLocalKey[T])(value: Any): Unit = underlyingStorage.put(key, value)
+  def retrieve[T](key: TraceLocalKey[T]): Option[T] = underlyingStorage.get(key).map(_.asInstanceOf[T])
 }
diff --git a/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala b/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala
index b58b247f..41d5bc83 100644
--- a/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala
@@ -19,6 +19,7 @@ package kamon.trace
 import kamon.testkit.BaseKamonSpec
 import kamon.trace.TraceLocal.AvailableToMdc
 import kamon.trace.logging.MdcKeysSupport
+import kamon.util.Supplier
 import org.scalatest.concurrent.PatienceConfiguration
 import org.scalatest.OptionValues
 import org.slf4j.MDC
@@ -26,7 +27,7 @@ import org.slf4j.MDC
 class TraceLocalSpec extends BaseKamonSpec("trace-local-spec") with PatienceConfiguration with OptionValues with MdcKeysSupport {
   val SampleTraceLocalKeyAvailableToMDC = AvailableToMdc("someKey")
 
-  object SampleTraceLocalKey extends TraceLocal.TraceLocalKey { type ValueType = String }
+  object SampleTraceLocalKey extends TraceLocal.TraceLocalKey[String]
 
   "the TraceLocal storage" should {
     "allow storing and retrieving values" in {
@@ -44,6 +45,20 @@ class TraceLocalSpec extends BaseKamonSpec("trace-local-spec") with PatienceConf
       }
     }
 
+    "throws an exception when trying to get a non existent key" in {
+      Tracer.withContext(newContext("non-existent-key")) {
+        intercept[NoSuchElementException] {
+          TraceLocal.get(SampleTraceLocalKey)
+        }
+      }
+    }
+
+    "return the given value when retrieving a non existent key" in {
+      Tracer.withContext(newContext("non-existent-key")) {
+        TraceLocal.getOrElse(SampleTraceLocalKey, new Supplier[String] { def get = "optionalValue" }) should equal("optionalValue")
+      }
+    }
+
     "return None when retrieving a key without a current TraceContext" in {
       TraceLocal.retrieve(SampleTraceLocalKey) should equal(None)
     }
diff --git a/kamon-play/src/test/scala/kamon/play/RequestInstrumentationSpec.scala b/kamon-play/src/test/scala/kamon/play/RequestInstrumentationSpec.scala
index aebcc4d0..6bae7550 100644
--- a/kamon-play/src/test/scala/kamon/play/RequestInstrumentationSpec.scala
+++ b/kamon-play/src/test/scala/kamon/play/RequestInstrumentationSpec.scala
@@ -179,9 +179,7 @@ class RequestInstrumentationSpec extends PlaySpec with OneServerPerSuite {
 
   object MockGlobalTest extends WithFilters(TraceLocalFilter)
 
-  object TraceLocalKey extends TraceLocal.TraceLocalKey {
-    type ValueType = String
-  }
+  object TraceLocalKey extends TraceLocal.TraceLocalKey[String]
 
   object TraceLocalFilter extends Filter {
     override def apply(next: (RequestHeader) ⇒ Future[Result])(header: RequestHeader): Future[Result] = {