= all: upgrade to be compatible with the latest code in core

5 files changed, 37 insertions, 44 deletions

diff --git a/kamon-play/src/main/scala/kamon/play/instrumentation/LoggerLikeInstrumentation.scala b/kamon-play/src/main/scala/kamon/play/instrumentation/LoggerLikeInstrumentation.scala
index 92686ff0..e2ffd3f9 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/LoggerLikeInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/LoggerLikeInstrumentation.scala
@@ -15,7 +15,7 @@
 
 package kamon.play.instrumentation
 
-import kamon.trace.{ TraceContext, TraceContextAware, TraceRecorder }
+import kamon.trace._
 import org.aspectj.lang.ProceedingJoinPoint
 import org.aspectj.lang.annotation._
 import org.slf4j.MDC
@@ -52,21 +52,24 @@ class LoggerLikeInstrumentation {
 object LoggerLikeInstrumentation {
 
   @inline final def withMDC[A](block: ⇒ A): A = {
-    val keys = TraceRecorder.currentContext.map(extractProperties).map(putAndExtractKeys)
+    val keys = putAndExtractKeys(extractProperties(TraceRecorder.currentContext))
 
-    try block finally keys.map(k ⇒ k.foreach(MDC.remove(_)))
+    try block finally keys.foreach(k ⇒ MDC.remove(k))
   }
 
   def putAndExtractKeys(values: Iterable[Map[String, Any]]): Iterable[String] = values.map {
     value ⇒ value.map { case (key, value) ⇒ MDC.put(key, value.toString); key }
   }.flatten
 
-  def extractProperties(ctx: TraceContext): Iterable[Map[String, Any]] = ctx.traceLocalStorage.underlyingStorage.values.map {
-    case traceLocalValue @ (p: Product) ⇒ {
-      val properties = p.productIterator
-      traceLocalValue.getClass.getDeclaredFields.filter(field ⇒ field.getName != "$outer").map(_.getName -> properties.next).toMap
-    }
-    case anything ⇒ Map.empty[String, Any]
+  def extractProperties(traceContext: TraceContext): Iterable[Map[String, Any]] = traceContext match {
+    case ctx: DefaultTraceContext ⇒
+      ctx.traceLocalStorage.underlyingStorage.values.collect {
+        case traceLocalValue @ (p: Product) ⇒ {
+          val properties = p.productIterator
+          traceLocalValue.getClass.getDeclaredFields.filter(field ⇒ field.getName != "$outer").map(_.getName -> properties.next).toMap
+        }
+      }
+    case EmptyTraceContext ⇒ Iterable.empty[Map[String, Any]]
   }
 }
 
diff --git a/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala b/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
index c761e72f..82a43926 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
@@ -54,16 +54,23 @@ class RequestInstrumentation {
   @Around("execution(* play.api.GlobalSettings+.doFilter(*)) && args(next)")
   def aroundDoFilter(pjp: ProceedingJoinPoint, next: EssentialAction): Any = {
     val essentialAction = (requestHeader: RequestHeader) ⇒ {
+      // TODO: Move to a Kamon-specific dispatcher.
       val executor = Kamon(Play)(Akka.system()).defaultDispatcher
 
       def onResult(result: Result): Result = {
-        TraceRecorder.finish()
-        TraceRecorder.currentContext.map { ctx ⇒
-          val playExtension = Kamon(Play)(ctx.system)
+
+        TraceRecorder.withTraceContextAndSystem { (ctx, system) ⇒
+          ctx.finish()
+
+          val playExtension = Kamon(Play)(system)
           recordHttpServerMetrics(result.header, ctx.name, playExtension)
-          if (playExtension.includeTraceToken) result.withHeaders(playExtension.traceTokenHeaderName -> ctx.token)
-          else result
-        }.getOrElse(result)
+
+          if (playExtension.includeTraceToken)
+            result.withHeaders(playExtension.traceTokenHeaderName -> ctx.token)
+          else
+            result
+
+        } getOrElse (result)
       }
 
       //override the current trace name
@@ -76,8 +83,8 @@ class RequestInstrumentation {
   }
 
   @Before("execution(* play.api.GlobalSettings+.onError(..)) && args(request, ex)")
-  def beforeOnError(request: TraceContextAware, ex: Throwable): Unit = request.traceContext.map {
-    ctx ⇒ recordHttpServerMetrics(InternalServerError.header, ctx.name, Kamon(Play)(ctx.system))
+  def beforeOnError(request: TraceContextAware, ex: Throwable): Unit = TraceRecorder.withTraceContextAndSystem { (ctx, system) ⇒
+    recordHttpServerMetrics(InternalServerError.header, ctx.name, Kamon(Play)(system))
   }
 
   private def recordHttpServerMetrics(header: ResponseHeader, traceName: String, playExtension: PlayExtension): Unit =
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 87467050..c2eafa2b 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/WSInstrumentation.scala
@@ -17,12 +17,10 @@
 package kamon.play.instrumentation
 
 import kamon.Kamon
-import kamon.metric.TraceMetrics.HttpClientRequest
 import kamon.play.Play
-import kamon.trace.TraceRecorder
+import kamon.trace.{ SegmentMetricIdentityLabel, SegmentMetricIdentity, 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 scala.concurrent.Future
@@ -35,28 +33,13 @@ class WSInstrumentation {
 
   @Around("onExecuteRequest(request)")
   def aroundExecuteRequest(pjp: ProceedingJoinPoint, request: WSRequest): Any = {
-
-    import kamon.play.instrumentation.WSInstrumentation._
-
-    TraceRecorder.currentContext.map { ctx ⇒
-      val executor = Kamon(Play)(ctx.system).defaultDispatcher
-      val segmentHandle = TraceRecorder.startSegment(HttpClientRequest(request.url), basicRequestAttributes(request))
+    TraceRecorder.withTraceContextAndSystem { (ctx, system) ⇒
+      val executor = Kamon(Play)(system).defaultDispatcher
+      val segment = ctx.startSegment(request.url, SegmentMetricIdentityLabel.HttpClient)
       val response = pjp.proceed().asInstanceOf[Future[WSResponse]]
 
-      response.map(result ⇒ segmentHandle.map(_.finish()))(executor)
+      response.map(result ⇒ segment.finish())(executor)
       response
-    }.getOrElse(pjp.proceed())
-  }
-}
-
-object WSInstrumentation {
-
-  def uri(request: WSRequest): java.net.URI = request.asInstanceOf[NingWSRequest].builder.build().getURI
-
-  def basicRequestAttributes(request: WSRequest): Map[String, String] = {
-    Map[String, String](
-      "host" -> uri(request).getHost,
-      "path" -> uri(request).getPath,
-      "method" -> request.method)
+    } getOrElse (pjp.proceed())
   }
 }
\ No newline at end of file
diff --git a/kamon-play/src/test/scala/kamon/play/RequestInstrumentationSpec.scala b/kamon-play/src/test/scala/kamon/play/RequestInstrumentationSpec.scala
index 2afd31fd..3feb6246 100644
--- a/kamon-play/src/test/scala/kamon/play/RequestInstrumentationSpec.scala
+++ b/kamon-play/src/test/scala/kamon/play/RequestInstrumentationSpec.scala
@@ -117,7 +117,7 @@ class RequestInstrumentationSpec extends PlaySpec with OneServerPerSuite {
 
     "respond to the Async Action with X-Trace-Token and the renamed trace" in {
       val result = Await.result(route(FakeRequest(GET, "/async-renamed").withHeaders(traceTokenHeader)).get, 10 seconds)
-      TraceRecorder.currentContext.map(_.name) must be(Some("renamed-trace"))
+      TraceRecorder.currentContext.name must be("renamed-trace")
       Some(result.header.headers(traceTokenHeaderName)) must be(expectedToken)
     }
 
diff --git a/kamon-play/src/test/scala/kamon/play/WSInstrumentationSpec.scala b/kamon-play/src/test/scala/kamon/play/WSInstrumentationSpec.scala
index b72659d2..bf1ead05 100644
--- a/kamon-play/src/test/scala/kamon/play/WSInstrumentationSpec.scala
+++ b/kamon-play/src/test/scala/kamon/play/WSInstrumentationSpec.scala
@@ -17,9 +17,9 @@
 package kamon.play
 
 import kamon.Kamon
-import kamon.metric.TraceMetrics.{ HttpClientRequest, TraceMetricsSnapshot }
+import kamon.metric.TraceMetrics.TraceMetricsSnapshot
 import kamon.metric.{ Metrics, TraceMetrics }
-import kamon.trace.TraceRecorder
+import kamon.trace.{ SegmentMetricIdentityLabel, SegmentMetricIdentity, TraceRecorder }
 import org.scalatest.{ Matchers, WordSpecLike }
 import org.scalatestplus.play.OneServerPerSuite
 import play.api.libs.ws.WS
@@ -49,7 +49,7 @@ class WSInstrumentationSpec extends WordSpecLike with Matchers with OneServerPer
       val snapshot = takeSnapshotOf("GET: /inside")
       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(SegmentMetricIdentity("http://localhost:19001/async", SegmentMetricIdentityLabel.HttpClient)).numberOfMeasurements should be(1)
     }
 
     "propagate the TraceContext outside an Action and complete the WS request" in {