wip

4 files changed, 24 insertions, 35 deletions

diff --git a/src/main/scala/akka/instrumentation/ActorRefTellInstrumentation.scala b/src/main/scala/akka/instrumentation/ActorRefTellInstrumentation.scala
index 00e4e066..f631b79a 100644
--- a/src/main/scala/akka/instrumentation/ActorRefTellInstrumentation.scala
+++ b/src/main/scala/akka/instrumentation/ActorRefTellInstrumentation.scala
@@ -20,7 +20,7 @@ class ActorRefTellInstrumentation {
   def around(pjp: ProceedingJoinPoint, message: Any, sender: ActorRef): Unit  = {
     import pjp._
 
-    Kamon.context match {
+    Kamon.context() match {
       case Some(ctx) => {
         val traceableMessage = TraceableMessage(ctx, message)
         proceed(getArgs.updated(0, traceableMessage))
diff --git a/src/main/scala/kamon/Kamon.scala b/src/main/scala/kamon/Kamon.scala
index b5998f81..4cba2342 100644
--- a/src/main/scala/kamon/Kamon.scala
+++ b/src/main/scala/kamon/Kamon.scala
@@ -6,15 +6,19 @@ object Kamon {
 
   implicit val actorSystem = ActorSystem("kamon")
 
-  private val ctx = new ThreadLocal[Option[TraceContext]] {
+  val ctx = new ThreadLocal[Option[TraceContext]] {
     override def initialValue() = None
   }
 
-  def context = ctx.get()
+  def context() = ctx.get()
   def clear = ctx.remove()
   def set(traceContext: TraceContext) = ctx.set(Some(traceContext))
 
-  def start: Unit = set(newTraceContext)
+  def start = set(newTraceContext)
+  def stop = ctx.get match {
+    case Some(context) => context.close
+    case None =>
+  }
 
   def newTraceContext(): TraceContext = TraceContext()
 
diff --git a/src/main/scala/kamon/executor/eventbus.scala b/src/main/scala/kamon/executor/eventbus.scala
index d83f2ac6..ed76334f 100644
--- a/src/main/scala/kamon/executor/eventbus.scala
+++ b/src/main/scala/kamon/executor/eventbus.scala
@@ -5,7 +5,7 @@ import akka.event.LookupClassification
 import akka.actor._
 import java.util.concurrent.TimeUnit
 
-import kamon.{Kamon, TraceContext}
+import kamon.{CodeBlockExecutionTime, Kamon, TraceContext}
 import akka.util.Timeout
 import scala.util.Success
 import scala.util.Failure
@@ -88,34 +88,20 @@ object TryAkka extends App{
   import akka.pattern.ask
   implicit val timeout = Timeout(10, TimeUnit.SECONDS)
   implicit def execContext = system.dispatcher
-  //for(i <- 1 to 8) {
-/*  val i = 1
-    TraceContext.start
-    val ping = system.actorOf(Props(new PingActor(system.actorOf(Props[PongActor], s"ping-${i}"))), s"pong-${i}")
-    val f = ping ? Pong()
-
-  f.map({
-    a => threadPrintln(s"In the map body, with the context: ${TraceContext.current}")
-  })
-  .flatMap({
-    (a: Any) => {
-      threadPrintln(s"Executing the flatMap, with the context: ${TraceContext.current}")
-      Future { s"In the flatMap body, with the context: ${TraceContext.current}" }
-    }
-  })
-    .onComplete({
-    case Success(p) => threadPrintln(s"On my main success, with String [$p] and the context: ${TraceContext.current}")
-    case Failure(t) => threadPrintln(s"Something went wrong in the main, with the context: ${TraceContext.current}")
-  })*/
-  //}
+
+
 
   Kamon.start
+
+  Kamon.context.get.append(CodeBlockExecutionTime("some-block", System.nanoTime(), System.nanoTime()))
   threadPrintln("Before doing it")
   val f = Future { threadPrintln("This is happening inside the future body") }
 
-  Kamon.context.get.close
+  Kamon.stop
 
 
+  Thread.sleep(3000)
+  system.shutdown()
 
 /*  appActorEventBus.subscribe(subscriber, NEW_POST_CHANNEL)
   appActorEventBus.publish(MessageEvent(NEW_POST_CHANNEL,PostMessage(text="hello world")))*/
diff --git a/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala b/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala
index f2e83824..4fe9e617 100644
--- a/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala
+++ b/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala
@@ -1,7 +1,6 @@
 package kamon.instrumentation
 
 import scala.concurrent.{Await, Promise, Future}
-import scala.concurrent.ExecutionContext.Implicits.global
 import org.scalatest.{OptionValues, WordSpec}
 import org.scalatest.matchers.MustMatchers
 import org.scalatest.concurrent.PatienceConfiguration
@@ -18,14 +17,14 @@ class RunnableInstrumentationSpec extends WordSpec with MustMatchers with ScalaF
   "a instrumented runnable" when {
     "created in a thread that does have a TraceContext" must {
       "preserve the TraceContext" which {
-        "should be available during the run method execution" in { new FutureWithContext {
+        "should be available during the run method execution" in { new FutureWithContextFixture {
 
             whenReady(futureWithContext) { result =>
               result.value must be === testContext
             }
         }}
 
-        "should be available during the execution of onComplete callbacks" in { new FutureWithContext {
+        "should be available during the execution of onComplete callbacks" in { new FutureWithContextFixture {
             val onCompleteContext = Promise[TraceContext]()
 
             futureWithContext.onComplete({
@@ -40,14 +39,14 @@ class RunnableInstrumentationSpec extends WordSpec with MustMatchers with ScalaF
     }
 
     "created in a thread that doest have a TraceContext" must {
-      "not capture any TraceContext for the body execution" in { new FutureWithoutContext{
+      "not capture any TraceContext for the body execution" in { new FutureWithoutContextFixture{
 
           whenReady(futureWithoutContext) { result =>
             result must be === None
           }
       }}
 
-      "not make any TraceContext available during the onComplete callback" in { new FutureWithoutContext {
+      "not make any TraceContext available during the onComplete callback" in { new FutureWithoutContextFixture {
         val onCompleteContext = Promise[Option[TraceContext]]()
 
         futureWithoutContext.onComplete({
@@ -65,17 +64,17 @@ class RunnableInstrumentationSpec extends WordSpec with MustMatchers with ScalaF
   /**
    *  We are using Futures for the test since they exercise Runnables in the back and also resemble the real use case we have.
    */
+  implicit val testActorSystem = ActorSystem("test-actorsystem")
+  implicit val execContext = testActorSystem.dispatcher
 
-
-  trait FutureWithContext {
-    implicit val as = ActorSystem("test-actorsystem")
+  class FutureWithContextFixture {
     val testContext = TraceContext()
     Kamon.set(testContext)
 
     val futureWithContext = Future { Kamon.context}
   }
 
-  trait FutureWithoutContext {
+  trait FutureWithoutContextFixture {
     Kamon.clear // Make sure no TraceContext is available
     val futureWithoutContext = Future { Kamon.context }
   }