diff options
Diffstat (limited to 'kamon-core/src/main/scala')
3 files changed, 44 insertions, 26 deletions
diff --git a/kamon-core/src/main/scala/kamon/TraceContextSwap.scala b/kamon-core/src/main/scala/kamon/TraceContextSwap.scala index c25e63d1..470b2f34 100644 --- a/kamon-core/src/main/scala/kamon/TraceContextSwap.scala +++ b/kamon-core/src/main/scala/kamon/TraceContextSwap.scala @@ -10,19 +10,23 @@ trait TraceContextSwap { def withContext[A](ctx: Option[TraceContext], body: => A): A = withContext(ctx, body, body) def withContext[A](ctx: Option[TraceContext], primary: => A, fallback: => A): A = { - ctx match { + + val previous = Tracer.context() + val r = ctx match { case Some(context) => { //MDC.put("uow", context.userContext.get.asInstanceOf[String]) Tracer.set(context) val bodyResult = primary - Tracer.clear + //Tracer.clear //MDC.remove("uow") bodyResult } case None => fallback } + previous.map(ctx => Tracer.set(ctx)) + r } } diff --git a/kamon-core/src/main/scala/kamon/instrumentation/ActorRefTellInstrumentation.scala b/kamon-core/src/main/scala/kamon/instrumentation/ActorRefTellInstrumentation.scala index 43841165..6126d642 100644 --- a/kamon-core/src/main/scala/kamon/instrumentation/ActorRefTellInstrumentation.scala +++ b/kamon-core/src/main/scala/kamon/instrumentation/ActorRefTellInstrumentation.scala @@ -31,7 +31,7 @@ class ActorCellInvokeInstrumentation { Tracer.clear } case None => - assert(Tracer.context() == None) + //assert(Tracer.context() == None) pjp.proceed() } Tracer.clear diff --git a/kamon-core/src/main/scala/kamon/instrumentation/RunnableInstrumentation.scala b/kamon-core/src/main/scala/kamon/instrumentation/RunnableInstrumentation.scala index bff118d9..456917e0 100644 --- a/kamon-core/src/main/scala/kamon/instrumentation/RunnableInstrumentation.scala +++ b/kamon-core/src/main/scala/kamon/instrumentation/RunnableInstrumentation.scala @@ -8,10 +8,12 @@ import scala.Some /** * Marker interface, just to make sure we don't instrument all the Runnables in the classpath. */ -trait TraceContextAwareRunnable extends Runnable {} +trait TraceContextAwareRunnable { + def traceContext: Option[TraceContext] +} -@Aspect("perthis(instrumentedRunnableCreation())") +@Aspect class RunnableInstrumentation { /** @@ -19,35 +21,32 @@ class RunnableInstrumentation { * while their run method is executed. */ @DeclareMixin("scala.concurrent.impl.CallbackRunnable || scala.concurrent.impl.Future.PromiseCompletingRunnable") - def onCompleteCallbacksRunnable: TraceContextAwareRunnable = null + def onCompleteCallbacksRunnable: TraceContextAwareRunnable = new TraceContextAwareRunnable { + val traceContext: Option[TraceContext] = Tracer.context() + } /** * Pointcuts */ - @Pointcut("execution(kamon.instrumentation.TraceContextAwareRunnable+.new(..))") - def instrumentedRunnableCreation(): Unit = {} - - @Pointcut("execution(* kamon.instrumentation.TraceContextAwareRunnable.run())") - def runnableExecution() = {} + @Pointcut("execution(kamon.instrumentation.TraceContextAwareRunnable+.new(..)) && this(runnable)") + def instrumentedRunnableCreation(runnable: TraceContextAwareRunnable): Unit = {} + @Pointcut("execution(* kamon.instrumentation.TraceContextAwareRunnable+.run()) && this(runnable)") + def runnableExecution(runnable: TraceContextAwareRunnable) = {} - /** - * Aspect members - */ - - private var traceContext = Tracer.context - /** - * Advices - */ import kamon.TraceContextSwap.withContext - @Before("instrumentedRunnableCreation()") - def beforeCreation = { - traceContext = Tracer.context + @After("instrumentedRunnableCreation(runnable)") + def beforeCreation(runnable: TraceContextAwareRunnable) = { + val x = runnable.traceContext + /*if(runnable.traceContext.isEmpty) + println("WTFWI from: " + (new Throwable).getStackTraceString) + else + println("NOWTF: " + (new Throwable).getStackTraceString)*/ /* if(traceContext.isEmpty) println("NO TRACE CONTEXT FOR RUNNABLE at: [[[%s]]]", (new Throwable).getStackTraceString)//println((new Throwable).getStackTraceString) else @@ -55,12 +54,27 @@ class RunnableInstrumentation { } - @Around("runnableExecution()") - def around(pjp: ProceedingJoinPoint) = { + @Around("runnableExecution(runnable)") + def around(pjp: ProceedingJoinPoint, runnable: TraceContextAwareRunnable) = { import pjp._ - if(traceContext.isEmpty) + + /*println("EXECUTING") + if(runnable.traceContext.isEmpty) println("NOMONEY") - withContext(traceContext, proceed()) + + runnable.traceContext match { + case Some(context) => { + //MDC.put("uow", context.userContext.get.asInstanceOf[String]) + Tracer.set(context) + val bodyResult = proceed() + Tracer.clear + //MDC.remove("uow") + + bodyResult + } + case None => proceed() + }*/ + withContext(runnable.traceContext, proceed()) } } |