1 files changed, 22 insertions, 27 deletions
diff --git a/kamon-spray/src/main/scala/spray/can/server/ServerRequestInstrumentation.scala b/kamon-spray/src/main/scala/spray/can/server/ServerRequestInstrumentation.scala
index 69b0160e..eb25412b 100644
--- a/kamon-spray/src/main/scala/spray/can/server/ServerRequestInstrumentation.scala
+++ b/kamon-spray/src/main/scala/spray/can/server/ServerRequestInstrumentation.scala
@@ -16,11 +16,10 @@
 package spray.can.server
 
 import org.aspectj.lang.annotation._
-import kamon.trace.{ TraceContext, TraceRecorder, TraceContextAware }
+import kamon.trace._
 import akka.actor.ActorSystem
 import spray.http.{ HttpResponse, HttpMessagePartWrapper, HttpRequest }
 import akka.event.Logging.Warning
-import scala.Some
 import kamon.Kamon
 import kamon.spray.{ SprayExtension, Spray }
 import org.aspectj.lang.ProceedingJoinPoint
@@ -40,7 +39,7 @@ class ServerRequestInstrumentation {
     val system: ActorSystem = openRequest.asInstanceOf[OpenRequest].context.actorContext.system
     val sprayExtension = Kamon(Spray)(system)
 
-    val defaultTraceName: String = request.method.value + ": " + request.uri.path
+    val defaultTraceName = sprayExtension.generateTraceName(request)
     val token = if (sprayExtension.includeTraceToken) {
       request.headers.find(_.name == sprayExtension.traceTokenHeaderName).map(_.value)
     } else None
@@ -67,40 +66,36 @@ class ServerRequestInstrumentation {
     val incomingContext = TraceRecorder.currentContext
     val storedContext = openRequest.traceContext
 
-    verifyTraceContextConsistency(incomingContext, storedContext)
-    incomingContext match {
-      case None ⇒ pjp.proceed()
-      case Some(traceContext) ⇒
-        val sprayExtension = Kamon(Spray)(traceContext.system)
+    // The stored context is always a DefaultTraceContext if the instrumentation is running
+    val system = storedContext.asInstanceOf[DefaultTraceContext].system
 
-        val proceedResult = if (sprayExtension.includeTraceToken) {
-          val responseWithHeader = includeTraceTokenIfPossible(response, sprayExtension.traceTokenHeaderName, traceContext.token)
-          pjp.proceed(Array(openRequest, responseWithHeader))
+    verifyTraceContextConsistency(incomingContext, storedContext, system)
 
-        } else pjp.proceed
+    if (incomingContext.isEmpty)
+      pjp.proceed()
+    else {
+      val sprayExtension = Kamon(Spray)(system)
 
-        TraceRecorder.finish()
-        recordHttpServerMetrics(response, traceContext.name, sprayExtension)
-        proceedResult
-    }
-  }
+      val proceedResult = if (sprayExtension.includeTraceToken) {
+        val responseWithHeader = includeTraceTokenIfPossible(response, sprayExtension.traceTokenHeaderName, incomingContext.token)
+        pjp.proceed(Array(openRequest, responseWithHeader))
 
-  def verifyTraceContextConsistency(incomingTraceContext: Option[TraceContext], storedTraceContext: Option[TraceContext]): Unit = {
-    for (original ← storedTraceContext) {
-      incomingTraceContext match {
-        case Some(incoming) if original.token != incoming.token ⇒
-          publishWarning(s"Different ids when trying to close a Trace, original: [$original] - incoming: [$incoming]", incoming.system)
+      } else pjp.proceed
 
-        case Some(_) ⇒ // nothing to do here.
-
-        case None ⇒
-          publishWarning(s"Trace context not present while closing the Trace: [$original]", original.system)
-      }
+      TraceRecorder.finish()
+      recordHttpServerMetrics(response, incomingContext.name, sprayExtension)
+      proceedResult
     }
+  }
 
+  def verifyTraceContextConsistency(incomingTraceContext: TraceContext, storedTraceContext: TraceContext, system: ActorSystem): Unit = {
     def publishWarning(text: String, system: ActorSystem): Unit =
       system.eventStream.publish(Warning("", classOf[ServerRequestInstrumentation], text))
 
+    if (incomingTraceContext.nonEmpty && incomingTraceContext.token != storedTraceContext.token)
+      publishWarning(s"Different trace token found when trying to close a trace, original: [${storedTraceContext.token}] - incoming: [${incomingTraceContext.token}]", system)
+    else
+      publishWarning(s"EmptyTraceContext present while closing the trace with token [${storedTraceContext.token}]", system)
   }
 
   def recordHttpServerMetrics(response: HttpMessagePartWrapper, traceName: String, sprayExtension: SprayExtension): Unit =