From 910830c979699db04a6745cf922631731fda4cd0 Mon Sep 17 00:00:00 2001
From: Ivan Topolnjak <ivantopo@gmail.com>
Date: Tue, 24 Dec 2013 09:33:51 -0300
Subject: 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 :).
---
 kamon-trace/src/main/scala/kamon/trace/Trace.scala | 25 +++++++++++++++-------
 1 file changed, 17 insertions(+), 8 deletions(-)

(limited to 'kamon-trace')

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
   }
 
-- 
cgit v1.2.3