! core: move the TraceContext manipulation API to the Tracer companion object.

4 files changed, 26 insertions, 27 deletions

diff --git a/kamon-spray/src/main/scala/kamon/spray/KamonTraceDirectives.scala b/kamon-spray/src/main/scala/kamon/spray/KamonTraceDirectives.scala
index 4eefee95..85f93ca8 100644
--- a/kamon-spray/src/main/scala/kamon/spray/KamonTraceDirectives.scala
+++ b/kamon-spray/src/main/scala/kamon/spray/KamonTraceDirectives.scala
@@ -17,11 +17,11 @@ package kamon.spray
 
 import spray.routing.directives.BasicDirectives
 import spray.routing._
-import kamon.trace.TraceContext
+import kamon.trace.Tracer
 
 trait KamonTraceDirectives extends BasicDirectives {
   def traceName(name: String): Directive0 = mapRequest { req ⇒
-    TraceContext.currentContext.rename(name)
+    Tracer.currentContext.rename(name)
     req
   }
 }
diff --git a/kamon-spray/src/main/scala/kamon/spray/instrumentation/ClientRequestInstrumentation.scala b/kamon-spray/src/main/scala/kamon/spray/instrumentation/ClientRequestInstrumentation.scala
index d1e9036d..22d6e354 100644
--- a/kamon-spray/src/main/scala/kamon/spray/instrumentation/ClientRequestInstrumentation.scala
+++ b/kamon-spray/src/main/scala/kamon/spray/instrumentation/ClientRequestInstrumentation.scala
@@ -47,7 +47,7 @@ class ClientRequestInstrumentation {
     // This read to requestContext.traceContext takes care of initializing the aspect timely.
     requestContext.traceContext
 
-    TraceContext.map { ctx ⇒
+    Tracer.currentContext.collect { ctx ⇒
       val sprayExtension = Kamon.extension(Spray)
 
       if (sprayExtension.settings.clientInstrumentationLevel == ClientInstrumentationLevel.HostLevelAPI) {
@@ -74,7 +74,7 @@ class ClientRequestInstrumentation {
 
   @Around("copyingRequestContext(old)")
   def aroundCopyingRequestContext(pjp: ProceedingJoinPoint, old: TraceContextAware): Any = {
-    TraceContext.withContext(old.traceContext) {
+    Tracer.withContext(old.traceContext) {
       pjp.proceed()
     }
   }
@@ -85,7 +85,7 @@ class ClientRequestInstrumentation {
   @Around("dispatchToCommander(requestContext, message)")
   def aroundDispatchToCommander(pjp: ProceedingJoinPoint, requestContext: TraceContextAware, message: Any): Any = {
     if (requestContext.traceContext.nonEmpty) {
-      TraceContext.withContext(requestContext.traceContext) {
+      Tracer.withContext(requestContext.traceContext) {
         if (message.isInstanceOf[HttpMessageEnd])
           requestContext.asInstanceOf[SegmentAware].segment.finish()
 
@@ -112,7 +112,7 @@ class ClientRequestInstrumentation {
     val originalSendReceive = pjp.proceed().asInstanceOf[HttpRequest ⇒ Future[HttpResponse]]
 
     (request: HttpRequest) ⇒ {
-      TraceContext.map { ctx ⇒
+      Tracer.currentContext.collect { ctx ⇒
         val sprayExtension = Kamon.extension(Spray)
         val segment =
           if (sprayExtension.settings.clientInstrumentationLevel == ClientInstrumentationLevel.RequestLevelAPI)
@@ -139,7 +139,7 @@ class ClientRequestInstrumentation {
   @Around("includingDefaultHeadersAtHttpHostConnector(request, defaultHeaders)")
   def aroundIncludingDefaultHeadersAtHttpHostConnector(pjp: ProceedingJoinPoint, request: HttpMessage, defaultHeaders: List[HttpHeader]): Any = {
 
-    val modifiedHeaders = TraceContext.map { ctx ⇒
+    val modifiedHeaders = Tracer.currentContext.collect { ctx ⇒
       val sprayExtension = Kamon.extension(Spray)
       if (sprayExtension.settings.includeTraceTokenHeader)
         RawHeader(sprayExtension.settings.traceTokenHeaderName, ctx.token) :: defaultHeaders
diff --git a/kamon-spray/src/main/scala/kamon/spray/instrumentation/ServerRequestInstrumentation.scala b/kamon-spray/src/main/scala/kamon/spray/instrumentation/ServerRequestInstrumentation.scala
index bf20d167..5be156cd 100644
--- a/kamon-spray/src/main/scala/kamon/spray/instrumentation/ServerRequestInstrumentation.scala
+++ b/kamon-spray/src/main/scala/kamon/spray/instrumentation/ServerRequestInstrumentation.scala
@@ -18,7 +18,6 @@ package spray.can.server.instrumentation
 import kamon.trace.TraceLocal.{ HttpContext, HttpContextKey }
 import org.aspectj.lang.annotation._
 import kamon.trace._
-import akka.actor.ActorSystem
 import spray.can.server.OpenRequest
 import spray.http.{ HttpResponse, HttpMessagePartWrapper, HttpRequest }
 import kamon.Kamon
@@ -48,7 +47,7 @@ class ServerRequestInstrumentation {
     } else None
 
     val newContext = token.map(customToken ⇒ tracer.newContext(defaultTraceName, customToken)) getOrElse (tracer.newContext(defaultTraceName))
-    TraceContext.setCurrentContext(newContext)
+    Tracer.setCurrentContext(newContext)
 
     // Necessary to force initialization of traceContext when initiating the request.
     openRequest.traceContext
@@ -59,7 +58,7 @@ class ServerRequestInstrumentation {
 
   @After("openNewRequest()")
   def afterOpenNewRequest(): Unit = {
-    TraceContext.clearCurrentContext
+    Tracer.clearCurrentContext
   }
 
   @Pointcut("execution(* spray.can.server.OpenRequestComponent$DefaultOpenRequest.handleResponseEndAndReturnNextOpenRequest(..)) && target(openRequest) && args(response)")
@@ -67,7 +66,7 @@ class ServerRequestInstrumentation {
 
   @Around("openRequestCreation(openRequest, response)")
   def afterFinishingRequest(pjp: ProceedingJoinPoint, openRequest: TraceContextAware, response: HttpMessagePartWrapper): Any = {
-    val incomingContext = TraceContext.currentContext
+    val incomingContext = Tracer.currentContext
     val storedContext = openRequest.traceContext
 
     // The stored context is always a DefaultTraceContext if the instrumentation is running
@@ -84,7 +83,7 @@ class ServerRequestInstrumentation {
 
       } else pjp.proceed
 
-      TraceContext.currentContext.finish()
+      Tracer.currentContext.finish()
 
       recordHttpServerMetrics(response, incomingContext.name, sprayExtension)
 
diff --git a/kamon-spray/src/test/scala/kamon/spray/ClientRequestInstrumentationSpec.scala b/kamon-spray/src/test/scala/kamon/spray/ClientRequestInstrumentationSpec.scala
index a99bbdcc..67e6725f 100644
--- a/kamon-spray/src/test/scala/kamon/spray/ClientRequestInstrumentationSpec.scala
+++ b/kamon-spray/src/test/scala/kamon/spray/ClientRequestInstrumentationSpec.scala
@@ -22,7 +22,7 @@ import org.scalatest.concurrent.ScalaFutures
 import org.scalatest.time.{ Millis, Seconds, Span }
 import spray.httpx.RequestBuilding
 import spray.http.{ HttpResponse, HttpRequest }
-import kamon.trace.{ TraceContext, SegmentCategory }
+import kamon.trace.{ Tracer, SegmentCategory }
 import com.typesafe.config.ConfigFactory
 import spray.can.Http
 import spray.http.HttpHeaders.RawHeader
@@ -57,12 +57,12 @@ class ClientRequestInstrumentationSpec extends BaseKamonSpec("client-request-ins
         val (_, server, bound) = buildSHostConnectorAndServer
 
         // Initiate a request within the context of a trace
-        val (testContext, responseFuture) = TraceContext.withContext(newContext("include-trace-token-header-at-request-level-api")) {
+        val (testContext, responseFuture) = Tracer.withContext(newContext("include-trace-token-header-at-request-level-api")) {
           val rF = sendReceive(system, ec) {
             Get(s"http://${bound.localAddress.getHostName}:${bound.localAddress.getPort}/dummy-path")
           }
 
-          (TraceContext.currentContext, rF)
+          (Tracer.currentContext, rF)
         }
 
         // Accept the connection at the server side
@@ -84,12 +84,12 @@ class ClientRequestInstrumentationSpec extends BaseKamonSpec("client-request-ins
         val (_, server, bound) = buildSHostConnectorAndServer
 
         // Initiate a request within the context of a trace
-        val (testContext, responseFuture) = TraceContext.withContext(newContext("do-not-include-trace-token-header-at-request-level-api")) {
+        val (testContext, responseFuture) = Tracer.withContext(newContext("do-not-include-trace-token-header-at-request-level-api")) {
           val rF = sendReceive(system, ec) {
             Get(s"http://${bound.localAddress.getHostName}:${bound.localAddress.getPort}/dummy-path")
           }
 
-          (TraceContext.currentContext, rF)
+          (Tracer.currentContext, rF)
         }
 
         // Accept the connection at the server side
@@ -114,12 +114,12 @@ class ClientRequestInstrumentationSpec extends BaseKamonSpec("client-request-ins
         val (_, _, bound) = buildSHostConnectorAndServer
 
         // Initiate a request within the context of a trace
-        val (testContext, responseFuture) = TraceContext.withContext(newContext("assign-name-to-segment-with-request-level-api")) {
+        val (testContext, responseFuture) = Tracer.withContext(newContext("assign-name-to-segment-with-request-level-api")) {
           val rF = sendReceive(transport.ref)(ec, 10.seconds) {
             Get(s"http://${bound.localAddress.getHostName}:${bound.localAddress.getPort}/request-level-api-segment")
           }
 
-          (TraceContext.currentContext, rF)
+          (Tracer.currentContext, rF)
         }
 
         // Receive the request and reply back
@@ -141,12 +141,12 @@ class ClientRequestInstrumentationSpec extends BaseKamonSpec("client-request-ins
         val (_, server, bound) = buildSHostConnectorAndServer
 
         // Initiate a request within the context of a trace
-        val (testContext, responseFuture) = TraceContext.withContext(newContext("rename-segment-with-request-level-api")) {
+        val (testContext, responseFuture) = Tracer.withContext(newContext("rename-segment-with-request-level-api")) {
           val rF = sendReceive(system, ec) {
             Get(s"http://${bound.localAddress.getHostName}:${bound.localAddress.getPort}/request-level-api-segment")
           }
 
-          (TraceContext.currentContext, rF)
+          (Tracer.currentContext, rF)
         }
 
         // Accept the connection at the server side
@@ -175,9 +175,9 @@ class ClientRequestInstrumentationSpec extends BaseKamonSpec("client-request-ins
         val client = TestProbe()
 
         // Initiate a request within the context of a trace
-        val testContext = TraceContext.withContext(newContext("include-trace-token-header-on-http-client-request")) {
+        val testContext = Tracer.withContext(newContext("include-trace-token-header-on-http-client-request")) {
           client.send(hostConnector, Get("/dummy-path"))
-          TraceContext.currentContext
+          Tracer.currentContext
         }
 
         // Accept the connection at the server side
@@ -202,9 +202,9 @@ class ClientRequestInstrumentationSpec extends BaseKamonSpec("client-request-ins
         val client = TestProbe()
 
         // Initiate a request within the context of a trace
-        val testContext = TraceContext.withContext(newContext("not-include-trace-token-header-on-http-client-request")) {
+        val testContext = Tracer.withContext(newContext("not-include-trace-token-header-on-http-client-request")) {
           client.send(hostConnector, Get("/dummy-path"))
-          TraceContext.currentContext
+          Tracer.currentContext
         }
 
         // Accept the connection at the server side
@@ -229,9 +229,9 @@ class ClientRequestInstrumentationSpec extends BaseKamonSpec("client-request-ins
         val client = TestProbe()
 
         // Initiate a request within the context of a trace
-        val testContext = TraceContext.withContext(newContext("create-segment-with-host-level-api")) {
+        val testContext = Tracer.withContext(newContext("create-segment-with-host-level-api")) {
           client.send(hostConnector, Get("/host-level-api-segment"))
-          TraceContext.currentContext
+          Tracer.currentContext
         }
 
         // Accept the connection at the server side