Added a simple test of futures instrumentation

2 files changed, 7 insertions, 11 deletions

diff --git a/src/main/scala/kamon/TraceContext.scala b/src/main/scala/kamon/TraceContext.scala
index 18a91145..5787167b 100644
--- a/src/main/scala/kamon/TraceContext.scala
+++ b/src/main/scala/kamon/TraceContext.scala
@@ -10,19 +10,15 @@ case class TraceContext(id: UUID, entries: List[TraceEntry]) {
 }
 
 object TraceContext {
-  private val context = new ThreadLocal[TraceContext]
-
-  def current = {
-    val ctx = context.get()
-    if(ctx ne null)
-      Some(ctx)
-    else
-      None
+  private val context = new ThreadLocal[Option[TraceContext]] {
+    override def initialValue(): Option[TraceContext] = None
   }
 
-  def clear = context.remove()
+  def current = context.get
 
-  def set(ctx: TraceContext) = context.set(ctx)
+  def clear = context.remove
+
+  def set(ctx: TraceContext) = context.set(Some(ctx))
 
   def start = set(TraceContext(UUID.randomUUID(), Nil))
 }
diff --git a/src/main/scala/kamon/instrumentation/PromiseCompletingRunnableInstrumentation.scala b/src/main/scala/kamon/instrumentation/PromiseCompletingRunnableInstrumentation.scala
index dde0d857..323951f9 100644
--- a/src/main/scala/kamon/instrumentation/PromiseCompletingRunnableInstrumentation.scala
+++ b/src/main/scala/kamon/instrumentation/PromiseCompletingRunnableInstrumentation.scala
@@ -7,7 +7,7 @@ import org.aspectj.lang.ProceedingJoinPoint
 @Aspect("perthis(promiseCompletingRunnableCreation())")
 class PromiseCompletingRunnableInstrumentation {
 
-  private var traceContext: Option[TraceContext] = None
+  @volatile private var traceContext: Option[TraceContext] = None
 
   @Pointcut("execution(scala.concurrent.impl.Future.PromiseCompletingRunnable.new(..))")
   def promiseCompletingRunnableCreation(): Unit = {}