2 files changed, 20 insertions, 28 deletions

diff --git a/kamon-core/src/main/scala/akka/instrumentation/ActorMessagePassingTracing.scala b/kamon-core/src/main/scala/akka/instrumentation/ActorMessagePassingTracing.scala
index d43de311..98700974 100644
--- a/kamon-core/src/main/scala/akka/instrumentation/ActorMessagePassingTracing.scala
+++ b/kamon-core/src/main/scala/akka/instrumentation/ActorMessagePassingTracing.scala
@@ -19,10 +19,9 @@ import org.aspectj.lang.annotation._
 import org.aspectj.lang.ProceedingJoinPoint
 import akka.actor._
 import akka.dispatch.{ Envelope, MessageDispatcher }
-import kamon.trace.{ TraceContext, ContextAware, Trace }
+import kamon.trace._
 import kamon.metrics.{ ActorMetrics, Metrics }
 import kamon.Kamon
-import kamon.metrics.ActorMetrics.ActorMetricRecorder
 import kamon.trace.TraceContext
 import kamon.metrics.ActorMetrics.ActorMetricRecorder
 
@@ -48,15 +47,15 @@ class BehaviourInvokeTracing {
   @Around("invokingActorBehaviourAtActorCell(cell, envelope)")
   def aroundBehaviourInvoke(pjp: ProceedingJoinPoint, cell: ActorCell, envelope: Envelope): Unit = {
     val timestampBeforeProcessing = System.nanoTime()
-    val contextAndTimestamp = envelope.asInstanceOf[ContextAndTimestampAware]
+    val contextAndTimestamp = envelope.asInstanceOf[TraceContextAware]
 
-    Trace.withContext(contextAndTimestamp.traceContext) {
+    TraceRecorder.withContext(contextAndTimestamp.traceContext) {
       pjp.proceed()
     }
 
     actorMetrics.map { am ⇒
       am.processingTime.record(System.nanoTime() - timestampBeforeProcessing)
-      am.timeInMailbox.record(timestampBeforeProcessing - contextAndTimestamp.timestamp)
+      am.timeInMailbox.record(timestampBeforeProcessing - contextAndTimestamp.captureMark)
       am.mailboxSize.record(cell.numberOfMessages)
     }
   }
@@ -74,21 +73,14 @@ class BehaviourInvokeTracing {
 class EnvelopeTraceContextMixin {
 
   @DeclareMixin("akka.dispatch.Envelope")
-  def mixin: ContextAndTimestampAware = new ContextAndTimestampAware {
-    val traceContext: Option[TraceContext] = Trace.context()
-    val timestamp: Long = System.nanoTime()
-  }
+  def mixin: TraceContextAware = new TraceContextAware {}
 
   @Pointcut("execution(akka.dispatch.Envelope.new(..)) && this(ctx)")
-  def envelopeCreation(ctx: ContextAware): Unit = {}
+  def envelopeCreation(ctx: TraceContextAware): Unit = {}
 
   @After("envelopeCreation(ctx)")
-  def afterEnvelopeCreation(ctx: ContextAware): Unit = {
+  def afterEnvelopeCreation(ctx: TraceContextAware): Unit = {
     // Necessary to force the initialization of ContextAware at the moment of creation.
     ctx.traceContext
   }
-}
-
-trait ContextAndTimestampAware extends ContextAware {
-  def timestamp: Long
-}
+}
\ No newline at end of file
diff --git a/kamon-core/src/main/scala/akka/instrumentation/ActorSystemMessagePassingTracing.scala b/kamon-core/src/main/scala/akka/instrumentation/ActorSystemMessagePassingTracing.scala
index 7d26016e..d4f8f769 100644
--- a/kamon-core/src/main/scala/akka/instrumentation/ActorSystemMessagePassingTracing.scala
+++ b/kamon-core/src/main/scala/akka/instrumentation/ActorSystemMessagePassingTracing.scala
@@ -1,21 +1,21 @@
 package akka.instrumentation
 
 import org.aspectj.lang.annotation._
-import kamon.trace.{ Trace, ContextAware }
 import akka.dispatch.sysmsg.EarliestFirstSystemMessageList
 import org.aspectj.lang.ProceedingJoinPoint
+import kamon.trace.{TraceRecorder, TraceContextAware}
 
 @Aspect
 class SystemMessageTraceContextMixin {
 
   @DeclareMixin("akka.dispatch.sysmsg.SystemMessage+")
-  def mixin: ContextAware = ContextAware.default
+  def mixin: TraceContextAware = new TraceContextAware {}
 
   @Pointcut("execution(akka.dispatch.sysmsg.SystemMessage+.new(..)) && this(ctx)")
-  def envelopeCreation(ctx: ContextAware): Unit = {}
+  def envelopeCreation(ctx: TraceContextAware): Unit = {}
 
   @After("envelopeCreation(ctx)")
-  def afterEnvelopeCreation(ctx: ContextAware): Unit = {
+  def afterEnvelopeCreation(ctx: TraceContextAware): Unit = {
     // Necessary to force the initialization of ContextAware at the moment of creation.
     ctx.traceContext
   }
@@ -25,23 +25,23 @@ class SystemMessageTraceContextMixin {
 class RepointableActorRefTraceContextMixin {
 
   @DeclareMixin("akka.actor.RepointableActorRef")
-  def mixin: ContextAware = ContextAware.default
+  def mixin: TraceContextAware = new TraceContextAware {}
 
   @Pointcut("execution(akka.actor.RepointableActorRef.new(..)) && this(ctx)")
-  def envelopeCreation(ctx: ContextAware): Unit = {}
+  def envelopeCreation(ctx: TraceContextAware): Unit = {}
 
   @After("envelopeCreation(ctx)")
-  def afterEnvelopeCreation(ctx: ContextAware): Unit = {
+  def afterEnvelopeCreation(ctx: TraceContextAware): Unit = {
     // Necessary to force the initialization of ContextAware at the moment of creation.
     ctx.traceContext
   }
 
   @Pointcut("execution(* akka.actor.RepointableActorRef.point(..)) && this(repointableActorRef)")
-  def repointableActorRefCreation(repointableActorRef: ContextAware): Unit = {}
+  def repointableActorRefCreation(repointableActorRef: TraceContextAware): Unit = {}
 
   @Around("repointableActorRefCreation(repointableActorRef)")
-  def afterRepointableActorRefCreation(pjp: ProceedingJoinPoint, repointableActorRef: ContextAware): Any = {
-    Trace.withContext(repointableActorRef.traceContext) {
+  def afterRepointableActorRefCreation(pjp: ProceedingJoinPoint, repointableActorRef: TraceContextAware): Any = {
+    TraceRecorder.withContext(repointableActorRef.traceContext) {
       pjp.proceed()
     }
   }
@@ -57,8 +57,8 @@ class ActorSystemMessagePassingTracing {
   @Around("systemMessageProcessing(messages)")
   def aroundSystemMessageInvoke(pjp: ProceedingJoinPoint, messages: EarliestFirstSystemMessageList): Any = {
     if (messages.nonEmpty) {
-      val ctx = messages.head.asInstanceOf[ContextAware].traceContext
-      Trace.withContext(ctx)(pjp.proceed())
+      val ctx = messages.head.asInstanceOf[TraceContextAware].traceContext
+      TraceRecorder.withContext(ctx)(pjp.proceed())
 
     } else pjp.proceed()
   }