WIP - first functional implementation of TraceContext

1 files changed, 8 insertions, 6 deletions

diff --git a/src/main/scala/akka/instrumentation/ActorRefTellInstrumentation.scala b/src/main/scala/akka/instrumentation/ActorRefTellInstrumentation.scala
index 9e816d11..00e4e066 100644
--- a/src/main/scala/akka/instrumentation/ActorRefTellInstrumentation.scala
+++ b/src/main/scala/akka/instrumentation/ActorRefTellInstrumentation.scala
@@ -3,10 +3,12 @@ package akka.instrumentation
 import org.aspectj.lang.annotation.{Around, Pointcut, Aspect}
 import org.aspectj.lang.ProceedingJoinPoint
 import akka.actor.{ActorRef}
-import kamon.TraceContext
-import kamon.actor.TraceableMessage
+import kamon.{Kamon, TraceContext}
 import akka.dispatch.Envelope
 
+case class TraceableMessage(traceContext: TraceContext, message: Any)
+
+
 @Aspect
 class ActorRefTellInstrumentation {
   println("Created ActorAspect")
@@ -18,9 +20,9 @@ class ActorRefTellInstrumentation {
   def around(pjp: ProceedingJoinPoint, message: Any, sender: ActorRef): Unit  = {
     import pjp._
 
-    TraceContext.current match {
+    Kamon.context match {
       case Some(ctx) => {
-        val traceableMessage = TraceableMessage(ctx.fork, message)
+        val traceableMessage = TraceableMessage(ctx, message)
         proceed(getArgs.updated(0, traceableMessage))
       }
       case None => proceed
@@ -42,12 +44,12 @@ class ActorCellInvokeInstrumentation {
 
     envelope match {
       case Envelope(TraceableMessage(ctx, msg), sender) => {
-        TraceContext.set(ctx)
+        Kamon.set(ctx)
 
         val originalEnvelope = envelope.copy(message = msg)
         proceed(getArgs.updated(0, originalEnvelope))
 
-        TraceContext.clear
+        Kamon.clear
       }
       case _ => proceed
     }