complete spray client instrumentation with experimental branch

1 files changed, 64 insertions, 18 deletions

diff --git a/kamon-core/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala b/kamon-core/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala
index 5117e7e7..2bd8643c 100644
--- a/kamon-core/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala
+++ b/kamon-core/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala
@@ -1,11 +1,14 @@
 package kamon.instrumentation
 
-import org.aspectj.lang.annotation.{DeclareMixin, After, Pointcut, Aspect}
+import org.aspectj.lang.annotation._
 import kamon.{TraceContext, Tracer}
-import kamon.trace.UowTracing.{Finish, Rename}
-import spray.http.HttpRequest
-import spray.can.server.{OpenRequest, OpenRequestComponent}
+import kamon.trace.UowTracing._
 import kamon.trace.context.TracingAwareContext
+import org.aspectj.lang.ProceedingJoinPoint
+import spray.http.HttpRequest
+import kamon.trace.UowTracing.Finish
+import kamon.trace.UowTracing.Rename
+import spray.http.HttpHeaders.Host
 
 //import spray.can.client.HttpHostConnector.RequestContext
 
@@ -13,35 +16,45 @@ trait ContextAware {
   def traceContext: Option[TraceContext]
 }
 
+trait TimedContextAware {
+  def timestamp: Long
+  def traceContext: Option[TraceContext]
+}
+
 @Aspect
 class SprayOpenRequestContextTracing {
   @DeclareMixin("spray.can.server.OpenRequestComponent.DefaultOpenRequest")
   def mixinContextAwareToOpenRequest: ContextAware = new ContextAware {
     val traceContext: Option[TraceContext] = Tracer.traceContext.value
   }
+
+  @DeclareMixin("spray.can.client.HttpHostConnector.RequestContext")
+  def mixinContextAwareToRequestContext: TimedContextAware = new TimedContextAware {
+    val timestamp: Long = System.nanoTime()
+    val traceContext: Option[TraceContext] = Tracer.traceContext.value
+  }
 }
 
 @Aspect
 class SprayServerInstrumentation {
 
-  @Pointcut("execution(spray.can.server.OpenRequestComponent$DefaultOpenRequest.new(..)) && this(openRequest) && args(enclosing, request, closeAfterResponseCompletion, timestamp)")
-  def openRequestInit(openRequest: OpenRequest, enclosing: OpenRequestComponent, request: HttpRequest, closeAfterResponseCompletion: Boolean, timestamp: Long): Unit = {}
+  @Pointcut("execution(spray.can.server.OpenRequestComponent$DefaultOpenRequest.new(..)) && this(openRequest) && args(*, request, *, *)")
+  def openRequestInit(openRequest: ContextAware, request: HttpRequest): Unit = {}
 
-  @After("openRequestInit(openRequest, enclosing, request, closeAfterResponseCompletion, timestamp)")
-  def afterInit(openRequest: OpenRequest, enclosing: OpenRequestComponent, request: HttpRequest, closeAfterResponseCompletion: Boolean, timestamp: Long): Unit = {
+  @After("openRequestInit(openRequest, request)")
+  def afterInit(openRequest: ContextAware, request: HttpRequest): Unit = {
     Tracer.start
-    val discard = openRequest.asInstanceOf[ContextAware].traceContext
+    openRequest.traceContext
 
     Tracer.context().map(_.tracer ! Rename(request.uri.path.toString()))
   }
 
   @Pointcut("execution(* spray.can.server.OpenRequestComponent$DefaultOpenRequest.handleResponseEndAndReturnNextOpenRequest(..)) && target(openRequest)")
-  def openRequestCreation(openRequest: OpenRequest): Unit = {}
+  def openRequestCreation(openRequest: ContextAware): Unit = {}
 
   @After("openRequestCreation(openRequest)")
-  def afterFinishingRequest(openRequest: OpenRequest): Unit = {
-    val original = openRequest.asInstanceOf[ContextAware].traceContext
-
+  def afterFinishingRequest(openRequest: ContextAware): Unit = {
+    val original = openRequest.traceContext
     Tracer.context().map(_.tracer ! Finish())
 
     if(Tracer.context() != original) {
@@ -49,13 +62,46 @@ class SprayServerInstrumentation {
     }
   }
 
-  @Pointcut("execution(spray.can.client.HttpHostConnector.RequestContext.new(..)) && this(ctx)")
-  def requestRecordInit(ctx: TracingAwareContext): Unit = {}
+  @Pointcut("execution(spray.can.client.HttpHostConnector.RequestContext.new(..)) && this(ctx) && args(request, *, *)")
+  def requestRecordInit(ctx: TimedContextAware, request: HttpRequest): Unit = {}
 
-  @After("requestRecordInit(ctx)")
-  def whenCreatedRequestRecord(ctx: TracingAwareContext): Unit = {
+  @After("requestRecordInit(ctx, request)")
+  def whenCreatedRequestRecord(ctx: TimedContextAware, request: HttpRequest): Unit = {
     // Necessary to force the initialization of TracingAwareRequestContext at the moment of creation.
-    ctx.traceContext
+    for{
+      tctx <- ctx.traceContext
+      host <- request.header[Host]
+    } tctx.tracer ! WebExternalStart(ctx.timestamp, host.host)
+  }
+
+
+
+  @Pointcut("execution(* spray.can.client.HttpHostConnectionSlot.dispatchToCommander(..)) && args(requestContext, message)")
+  def dispatchToCommander(requestContext: TimedContextAware, message: Any): Unit = {}
+
+  @Around("dispatchToCommander(requestContext, message)")
+  def aroundDispatchToCommander(pjp: ProceedingJoinPoint, requestContext: TimedContextAware, message: Any) = {
+    println("Completing the request with context: " + requestContext.traceContext)
+
+    Tracer.traceContext.withValue(requestContext.traceContext) {
+      requestContext.traceContext.map {
+        tctx => tctx.tracer ! WebExternalFinish(requestContext.timestamp)
+      }
+      pjp.proceed()
+    }
+
+  }
+
+
+  @Pointcut("execution(* spray.can.client.HttpHostConnector.RequestContext.copy(..)) && this(old)")
+  def copyingRequestContext(old: TimedContextAware): Unit = {}
+
+  @Around("copyingRequestContext(old)")
+  def aroundCopyingRequestContext(pjp: ProceedingJoinPoint, old: TimedContextAware) = {
+    println("Instrumenting the request context copy.")
+    Tracer.traceContext.withValue(old.traceContext) {
+      pjp.proceed()
+    }
   }
 }