Added a simple test of futures instrumentation

3 files changed, 13 insertions, 19 deletions

diff --git a/src/main/resources/META-INF/aop.xml b/src/main/resources/META-INF/aop.xml
index ae3914cd..bd67064d 100644
--- a/src/main/resources/META-INF/aop.xml
+++ b/src/main/resources/META-INF/aop.xml
@@ -1,9 +1,6 @@
 <!DOCTYPE aspectj PUBLIC "-//AspectJ//DTD//EN" "http://www.eclipse.org/aspectj/dtd/aspectj.dtd">
 
 <aspectj>
-
-    <weaver options="-verbose -showWeaveInfo"/>
-
     <aspects>
         <!--<aspect name="akka.ActorSystemAspect"/>
         &lt;!&ndash;<aspect name="akka.MailboxAspect"/>&ndash;&gt;
@@ -14,11 +11,12 @@
         <aspect name="kamon.instrumentation.PromiseCompletingRunnableInstrumentation" />
 
         <include within="*"/>
-        <exclude within="javax.*"/>
-        <exclude within="org.aspectj.*"/>
-        <exclude within="scala.*"/>
-        <exclude within="scalaz.*"/>
-        <exclude within="scalad.*"/>
+        <exclude within="javax..*"/>
+        <exclude within="org.aspectj..*"/>
+        <exclude within="scala..*"/>
+        <exclude within="scalaz..*"/>
+        <exclude within="scalad..*"/>
+        <exclude within="play..*"/>
     </aspects>
 
 </aspectj>
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 = {}