diff options
Diffstat (limited to 'kamon-core/src')
3 files changed, 77 insertions, 0 deletions
diff --git a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala index 9980a022..307cf17a 100644 --- a/kamon-core/src/main/scala/kamon/trace/TraceContext.scala +++ b/kamon-core/src/main/scala/kamon/trace/TraceContext.scala @@ -32,6 +32,8 @@ trait TraceContext { def levelOfDetail: TracingLevelOfDetail def startSegment(identity: SegmentIdentity, metadata: Map[String, String]): SegmentCompletionHandle def finish(metadata: Map[String, String]) + + private[kamon] val traceLocalStorage: TraceLocalStorage = new TraceLocalStorage } object TraceContext { diff --git a/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala b/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala new file mode 100644 index 00000000..c79fa632 --- /dev/null +++ b/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala @@ -0,0 +1,25 @@ +package kamon.trace + +import scala.collection.concurrent.TrieMap +import kamon.trace.TraceLocal.TraceLocalKey + +object TraceLocal { + trait TraceLocalKey { + type ValueType + } + + def store(key: TraceLocalKey)(value: key.ValueType): Unit = + TraceRecorder.currentContext.map(_.traceLocalStorage.store(key)(value)) + + def retrieve(key: TraceLocalKey): Option[key.ValueType] = + TraceRecorder.currentContext.flatMap(_.traceLocalStorage.retrieve(key)) + +} + +class TraceLocalStorage { + val underlyingStorage = TrieMap[TraceLocal.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]) +} diff --git a/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala b/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala new file mode 100644 index 00000000..53ff8a24 --- /dev/null +++ b/kamon-core/src/test/scala/kamon/trace/TraceLocalSpec.scala @@ -0,0 +1,50 @@ +package kamon.trace + +import akka.testkit.TestKit +import akka.actor.ActorSystem +import org.scalatest.{ OptionValues, Matchers, WordSpecLike } +import org.scalatest.concurrent.PatienceConfiguration + +class TraceLocalSpec extends TestKit(ActorSystem("trace-local-spec")) with WordSpecLike with Matchers + with PatienceConfiguration with OptionValues { + + object SampleTraceLocalKey extends TraceLocal.TraceLocalKey { type ValueType = String } + + "the TraceLocal storage" should { + "allow storing and retrieving values" in { + TraceRecorder.withNewTraceContext("store-and-retrieve-trace-local") { + val testString = "Hello World" + + TraceLocal.store(SampleTraceLocalKey)(testString) + TraceLocal.retrieve(SampleTraceLocalKey).value should equal(testString) + } + } + + "return None when retrieving a non existent key" in { + TraceRecorder.withNewTraceContext("non-existent-key") { + TraceLocal.retrieve(SampleTraceLocalKey) should equal(None) + } + } + + "return None when retrieving a key without a current TraceContext" in { + TraceLocal.retrieve(SampleTraceLocalKey) should equal(None) + } + + "be attached to the TraceContext when it is propagated" in { + val testString = "Hello World" + val testContext = TraceRecorder.withNewTraceContext("manually-propagated-trace-local") { + TraceLocal.store(SampleTraceLocalKey)(testString) + TraceLocal.retrieve(SampleTraceLocalKey).value should equal(testString) + TraceRecorder.currentContext + } + + /** No TraceLocal should be available here */ + TraceLocal.retrieve(SampleTraceLocalKey) should equal(None) + + TraceRecorder.withTraceContext(testContext) { + TraceLocal.retrieve(SampleTraceLocalKey).value should equal(testString) + } + } + } + +} |