fixed the instrumentation to work nicely with spray

2 files changed, 42 insertions, 14 deletions

diff --git a/kamon-core/src/test/scala/akka/instrumentation/ActorInstrumentationSpec.scala b/kamon-core/src/test/scala/akka/instrumentation/ActorInstrumentationSpec.scala
index 0026d953..ccc7740b 100644
--- a/kamon-core/src/test/scala/akka/instrumentation/ActorInstrumentationSpec.scala
+++ b/kamon-core/src/test/scala/akka/instrumentation/ActorInstrumentationSpec.scala
@@ -1,13 +1,18 @@
 package akka.instrumentation
 
 import org.scalatest.{WordSpecLike, Matchers}
-import akka.actor.{Actor, Props, ActorSystem}
+import akka.actor.{ActorRef, Actor, Props, ActorSystem}
 
 import akka.testkit.{ImplicitSender, TestKit}
-import kamon.{TraceContext, Kamon}
+import kamon.{TraceContext, Tracer}
+import akka.pattern.{pipe, ask}
+import akka.util.Timeout
+import scala.concurrent.duration._
+import akka.routing.RoundRobinRouter
 
 
 class ActorInstrumentationSpec extends TestKit(ActorSystem("ActorInstrumentationSpec")) with WordSpecLike with Matchers with ImplicitSender {
+  implicit val executionContext = system.dispatcher
 
   "an instrumented actor ref" when {
     "used inside the context of a transaction" should {
@@ -17,28 +22,51 @@ class ActorInstrumentationSpec extends TestKit(ActorSystem("ActorInstrumentation
         expectMsg(Some(testTraceContext))
       }
 
-      "propagate the trace context using tell" in {
+      "propagate the trace context using tell" in new TraceContextEchoFixture {
+        echo.tell("test", testActor)
 
+        expectMsg(Some(testTraceContext))
+      }
+
+      "propagate the trace context using ask" in new TraceContextEchoFixture {
+        implicit val timeout = Timeout(1 seconds)
+        (echo ? "test") pipeTo(testActor)
+
+        expectMsg(Some(testTraceContext))
       }
 
-      "propagate the trace context using ask" in {
+      "propagate the trace context to actors behind a rounter" in new RoutedTraceContextEchoFixture {
+        val contexts: Seq[Option[TraceContext]] = for(_ <- 1 to 10) yield Some(tellWithNewContext(echo, "test"))
 
+        expectMsgAllOf(contexts: _*)
       }
     }
   }
 
   trait TraceContextEchoFixture {
-    val testTraceContext = Kamon.newTraceContext()
+    val testTraceContext = Tracer.newTraceContext()
     val echo = system.actorOf(Props[TraceContextEcho])
 
-    Kamon.set(testTraceContext)
+    Tracer.set(testTraceContext)
+  }
+
+  trait RoutedTraceContextEchoFixture extends TraceContextEchoFixture {
+    override val echo = system.actorOf(Props[TraceContextEcho].withRouter(RoundRobinRouter(nrOfInstances = 10)))
+
+    def tellWithNewContext(target: ActorRef, message: Any): TraceContext = {
+      val context = Tracer.newTraceContext()
+      Tracer.set(context)
+
+      target ! message
+      context
+    }
   }
 
 }
 
 class TraceContextEcho extends Actor {
   def receive = {
-    case msg ⇒ sender ! Kamon.context()
+    case msg: String ⇒ sender ! Tracer.context()
   }
 }
 
diff --git a/kamon-core/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala b/kamon-core/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala
index de65aaca..fe89695b 100644
--- a/kamon-core/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala
+++ b/kamon-core/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala
@@ -3,7 +3,7 @@ package kamon.instrumentation
 import scala.concurrent.{Await, Promise, Future}
 import org.scalatest.{Matchers, OptionValues, WordSpec}
 import org.scalatest.concurrent.{ScalaFutures, PatienceConfiguration}
-import kamon.{Kamon, TraceContext}
+import kamon.{Tracer, Kamon, TraceContext}
 import java.util.UUID
 import scala.util.Success
 import scala.concurrent.duration._
@@ -27,7 +27,7 @@ class RunnableInstrumentationSpec extends WordSpec with Matchers with ScalaFutur
             val onCompleteContext = Promise[TraceContext]()
 
             futureWithContext.onComplete({
-              case _ => onCompleteContext.complete(Success(Kamon.context.get))
+              case _ => onCompleteContext.complete(Success(Tracer.context.get))
             })
 
             whenReady(onCompleteContext.future) { result =>
@@ -49,7 +49,7 @@ class RunnableInstrumentationSpec extends WordSpec with Matchers with ScalaFutur
         val onCompleteContext = Promise[Option[TraceContext]]()
 
         futureWithoutContext.onComplete({
-          case _ => onCompleteContext.complete(Success(Kamon.context))
+          case _ => onCompleteContext.complete(Success(Tracer.context))
         })
 
         whenReady(onCompleteContext.future) { result =>
@@ -68,14 +68,14 @@ class RunnableInstrumentationSpec extends WordSpec with Matchers with ScalaFutur
 
   class FutureWithContextFixture {
     val testContext = TraceContext()
-    Kamon.set(testContext)
+    Tracer.set(testContext)
 
-    val futureWithContext = Future { Kamon.context}
+    val futureWithContext = Future { Tracer.context}
   }
 
   trait FutureWithoutContextFixture {
-    Kamon.clear // Make sure no TraceContext is available
-    val futureWithoutContext = Future { Kamon.context }
+    Tracer.clear // Make sure no TraceContext is available
+    val futureWithoutContext = Future { Tracer.context }
   }
 }