2 files changed, 28 insertions, 9 deletions
diff --git a/kamon-core/src/main/scala/kamon/instrumentation/ActorRefTellInstrumentation.scala b/kamon-core/src/main/scala/kamon/instrumentation/ActorRefTellInstrumentation.scala
index fdd7b696..d92d7f6c 100644
--- a/kamon-core/src/main/scala/kamon/instrumentation/ActorRefTellInstrumentation.scala
+++ b/kamon-core/src/main/scala/kamon/instrumentation/ActorRefTellInstrumentation.scala
@@ -34,6 +34,7 @@ class ActorRefTellInstrumentation {
 @Aspect("""perthis(actorCellCreation(akka.actor.ActorSystem, akka.actor.ActorRef, akka.actor.Props, akka.dispatch.MessageDispatcher, akka.actor.ActorRef))""")
 class ActorCellInvokeInstrumentation {
   var instrumentation =  ActorReceiveInvokeInstrumentation.noopPreReceive
+  var self: ActorRef = _
 
   // AKKA 2.2 introduces the dispatcher parameter. Maybe we could provide a dual pointcut.
   @Pointcut("execution(akka.actor.ActorCell.new(..)) && args(system, ref, props, dispatcher, parent)")
@@ -42,6 +43,7 @@ class ActorCellInvokeInstrumentation {
   @After("actorCellCreation(system, ref, props, dispatcher, parent)")
   def registerMetricsInRegistry(system: ActorSystem, ref: ActorRef, props: Props, dispatcher: MessageDispatcher, parent: ActorRef): Unit = {
     instrumentation = kamon.Instrument.instrumentation.receiveInvokeInstrumentation(system, ref, props, dispatcher, parent)
+    self = ref
   }
 
 
@@ -53,7 +55,7 @@ class ActorCellInvokeInstrumentation {
     import ProceedingJoinPointPimp._
 
     val (originalEnvelope, ctx) = instrumentation.preReceive(envelope)
-    //println("Test")
+    //println(s"====>[$ctx] ## [${originalEnvelope.sender}] => [$self] --- ${originalEnvelope.message}")
     ctx match {
       case Some(c) => {
         //MDC.put("uow", c.userContext.get.asInstanceOf[String])
diff --git a/kamon-core/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala b/kamon-core/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala
index f8ab709e..9422a9f7 100644
--- a/kamon-core/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala
+++ b/kamon-core/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala
@@ -4,26 +4,39 @@ import org.aspectj.lang.annotation.{DeclareMixin, After, Pointcut, Aspect}
 import kamon.{TraceContext, Tracer}
 import kamon.trace.UowTracing.{WebExternal, Finish, Rename}
 import spray.http.HttpRequest
-import spray.can.server.OpenRequestComponent
+import spray.can.server.{OpenRequest, OpenRequestComponent}
 import spray.can.client.HttpHostConnector.RequestContext
 import spray.http.HttpHeaders.Host
 
+trait ContextAware {
+  def traceContext: Option[TraceContext]
+}
+
+@Aspect
+class SprayOpenRequestContextTracing {
+  @DeclareMixin("spray.can.server.OpenRequestComponent.DefaultOpenRequest")
+  def mixinContextAwareToOpenRequest: ContextAware = new ContextAware {
+    val traceContext: Option[TraceContext] = Tracer.context()
+  }
+}
+
 @Aspect
 class SprayServerInstrumentation {
 
-  @Pointcut("execution(spray.can.server.OpenRequestComponent$DefaultOpenRequest.new(..)) && args(enclosing, request, closeAfterResponseCompletion, timestamp)")
-  def openRequestInit(enclosing: OpenRequestComponent, request: HttpRequest, closeAfterResponseCompletion: Boolean, timestamp: Long): Unit = {}
+  @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 = {}
 
-  @After("openRequestInit(enclosing, request, closeAfterResponseCompletion, timestamp)")
-  def afterInit(enclosing: OpenRequestComponent, request: HttpRequest, closeAfterResponseCompletion: Boolean, timestamp: Long): Unit = {
+  @After("openRequestInit(openRequest, enclosing, request, closeAfterResponseCompletion, timestamp)")
+  def afterInit(openRequest: OpenRequest, enclosing: OpenRequestComponent, request: HttpRequest, closeAfterResponseCompletion: Boolean, timestamp: Long): Unit = {
   //@After("openRequestInit()")
   //def afterInit(): Unit = {
     Tracer.start
-    //println("Created the context: " + Tracer.context() + " for the transaction: " + request.uri.path.toString())
+    //openRequest.traceContext
+    println("Created the context: " + Tracer.context() + " for the transaction: " + request)
     Tracer.context().map(_.entries ! Rename(request.uri.path.toString()))
   }
 
-  @Pointcut("execution(* spray.can.server.OpenRequest.handleResponseEndAndReturnNextOpenRequest(..))")
+  @Pointcut("execution(* spray.can.server.OpenRequestComponent$DefaultOpenRequest.handleResponseEndAndReturnNextOpenRequest(..))")
   def openRequestCreation(): Unit = {}
 
   @After("openRequestCreation()")
@@ -31,6 +44,10 @@ class SprayServerInstrumentation {
     println("Finishing a request: " + Tracer.context())
 
     Tracer.context().map(_.entries ! Finish())
+
+    if(Tracer.context().isEmpty) {
+      println("WOOOOOPAAAAAAAAA")
+    }
   }
 
 
@@ -55,7 +72,7 @@ class SprayServerInstrumentation {
 
   @After("requestRecordInit2(ctx, msg)")
   def whenCreatedRequestRecord2(ctx: TracingAwareRequestContext, msg: Any): Unit = {
-    println("=======> Spent in WEB External: " + (System.nanoTime() - ctx.timestamp))
+    //println("=======> Spent in WEB External: " + (System.nanoTime() - ctx.timestamp))
 
     // TODO: REMOVE THIS:
     val request = (ctx.asInstanceOf[RequestContext]).request