Added a simple test of futures instrumentation

5 files changed, 61 insertions, 34 deletions

diff --git a/project/AspectJ.scala b/project/AspectJ.scala
index 7ba359eb..d708b0db 100644
--- a/project/AspectJ.scala
+++ b/project/AspectJ.scala
@@ -8,7 +8,8 @@ import com.typesafe.sbt.SbtAspectj.AspectjKeys._
 object AspectJ {
 
   lazy val aspectJSettings = SbtAspectj.aspectjSettings ++ Seq(
-           fork in Test   := true,
-    javaOptions in Test   <++= weaverOptions in Aspectj
+          showWeaveInfo     :=  false,
+           fork in Test     :=  true,
+    javaOptions in Test   <++=  weaverOptions in Aspectj
   )
 }
\ No newline at end of file
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 = {}
diff --git a/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala b/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
index 97e81b44..6f2a3678 100644
--- a/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
+++ b/src/test/scala/kamon/instrumentation/FutureInstrumentationSpec.scala
@@ -1,35 +1,67 @@
 package kamon.instrumentation
 
-import scala.concurrent.Future
+import scala.concurrent.{Promise, Future}
 import scala.concurrent.ExecutionContext.Implicits.global
-import org.scalatest.WordSpec
+import org.scalatest.{OptionValues, WordSpec}
 import org.scalatest.matchers.MustMatchers
 import org.scalatest.concurrent.PatienceConfiguration
 import kamon.TraceContext
 import java.util.UUID
+import scala.util.Success
 
 
-class FutureInstrumentationSpec extends WordSpec with MustMatchers with ScalaFutures with PatienceConfiguration {
+class FutureInstrumentationSpec extends WordSpec with MustMatchers with ScalaFutures with PatienceConfiguration with OptionValues {
 
-  "a instrumented Future" should {
-    "preserve the transaction context available during the future creation" in {
-      new ContextAwareTest {
-        val future = Future { TraceContext.current.get }
+  "a instrumented Future" when {
+    "created in a thread that does have a TraceContext" must {
+      "preserve the TraceContext" which {
+        "should be available during the body's execution" in { new FutureWithContext {
 
-        whenReady(future) { result =>
-          result must be === context
+            whenReady(futureWithContext) { result =>
+              result.value must be === testContext
+            }
+          }
+        }
+
+        "should be available during the execution of onComplete callbacks" in { new FutureWithContext {
+            val onCompleteContext = Promise[TraceContext]()
+
+            futureWithContext.onComplete({
+              case _ => onCompleteContext.complete(Success(TraceContext.current.get))
+            })
+
+            whenReady(onCompleteContext.future) { result =>
+              result must be === testContext
+            }
+          }
         }
       }
     }
 
-    "use the same context available at creation when executing the onComplete callback" in {
+    "created in a thread that doest have a TraceContext" must {
+      "not capture any TraceContext" in { new FutureWithoutContext{
 
+          whenReady(futureWithoutContext) { result =>
+            result must be === None
+          }
+        }
+      }
     }
   }
 
-  trait ContextAwareTest {
-    val context = TraceContext(UUID.randomUUID(), Nil)
-    TraceContext.set(context)
+
+
+
+  trait FutureWithContext {
+    val testContext = TraceContext(UUID.randomUUID(), Nil)
+    TraceContext.set(testContext)
+
+    val futureWithContext = Future { TraceContext.current }
+  }
+
+  trait FutureWithoutContext {
+    TraceContext.clear
+    val futureWithoutContext = Future { TraceContext.current }
   }
 }