From 96e89d0908c479360d4a915c916e341e148a91c7 Mon Sep 17 00:00:00 2001
From: Ivan Topolnak <ivantopo@gmail.com>
Date: Wed, 25 Sep 2013 19:45:54 -0300
Subject: wip

---
 .../SprayServerInstrumentation.scala               | 33 ++++++++++++++++------
 1 file changed, 25 insertions(+), 8 deletions(-)

(limited to 'kamon-core/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala')

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
-- 
cgit v1.2.3