= play: * fix tests

* use kamon default dispatcher

2 files changed, 27 insertions, 26 deletions

diff --git a/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala b/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
index 3e4d6110..14bcd8c9 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
@@ -16,15 +16,16 @@
 
 package kamon.play.instrumentation
 
+import kamon.Kamon
 import kamon.metric.TraceMetrics.HttpClientRequest
-import kamon.trace.{ TraceContext, TraceRecorder }
+import kamon.play.Play
+import kamon.trace.TraceRecorder
 import org.aspectj.lang.ProceedingJoinPoint
 import org.aspectj.lang.annotation.{ Around, Aspect, Pointcut }
 import play.api.libs.ws.ning.NingWSRequest
 import play.api.libs.ws.{ WSRequest, WSResponse }
 import play.libs.Akka
 
-import scala.concurrent.ExecutionContext.Implicits.global
 import scala.concurrent.Future
 
 @Aspect
@@ -38,16 +39,17 @@ class WSInstrumentation {
 
     import kamon.play.instrumentation.WSInstrumentation._
 
-    withOrNewTraceContext(TraceRecorder.currentContext)(request) {
-      val response = pjp.proceed().asInstanceOf[Future[WSResponse]]
-      val segmentHandle = TraceRecorder.startSegment(HttpClientRequest(request.url), basicRequestAttributes(request))
+    TraceRecorder.currentContext match {
+      case ctx @ Some(_) ⇒
+        TraceRecorder.withTraceContext(ctx) {
+          val executor = Kamon(Play)(Akka.system()).defaultDispatcher
+          val segmentHandle = TraceRecorder.startSegment(HttpClientRequest(request.url), basicRequestAttributes(request))
+          val response = pjp.proceed().asInstanceOf[Future[WSResponse]]
 
-      response.map {
-        r ⇒
-          segmentHandle.map(_.finish())
-          TraceRecorder.finish()
-      }
-      response
+          response.map(result ⇒ segmentHandle.map(_.finish()))(executor)
+          response
+        }
+      case None ⇒ pjp.proceed()
     }
   }
 }
@@ -62,9 +64,4 @@ object WSInstrumentation {
       "path" -> uri(request).getPath,
       "method" -> request.method)
   }
-
-  def withOrNewTraceContext[T](context: Option[TraceContext])(request: WSRequest)(thunk: ⇒ T): T = {
-    if (context.isDefined) TraceRecorder.withTraceContext(context) { thunk }
-    else TraceRecorder.withNewTraceContext(request.url, metadata = basicRequestAttributes(request)) { thunk }(Akka.system())
-  }
 }
\ No newline at end of file
diff --git a/kamon-play/src/test/scala/kamon/play/WSInstrumentationSpec.scala b/kamon-play/src/test/scala/kamon/play/WSInstrumentationSpec.scala
index cbd95db3..3378d477 100644
--- a/kamon-play/src/test/scala/kamon/play/WSInstrumentationSpec.scala
+++ b/kamon-play/src/test/scala/kamon/play/WSInstrumentationSpec.scala
@@ -19,6 +19,7 @@ package kamon.play
 import kamon.Kamon
 import kamon.metric.TraceMetrics.{ HttpClientRequest, TraceMetricsSnapshot }
 import kamon.metric.{ Metrics, TraceMetrics }
+import kamon.trace.TraceRecorder
 import org.scalatest.{ Matchers, WordSpecLike }
 import org.scalatestplus.play.OneServerPerSuite
 import play.api.libs.ws.WS
@@ -42,22 +43,25 @@ class WSInstrumentationSpec extends WordSpecLike with Matchers with OneServerPer
   })
 
   "the WS instrumentation" should {
-    "propagate the TraceContext outside an Action and complete the WS request" in {
-      Await.result(WS.url("http://localhost:19001/outside").get(), 10 seconds)
+    "propagate the TraceContext inside an Action and complete the WS request" in {
+      Await.result(route(FakeRequest(GET, "/inside")).get, 10 seconds)
 
-      val snapshot = takeSnapshotOf("http://localhost:19001/outside")
+      val snapshot = takeSnapshotOf("GET: /inside")
       snapshot.elapsedTime.numberOfMeasurements should be(1)
       snapshot.segments.size should be(1)
-      snapshot.segments(HttpClientRequest("http://localhost:19001/outside")).numberOfMeasurements should be(1)
+      snapshot.segments(HttpClientRequest("http://localhost:19001/async")).numberOfMeasurements should be(1)
     }
 
-    "propagate the TraceContext inside an Action and complete the WS request" in {
-      Await.result(route(FakeRequest(GET, "/inside")).get, 10 seconds)
+    "propagate the TraceContext outside an Action and complete the WS request" in {
+      TraceRecorder.withNewTraceContext("trace-outside-action") {
+        Await.result(WS.url("http://localhost:19001/outside").get(), 10 seconds)
+        TraceRecorder.finish()
+      }(Akka.system())
 
-      val snapshot = takeSnapshotOf("GET: /inside")
-      snapshot.elapsedTime.numberOfMeasurements should be(2)
+      val snapshot = takeSnapshotOf("trace-outside-action")
+      snapshot.elapsedTime.numberOfMeasurements should be(1)
       snapshot.segments.size should be(1)
-      snapshot.segments(HttpClientRequest("http://localhost:19001/async")).numberOfMeasurements should be(1)
+      snapshot.segments(HttpClientRequest("http://localhost:19001/outside")).numberOfMeasurements should be(1)
     }
 
   }
@@ -69,8 +73,8 @@ class WSInstrumentationSpec extends WordSpecLike with Matchers with OneServerPer
   }
 
   def callWSinsideController(url: String) = Action.async {
-    import play.api.libs.concurrent.Execution.Implicits.defaultContext
     import play.api.Play.current
+    import play.api.libs.concurrent.Execution.Implicits.defaultContext
 
     WS.url(url).get().map { response ⇒
       Ok("Ok")