1 files changed, 13 insertions, 46 deletions

diff --git a/kamon-spray/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala b/kamon-spray/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala
index 743769e2..08cb53ff 100644
--- a/kamon-spray/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala
+++ b/kamon-spray/src/main/scala/kamon/instrumentation/SprayServerInstrumentation.scala
@@ -1,61 +1,28 @@
-package kamon.instrumentation
+package spray.can.server
 
 import org.aspectj.lang.annotation._
 import org.aspectj.lang.ProceedingJoinPoint
 import spray.http.HttpRequest
 import spray.http.HttpHeaders.Host
+import kamon.trace.{TraceContext, Trace, ContextAware, TimedContextAware}
 
 //import spray.can.client.HttpHostConnector.RequestContext
 
-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
+    val traceContext: Option[TraceContext] = Trace.context()
   }
 }
 
 @Aspect
 class SprayServerInstrumentation {
 
-  @Pointcut("execution(spray.can.server.OpenRequestComponent$DefaultOpenRequest.new(..)) && this(openRequest) && args(*, request, *, *)")
-  def openRequestInit(openRequest: ContextAware, request: HttpRequest): Unit = {}
-
-  @After("openRequestInit(openRequest, request)")
-  def afterInit(openRequest: ContextAware, request: HttpRequest): Unit = {
-    Tracer.start
-    openRequest.traceContext
-
-    Tracer.context().map(_.tracer ! Rename(request.uri.path.toString()))
-  }
-
-  @Pointcut("execution(* spray.can.server.OpenRequestComponent$DefaultOpenRequest.handleResponseEndAndReturnNextOpenRequest(..)) && target(openRequest)")
-  def openRequestCreation(openRequest: ContextAware): Unit = {}
-
-  @After("openRequestCreation(openRequest)")
-  def afterFinishingRequest(openRequest: ContextAware): Unit = {
-    val original = openRequest.traceContext
-    Tracer.context().map(_.tracer ! Finish())
-
-    if(Tracer.context() != original) {
-      println(s"OMG DIFFERENT Original: [${original}] - Came in: [${Tracer.context}]")
-    }
-  }
 
   @Pointcut("execution(spray.can.client.HttpHostConnector.RequestContext.new(..)) && this(ctx) && args(request, *, *, *)")
   def requestRecordInit(ctx: TimedContextAware, request: HttpRequest): Unit = {}
@@ -63,10 +30,10 @@ class SprayServerInstrumentation {
   @After("requestRecordInit(ctx, request)")
   def whenCreatedRequestRecord(ctx: TimedContextAware, request: HttpRequest): Unit = {
     // Necessary to force the initialization of TracingAwareRequestContext at the moment of creation.
-    for{
+    /*for{
       tctx <- ctx.traceContext
       host <- request.header[Host]
-    } tctx.tracer ! WebExternalStart(ctx.timestamp, host.host)
+    } tctx.tracer ! WebExternalStart(ctx.timestamp, host.host)*/
   }
 
 
@@ -78,12 +45,12 @@ class SprayServerInstrumentation {
   def aroundDispatchToCommander(pjp: ProceedingJoinPoint, requestContext: TimedContextAware, message: Any) = {
     println("Completing the request with context: " + requestContext.traceContext)
 
-    Tracer.traceContext.withValue(requestContext.traceContext) {
+    /*Tracer.context.withValue(requestContext.traceContext) {
       requestContext.traceContext.map {
-        tctx => tctx.tracer ! WebExternalFinish(requestContext.timestamp)
+        tctx => //tctx.tracer ! WebExternalFinish(requestContext.timestamp)
       }
       pjp.proceed()
-    }
+    }*/
 
   }
 
@@ -94,17 +61,17 @@ class SprayServerInstrumentation {
   @Around("copyingRequestContext(old)")
   def aroundCopyingRequestContext(pjp: ProceedingJoinPoint, old: TimedContextAware) = {
     println("Instrumenting the request context copy.")
-    Tracer.traceContext.withValue(old.traceContext) {
+    /*Tracer.traceContext.withValue(old.traceContext) {
       pjp.proceed()
-    }
+    }*/
   }
 }
 
-case class DefaultTracingAwareRequestContext(traceContext: Option[TraceContext] = Tracer.context(), timestamp: Long = System.nanoTime) extends TracingAwareContext
-
 @Aspect
 class SprayRequestContextTracing {
 
   @DeclareMixin("spray.can.client.HttpHostConnector.RequestContext")
-  def mixin: TracingAwareContext = DefaultTracingAwareRequestContext()
+  def mixin: ContextAware = new ContextAware {
+    val traceContext: Option[TraceContext] = Trace.context()
+  }
 }
\ No newline at end of file