diff options
Diffstat (limited to 'kamon-spray/src/main/scala/spray/can/server/ServerRequestInstrumentation.scala')
-rw-r--r-- | kamon-spray/src/main/scala/spray/can/server/ServerRequestInstrumentation.scala | 47 |
1 files changed, 21 insertions, 26 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..74d98564 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 @@ -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 = |