From 6ba7487204e0339437c7279e6330fc03419a2c5c Mon Sep 17 00:00:00 2001
From: Ivan Topolnjak <ivantopo@gmail.com>
Date: Sat, 23 Nov 2013 21:30:34 -0300
Subject: Put FutureTracing instrumentation in place, simpler and more explicit
 test

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

(limited to 'kamon-trace/src/main/scala/kamon/trace/instrumentation/FutureTracing.scala')

diff --git a/kamon-trace/src/main/scala/kamon/trace/instrumentation/FutureTracing.scala b/kamon-trace/src/main/scala/kamon/trace/instrumentation/FutureTracing.scala
new file mode 100644
index 00000000..10daa2da
--- /dev/null
+++ b/kamon-trace/src/main/scala/kamon/trace/instrumentation/FutureTracing.scala
@@ -0,0 +1,40 @@
+package kamon.trace.instrumentation
+
+import org.aspectj.lang.annotation._
+import org.aspectj.lang.ProceedingJoinPoint
+import kamon.trace.{ContextAware, TraceContext, Trace}
+
+@Aspect
+class FutureTracing {
+
+  /**
+   *  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: ContextAware = ContextAware.default
+
+
+  @Pointcut("execution(kamon.trace.ContextAware+.new(..)) && this(runnable)")
+  def instrumentedRunnableCreation(runnable: ContextAware): Unit = {}
+
+  @Pointcut("execution(* kamon.trace.ContextAware+.run()) && this(runnable)")
+  def futureRunnableExecution(runnable: ContextAware) = {}
+
+
+  @After("instrumentedRunnableCreation(runnable)")
+  def beforeCreation(runnable: ContextAware): Unit = {
+    // Force traceContext initialization.
+    runnable.traceContext
+  }
+
+  @Around("futureRunnableExecution(runnable)")
+  def around(pjp: ProceedingJoinPoint, runnable: ContextAware): Any = {
+    import pjp._
+
+    Trace.withValue(runnable.traceContext) {
+      proceed()
+    }
+  }
+
+}
\ No newline at end of file
-- 
cgit v1.2.3