From df99b59fd05c5f5e6a4b48bb5e3485449a6d6eda Mon Sep 17 00:00:00 2001 From: Ivan Topolnak Date: Wed, 2 Oct 2013 13:49:34 -0300 Subject: wip --- .../src/main/scala/kamon/TraceContextSwap.scala | 8 ++- .../ActorRefTellInstrumentation.scala | 2 +- .../instrumentation/RunnableInstrumentation.scala | 60 +++++++++++++--------- .../instrumentation/ActorInstrumentationSpec.scala | 4 +- .../RunnableInstrumentationSpec.scala | 6 +-- 5 files changed, 49 insertions(+), 31 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()) } } diff --git a/kamon-core/src/test/scala/akka/instrumentation/ActorInstrumentationSpec.scala b/kamon-core/src/test/scala/akka/instrumentation/ActorInstrumentationSpec.scala index 454b4514..21be4a73 100644 --- a/kamon-core/src/test/scala/akka/instrumentation/ActorInstrumentationSpec.scala +++ b/kamon-core/src/test/scala/akka/instrumentation/ActorInstrumentationSpec.scala @@ -42,7 +42,7 @@ class ActorInstrumentationSpec extends TestKit(ActorSystem("ActorInstrumentation expectMsgAllOf(contexts: _*) } - "propagate with many asks" in { + /*"propagate with many asks" in { val echo = system.actorOf(Props[TraceContextEcho]) val iterations = 50000 implicit val timeout = Timeout(10 seconds) @@ -59,7 +59,7 @@ class ActorInstrumentationSpec extends TestKit(ActorSystem("ActorInstrumentation assert(iterations == allResults.collect { case Some(_) => 1 }.sum) - } + }*/ } } diff --git a/kamon-core/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala b/kamon-core/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala index 86bef0d8..789c7c77 100644 --- a/kamon-core/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala +++ b/kamon-core/src/test/scala/kamon/instrumentation/RunnableInstrumentationSpec.scala @@ -16,16 +16,16 @@ class RunnableInstrumentationSpec extends WordSpec with Matchers with ScalaFutur "a instrumented runnable" when { "created in a thread that does have a TraceContext" must { "preserve the TraceContext" which { - "should be available during the run method execution" in { new FutureWithContextFixture { + "should be available during the run method execution" in new FutureWithContextFixture { whenReady(futureWithContext) { result => result.value should equal(testContext) } - }} + } "should be available during the execution of onComplete callbacks" in { new FutureWithContextFixture { val onCompleteContext = Promise[TraceContext]() - Tracer.clear + futureWithContext.onComplete({ case _ => onCompleteContext.complete(Success(Tracer.context.get)) }) -- cgit v1.2.3