diff options
Diffstat (limited to 'kamon-trace/src/main/scala/kamon/trace/instrumentation')
3 files changed, 83 insertions, 3 deletions
diff --git a/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingInstrumentation.scala b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingInstrumentation.scala new file mode 100644 index 00000000..77993cdd --- /dev/null +++ b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorLoggingInstrumentation.scala @@ -0,0 +1,31 @@ +package kamon.trace.instrumentation + +import org.aspectj.lang.annotation.{Around, Pointcut, DeclareMixin, Aspect} +import org.aspectj.lang.ProceedingJoinPoint +import org.slf4j.MDC +import kamon.trace.{TraceContext, ContextAware, Trace} + +@Aspect +class ActorLoggingInstrumentation { + + + @DeclareMixin("akka.event.Logging.LogEvent+") + def traceContextMixin: ContextAware = new ContextAware { + def traceContext: Option[TraceContext] = Trace.context() + } + + @Pointcut("execution(* akka.event.slf4j.Slf4jLogger.withMdc(..)) && args(logSource, logEvent, logStatement)") + def withMdcInvocation(logSource: String, logEvent: ContextAware, logStatement: () => _): Unit = {} + + @Around("withMdcInvocation(logSource, logEvent, logStatement)") + def putTraceContextInMDC(pjp: ProceedingJoinPoint, logSource: String, logEvent: ContextAware, logStatement: () => _): Unit = { + logEvent.traceContext match { + case Some(ctx) => + MDC.put("uow", ctx.uow) + pjp.proceed() + MDC.remove("uow") + + case None => pjp.proceed() + } + } +} diff --git a/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorRefTellInstrumentation.scala b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorRefTellInstrumentation.scala new file mode 100644 index 00000000..3caba77c --- /dev/null +++ b/kamon-trace/src/main/scala/kamon/trace/instrumentation/ActorRefTellInstrumentation.scala @@ -0,0 +1,49 @@ +package kamon.trace.instrumentation + +import org.aspectj.lang.annotation._ +import org.aspectj.lang.ProceedingJoinPoint +import akka.actor.{Props, ActorSystem, ActorRef} +import akka.dispatch.{Envelope, MessageDispatcher} +import com.codahale.metrics.Timer +import kamon.trace.{ContextAware, TraceContext, Trace} + +case class TraceableMessage(traceContext: Option[TraceContext], message: Any, timer: Timer.Context) +case class DefaultTracingAwareEnvelopeContext(traceContext: Option[TraceContext] = Trace.traceContext.value, timestamp: Long = System.nanoTime) extends ContextAware + +@Aspect +class ActorCellInvokeInstrumentation { + + @Pointcut("execution(akka.actor.ActorCell.new(..)) && args(system, ref, props, dispatcher, parent)") + def actorCellCreation(system: ActorSystem, ref: ActorRef, props: Props, dispatcher: MessageDispatcher, parent: ActorRef): Unit = {} + + @Pointcut("(execution(* akka.actor.ActorCell.invoke(*)) || execution(* akka.routing.RoutedActorCell.sendMessage(*))) && args(envelope)") + def invokingActorBehaviourAtActorCell(envelope: Envelope) = {} + + @Around("invokingActorBehaviourAtActorCell(envelope)") + def around(pjp: ProceedingJoinPoint, envelope: Envelope): Unit = { + //safe cast + val msgContext = envelope.asInstanceOf[ContextAware].traceContext + + Trace.traceContext.withValue(msgContext) { + pjp.proceed() + } + } +} + +@Aspect +class EnvelopeTracingContext { + + @DeclareMixin("akka.dispatch.Envelope") + def mixin: ContextAware = new ContextAware { + val traceContext: Option[TraceContext] = Trace.context() + } + + @Pointcut("execution(akka.dispatch.ContextAware.new(..)) && this(ctx)") + def requestRecordInit(ctx: ContextAware): Unit = {} + + @After("requestRecordInit(ctx)") + def whenCreatedRequestRecord(ctx: ContextAware): Unit = { + // Necessary to force the initialization of TracingAwareRequestContext at the moment of creation. + ctx.traceContext + } +} diff --git a/kamon-trace/src/main/scala/kamon/trace/instrumentation/RunnableTracing.scala b/kamon-trace/src/main/scala/kamon/trace/instrumentation/RunnableTracing.scala index 236fd4fc..3e5a7cce 100644 --- a/kamon-trace/src/main/scala/kamon/trace/instrumentation/RunnableTracing.scala +++ b/kamon-trace/src/main/scala/kamon/trace/instrumentation/RunnableTracing.scala @@ -2,7 +2,7 @@ package kamon.trace.instrumentation import org.aspectj.lang.annotation._ import org.aspectj.lang.ProceedingJoinPoint -import kamon.trace.TraceContext +import kamon.trace.{TraceContext, Trace} @Aspect class RunnableTracing { @@ -13,7 +13,7 @@ class RunnableTracing { */ @DeclareMixin("scala.concurrent.impl.CallbackRunnable || scala.concurrent.impl.Future.PromiseCompletingRunnable") def onCompleteCallbacksRunnable: TraceContextAwareRunnable = new TraceContextAwareRunnable { - val traceContext: Option[TraceContext] = Tracer.traceContext.value + val traceContext: Option[TraceContext] = Trace.traceContext.value } @@ -40,7 +40,7 @@ class RunnableTracing { def around(pjp: ProceedingJoinPoint, runnable: TraceContextAwareRunnable): Any = { import pjp._ - Tracer.traceContext.withValue(runnable.traceContext) { + Trace.traceContext.withValue(runnable.traceContext) { proceed() } } |