stop using DynamicVariable for storing the TraceContext

DynamicVariable uses a InheritableThreadLocal which in some cases leads to
capturing a TraceContext upon FJP resize that we don't want in the new Thread,
or at least that is my observation up to now :).

2 files changed, 20 insertions, 8 deletions

diff --git a/kamon-playground/src/main/scala/test/SimpleRequestProcessor.scala b/kamon-playground/src/main/scala/test/SimpleRequestProcessor.scala
index 3afef0ca..333f5031 100644
--- a/kamon-playground/src/main/scala/test/SimpleRequestProcessor.scala
+++ b/kamon-playground/src/main/scala/test/SimpleRequestProcessor.scala
@@ -34,9 +34,12 @@ object SimpleRequestProcessor extends App with SimpleRoutingApp with RequestBuil
   import system.dispatcher
 
   val act = system.actorOf(Props(new Actor {
+    println("Initializing from: " + (new Throwable).getStackTraceString)
     def receive: Actor.Receive = { case any ⇒ sender ! any }
   }), "com")
 
+  Thread.sleep(10000)
+
   implicit val timeout = Timeout(30 seconds)
 
   val pipeline = sendReceive
diff --git a/kamon-trace/src/main/scala/kamon/trace/Trace.scala b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
index 6eff69e5..0985f547 100644
--- a/kamon-trace/src/main/scala/kamon/trace/Trace.scala
+++ b/kamon-trace/src/main/scala/kamon/trace/Trace.scala
@@ -16,7 +16,6 @@
 package kamon.trace
 
 import kamon.Kamon
-import scala.util.DynamicVariable
 import akka.actor._
 import scala.Some
 import kamon.trace.Trace.Register
@@ -31,30 +30,40 @@ object Trace extends ExtensionId[TraceExtension] with ExtensionIdProvider {
   case object Register
 
   /** User API */
-  private[trace] val traceContext = new DynamicVariable[Option[TraceContext]](None)
+  //private[trace] val traceContext = new DynamicVariable[Option[TraceContext]](None)
+  private[trace] val traceContext = new ThreadLocal[Option[TraceContext]] {
+    override def initialValue(): Option[TraceContext] = None
+  }
   private[trace] val tranid = new AtomicLong()
 
-  def context() = traceContext.value
-  def set(ctx: TraceContext) = traceContext.value = Some(ctx)
+  def context() = traceContext.get
+  private def set(ctx: Option[TraceContext]) = traceContext.set(ctx)
 
-  def clear: Unit = traceContext.value = None
+  def clear: Unit = traceContext.remove()
   def start(name: String)(implicit system: ActorSystem): TraceContext = {
     val ctx = newTraceContext(name)
     ctx.start(name)
-    set(ctx)
+    set(Some(ctx))
 
     ctx
   }
 
-  def withContext[T](ctx: Option[TraceContext])(thunk: ⇒ T): T = traceContext.withValue(ctx)(thunk)
+  def withContext[T](ctx: Option[TraceContext])(thunk: ⇒ T): T = {
+    val oldval = context
+    set(ctx)
+
+    try thunk
+    finally set(oldval)
+  }
 
   def transformContext(f: TraceContext ⇒ TraceContext): Unit = {
-    context.map(f).foreach(set(_))
+    context.map(f).foreach(ctx ⇒ set(Some(ctx)))
   }
 
   def finish(): Option[TraceContext] = {
     val ctx = context()
     ctx.map(_.finish)
+    clear
     ctx
   }