From 148827486f116c4196888022f04ad053f4fb6e99 Mon Sep 17 00:00:00 2001
From: Ivan Topolnak <ivantopo@gmail.com>
Date: Thu, 30 May 2013 18:32:41 -0300
Subject: WIP - first functional implementation of TraceContext

---
 .../instrumentation/RunnableInstrumentation.scala  | 55 ++++++++++++++++++++++
 1 file changed, 55 insertions(+)
 create mode 100644 src/main/scala/kamon/instrumentation/RunnableInstrumentation.scala

(limited to 'src/main/scala/kamon/instrumentation/RunnableInstrumentation.scala')

diff --git a/src/main/scala/kamon/instrumentation/RunnableInstrumentation.scala b/src/main/scala/kamon/instrumentation/RunnableInstrumentation.scala
new file mode 100644
index 00000000..ef908625
--- /dev/null
+++ b/src/main/scala/kamon/instrumentation/RunnableInstrumentation.scala
@@ -0,0 +1,55 @@
+package kamon.instrumentation
+
+import org.aspectj.lang.annotation._
+import kamon.{Kamon, TraceContext}
+import org.aspectj.lang.ProceedingJoinPoint
+import scala.Some
+
+/**
+ *  Marker interface, just to make sure we don't instrument all the Runnables in the classpath.
+ */
+trait TraceContextAwareRunnable extends Runnable {}
+
+
+@Aspect("perthis(instrumentedRunnableCreation())")
+class RunnableInstrumentation {
+
+  /**
+   *  These are the Runnables that need to be instrumented and make the TraceContext available
+   *  while their run method is executed.
+   */
+  @DeclareMixin("scala.concurrent.impl.CallbackRunnable || scala.concurrent.impl.Future.PromiseCompletingRunnable")
+  def onCompleteCallbacksRunnable: TraceContextAwareRunnable = null
+
+
+  /**
+   *  Pointcuts
+   */
+
+  @Pointcut("execution(kamon.instrumentation.TraceContextAwareRunnable+.new(..))")
+  def instrumentedRunnableCreation(): Unit = {}
+
+  @Pointcut("execution(* kamon.instrumentation.TraceContextAwareRunnable.run())")
+  def runnableExecution() = {}
+
+
+  /**
+   *  Aspect members
+   */
+
+  private val traceContext = Kamon.context
+
+
+  /**
+   *  Advices
+   */
+  import kamon.TraceContextSwap.withContext
+
+  @Around("runnableExecution()")
+  def around(pjp: ProceedingJoinPoint) = {
+    import pjp._
+
+    withContext(traceContext, proceed())
+  }
+
+}
-- 
cgit v1.2.3