Merge branch 'wip/newrelic-agent-rewrite'

2 files changed, 15 insertions, 6 deletions

diff --git a/kamon-spray/src/main/scala/kamon/spray/Spray.scala b/kamon-spray/src/main/scala/kamon/spray/Spray.scala
index c1c81116..72b997d7 100644
--- a/kamon-spray/src/main/scala/kamon/spray/Spray.scala
+++ b/kamon-spray/src/main/scala/kamon/spray/Spray.scala
@@ -21,6 +21,7 @@ import akka.actor
 import kamon.Kamon
 import kamon.http.HttpServerMetrics
 import kamon.metric.Metrics
+import spray.http.HttpHeaders.Host
 import spray.http.HttpRequest
 
 object Spray extends ExtensionId[SprayExtension] with ExtensionIdProvider {
@@ -66,7 +67,14 @@ trait SprayNameGenerator {
 }
 
 class DefaultSprayNameGenerator extends SprayNameGenerator {
-  def generateRequestLevelApiSegmentName(request: HttpRequest): String = request.method.value + ": " + request.uri.path
+  def hostFromHeaders(request: HttpRequest): Option[String] = request.header[Host].map(_.host)
+
+  def generateRequestLevelApiSegmentName(request: HttpRequest): String = {
+    val uriAddress = request.uri.authority.host.address
+    if (uriAddress.equals("")) hostFromHeaders(request).getOrElse("unknown-host") else uriAddress
+  }
+
+  def generateHostLevelApiSegmentName(request: HttpRequest): String = hostFromHeaders(request).getOrElse("unknown-host")
+
   def generateTraceName(request: HttpRequest): String = request.method.value + ": " + request.uri.path
-  def generateHostLevelApiSegmentName(request: HttpRequest): String = request.uri.authority.host.address
 }
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 eb25412b..93a9cf55 100644
--- a/kamon-spray/src/main/scala/spray/can/server/ServerRequestInstrumentation.scala
+++ b/kamon-spray/src/main/scala/spray/can/server/ServerRequestInstrumentation.scala
@@ -90,11 +90,12 @@ class ServerRequestInstrumentation {
 
   def verifyTraceContextConsistency(incomingTraceContext: TraceContext, storedTraceContext: TraceContext, system: ActorSystem): Unit = {
     def publishWarning(text: String, system: ActorSystem): Unit =
-      system.eventStream.publish(Warning("", classOf[ServerRequestInstrumentation], text))
+      system.eventStream.publish(Warning("ServerRequestInstrumentation", 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
+    if (incomingTraceContext.nonEmpty) {
+      if (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)
   }