+ play, spray, newrelic: store in TraceLocal useful data to diagnose errors and closes #6

6 files changed, 78 insertions, 20 deletions

diff --git a/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala b/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
index a7296c31..c5fb100c 100644
--- a/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TraceLocal.scala
@@ -38,6 +38,10 @@ object TraceLocal {
     def apply(mdcKey: String): AvailableToMdc = fromKey(mdcKey)
   }
 
+  case class HttpContext(agent: String, uri: String, xforwarded: String)
+
+  object HttpContextKey extends TraceLocal.TraceLocalKey { type ValueType = HttpContext }
+
   def store(key: TraceLocalKey)(value: key.ValueType): Unit = TraceRecorder.currentContext match {
     case ctx: DefaultTraceContext ⇒ ctx.traceLocalStorage.store(key)(value)
     case EmptyTraceContext        ⇒ // Can't store in the empty context.
diff --git a/kamon-newrelic/src/main/scala/kamon/newrelic/NewRelicErrorLogger.scala b/kamon-newrelic/src/main/scala/kamon/newrelic/NewRelicErrorLogger.scala
index 08fdc8c4..d6920ad2 100644
--- a/kamon-newrelic/src/main/scala/kamon/newrelic/NewRelicErrorLogger.scala
+++ b/kamon-newrelic/src/main/scala/kamon/newrelic/NewRelicErrorLogger.scala
@@ -21,7 +21,8 @@ import java.util
 import akka.actor.{ Actor, ActorLogging }
 import akka.event.Logging.{ Error, InitializeLogger, LoggerInitialized }
 import com.newrelic.api.agent.{ NewRelic ⇒ NR }
-import kamon.trace.{ TraceRecorder, TraceContextAware }
+import kamon.trace.TraceLocal.{ HttpContext, HttpContextKey }
+import kamon.trace.{ TraceLocal, TraceRecorder, TraceContextAware }
 
 trait CustomParamsSupport {
   this: NewRelicErrorLogger ⇒
@@ -41,8 +42,16 @@ class NewRelicErrorLogger extends Actor with ActorLogging with CustomParamsSuppo
   def notifyError(error: Error): Unit = runInFakeTransaction {
     val params = new util.HashMap[String, String]()
     val ctx = error.asInstanceOf[TraceContextAware].traceContext
+    val httpContext = TraceLocal.retrieve(HttpContextKey)
 
     params put ("TraceToken", ctx.token)
+
+    httpContext.map { httpCtx ⇒
+      params put ("User-Agent", httpCtx.agent)
+      params put ("X-Forwarded-For", httpCtx.xforwarded)
+      params put ("Request-URI", httpCtx.uri)
+    }
+
     customParams foreach { case (k, v) ⇒ params.put(k, v) }
 
     if (error.cause == Error.NoCause) NR.noticeError(error.message.toString, params)
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 989ef43e..527e0a9e 100644
--- a/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
+++ b/kamon-play/src/main/scala/kamon/play/instrumentation/RequestInstrumentation.scala
@@ -17,7 +17,8 @@ package kamon.play.instrumentation
 
 import kamon.Kamon
 import kamon.play.{ Play, PlayExtension }
-import kamon.trace.{ TraceContextAware, TraceRecorder }
+import kamon.trace.TraceLocal.{ HttpContextKey, HttpContext }
+import kamon.trace.{ TraceLocal, TraceContextAware, TraceRecorder }
 import org.aspectj.lang.ProceedingJoinPoint
 import org.aspectj.lang.annotation._
 import play.api.Routes
@@ -28,7 +29,7 @@ import play.libs.Akka
 @Aspect
 class RequestInstrumentation {
 
-  import kamon.play.instrumentation.RequestInstrumentation._
+  import RequestInstrumentation._
 
   @DeclareMixin("play.api.mvc.RequestHeader+")
   def mixinContextAwareNewRequest: TraceContextAware = TraceContextAware.default
@@ -58,23 +59,22 @@ class RequestInstrumentation {
       val executor = Kamon(Play)(Akka.system()).defaultDispatcher
 
       def onResult(result: Result): Result = {
-
         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
+          if (playExtension.includeTraceToken) result.withHeaders(playExtension.traceTokenHeaderName -> ctx.token)
+          else result
 
-        } getOrElse (result)
+        } getOrElse result
       }
+      //store in TraceLocal useful data to diagnose errors
+      storeDiagnosticData(requestHeader)
 
       //override the current trace name
-      normaliseTraceName(requestHeader).map(TraceRecorder.rename(_))
+      normaliseTraceName(requestHeader).map(TraceRecorder.rename)
 
       // Invoke the action
       next(requestHeader).map(onResult)(executor)
@@ -87,8 +87,15 @@ class RequestInstrumentation {
     recordHttpServerMetrics(InternalServerError.header, ctx.name, Kamon(Play)(system))
   }
 
-  private def recordHttpServerMetrics(header: ResponseHeader, traceName: String, playExtension: PlayExtension): Unit =
+  def recordHttpServerMetrics(header: ResponseHeader, traceName: String, playExtension: PlayExtension): Unit =
     playExtension.httpServerMetrics.recordResponse(traceName, header.status.toString)
+
+  def storeDiagnosticData(request: RequestHeader): Unit = {
+    val agent = request.headers.get(UserAgent).getOrElse(Unknown)
+    val forwarded = request.headers.get(XForwardedFor).getOrElse(Unknown)
+
+    TraceLocal.store(HttpContextKey)(HttpContext(agent, request.uri, forwarded))
+  }
 }
 
 object RequestInstrumentation {
@@ -96,6 +103,10 @@ object RequestInstrumentation {
   import java.util.Locale
   import scala.collection.concurrent.TrieMap
 
+  val UserAgent = "User-Agent"
+  val XForwardedFor = "X-Forwarded-For"
+  val Unknown = "unknown"
+
   private val cache = TrieMap.empty[String, String]
 
   def normaliseTraceName(requestHeader: RequestHeader): Option[String] = requestHeader.tags.get(Routes.ROUTE_VERB).map({ verb ⇒
diff --git a/kamon-play/src/test/scala/kamon/play/RequestInstrumentationSpec.scala b/kamon-play/src/test/scala/kamon/play/RequestInstrumentationSpec.scala
index 3feb6246..3c3ce9f2 100644
--- a/kamon-play/src/test/scala/kamon/play/RequestInstrumentationSpec.scala
+++ b/kamon-play/src/test/scala/kamon/play/RequestInstrumentationSpec.scala
@@ -19,7 +19,9 @@ import kamon.Kamon
 import kamon.http.HttpServerMetrics
 import kamon.metric.{ CollectionContext, Metrics, TraceMetrics }
 import kamon.play.action.TraceName
+import kamon.trace.TraceLocal.HttpContextKey
 import kamon.trace.{ TraceLocal, TraceRecorder }
+import org.scalatest.Matchers
 import org.scalatestplus.play._
 import play.api.DefaultGlobal
 import play.api.http.Writeable
@@ -127,7 +129,7 @@ class RequestInstrumentationSpec extends PlaySpec with OneServerPerSuite {
     }
 
     "response to the getRouted Action and normalise the current TraceContext name" in {
-      Await.result(WS.url("http://localhost:19001/getRouted").get, 10 seconds)
+      Await.result(WS.url("http://localhost:19001/getRouted").get(), 10 seconds)
       Kamon(Metrics)(Akka.system()).storage.get(TraceMetrics("getRouted.get")) must not be (empty)
     }
 
@@ -137,10 +139,20 @@ class RequestInstrumentationSpec extends PlaySpec with OneServerPerSuite {
     }
 
     "response to the showRouted Action and normalise the current TraceContext name" in {
-      Await.result(WS.url("http://localhost:19001/showRouted/2").get, 10 seconds)
+      Await.result(WS.url("http://localhost:19001/showRouted/2").get(), 10 seconds)
       Kamon(Metrics)(Akka.system()).storage.get(TraceMetrics("show.some.id.get")) must not be (empty)
     }
 
+    "include HttpContext information for help to diagnose possible errors" in {
+      Await.result(WS.url("http://localhost:19001/getRouted").get(), 10 seconds)
+      route(FakeRequest(GET, "/default").withHeaders("User-Agent" -> "Fake-Agent"))
+
+      val httpCtx = TraceLocal.retrieve(HttpContextKey).get
+      httpCtx.agent must be("Fake-Agent")
+      httpCtx.uri must be("/default")
+      httpCtx.xforwarded must be("unknown")
+    }
+
     "record http server metrics for all processed requests" in {
       val collectionContext = CollectionContext(100)
       Kamon(Metrics)(Akka.system()).register(HttpServerMetrics, HttpServerMetrics.Factory).get.collect(collectionContext)
@@ -180,7 +192,10 @@ class RequestInstrumentationSpec extends PlaySpec with OneServerPerSuite {
         TraceLocal.store(TraceLocalKey)(header.headers.get(traceLocalStorageKey).getOrElse("unknown"))
 
         next(header).map {
-          result ⇒ result.withHeaders((traceLocalStorageKey -> TraceLocal.retrieve(TraceLocalKey).get))
+          result ⇒
+            {
+              result.withHeaders(traceLocalStorageKey -> TraceLocal.retrieve(TraceLocalKey).get)
+            }
         }
       }
     }
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 93a9cf55..2bdd96e7 100644
--- a/kamon-spray/src/main/scala/spray/can/server/ServerRequestInstrumentation.scala
+++ b/kamon-spray/src/main/scala/spray/can/server/ServerRequestInstrumentation.scala
@@ -15,6 +15,7 @@
  * ========================================================== */
 package spray.can.server
 
+import kamon.trace.TraceLocal.{ HttpContext, HttpContextKey }
 import org.aspectj.lang.annotation._
 import kamon.trace._
 import akka.actor.ActorSystem
@@ -28,6 +29,8 @@ import spray.http.HttpHeaders.RawHeader
 @Aspect
 class ServerRequestInstrumentation {
 
+  import ServerRequestInstrumentation._
+
   @DeclareMixin("spray.can.server.OpenRequestComponent.DefaultOpenRequest")
   def mixinContextAwareToOpenRequest: TraceContextAware = TraceContextAware.default
 
@@ -83,7 +86,12 @@ class ServerRequestInstrumentation {
       } else pjp.proceed
 
       TraceRecorder.finish()
+
       recordHttpServerMetrics(response, incomingContext.name, sprayExtension)
+
+      //store in TraceLocal useful data to diagnose errors
+      storeDiagnosticData(openRequest)
+
       proceedResult
     }
   }
@@ -110,4 +118,19 @@ class ServerRequestInstrumentation {
       case response: HttpResponse ⇒ response.withHeaders(response.headers ::: RawHeader(traceTokenHeaderName, token) :: Nil)
       case other                  ⇒ other
     }
+
+  def storeDiagnosticData(currentContext: TraceContextAware): Unit = {
+    val request = currentContext.asInstanceOf[OpenRequest].request
+    val headers = request.headers.map(header ⇒ header.name -> header.value).toMap
+    val agent = headers.getOrElse(UserAgent, Unknown)
+    val forwarded = headers.getOrElse(XForwardedFor, Unknown)
+
+    TraceLocal.store(HttpContextKey)(HttpContext(agent, request.uri.toRelative.toString(), forwarded))
+  }
+}
+
+object ServerRequestInstrumentation {
+  val UserAgent = "User-Agent"
+  val XForwardedFor = "X-Forwarded-For"
+  val Unknown = "unknown"
 }
diff --git a/kamon-spray/src/test/scala/kamon/spray/SprayServerTracingSpec.scala b/kamon-spray/src/test/scala/kamon/spray/SprayServerTracingSpec.scala
index 48253b1d..30d42eea 100644
--- a/kamon-spray/src/test/scala/kamon/spray/SprayServerTracingSpec.scala
+++ b/kamon-spray/src/test/scala/kamon/spray/SprayServerTracingSpec.scala
@@ -24,11 +24,7 @@ import kamon.Kamon
 import org.scalatest.concurrent.{ PatienceConfiguration, ScalaFutures }
 import spray.http.HttpHeaders.RawHeader
 import spray.http.{ HttpResponse, HttpRequest }
-import kamon.metric.{ TraceMetrics, Metrics }
-import kamon.metric.Subscriptions.TickMetricSnapshot
 import com.typesafe.config.ConfigFactory
-import kamon.metric.TraceMetrics.ElapsedTime
-import kamon.metric.instrument.Histogram
 
 class SprayServerTracingSpec extends TestKitBase with WordSpecLike with Matchers with RequestBuilding
     with ScalaFutures with PatienceConfiguration with TestServer {
@@ -81,7 +77,7 @@ class SprayServerTracingSpec extends TestKitBase with WordSpecLike with Matchers
       server.reply(HttpResponse(entity = "ok"))
       val response = client.expectMsgType[HttpResponse]
 
-      response.headers.filter(_.name == Kamon(Spray).traceTokenHeaderName).size should be(1)
+      response.headers.count(_.name == Kamon(Spray).traceTokenHeaderName) should be(1)
 
     }
 
@@ -96,7 +92,7 @@ class SprayServerTracingSpec extends TestKitBase with WordSpecLike with Matchers
       server.reply(HttpResponse(entity = "ok"))
       val response = client.expectMsgType[HttpResponse]
 
-      response.headers should not contain (RawHeader(Kamon(Spray).traceTokenHeaderName, "propagation-disabled"))
+      response.headers should not contain RawHeader(Kamon(Spray).traceTokenHeaderName, "propagation-disabled")
     }
   }