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

3 files changed, 11 insertions, 12 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)