= core: Inline variant of withTraceContext using macros.

7 files changed, 71 insertions, 10 deletions

diff --git a/kamon-core/src/main/scala/kamon/instrumentation/FutureTracing.scala b/kamon-core/src/main/scala/kamon/instrumentation/FutureTracing.scala
index 634c94a1..b9424579 100644
--- a/kamon-core/src/main/scala/kamon/instrumentation/FutureTracing.scala
+++ b/kamon-core/src/main/scala/kamon/instrumentation/FutureTracing.scala
@@ -39,9 +39,9 @@ class FutureTracing {
 
   @Around("futureRelatedRunnableExecution(runnable)")
   def aroundExecution(pjp: ProceedingJoinPoint, runnable: TraceContextAware): Any = {
-    TraceRecorder.withTraceContext(runnable.traceContext) {
+    TraceRecorder.withInlineTraceContextReplacement(runnable.traceContext) {
       pjp.proceed()
     }
   }
 
-}
\ No newline at end of file
+}
diff --git a/kamon-core/src/main/scala/kamon/trace/TraceRecorder.scala b/kamon-core/src/main/scala/kamon/trace/TraceRecorder.scala
index 0e264cd2..c57ac064 100644
--- a/kamon-core/src/main/scala/kamon/trace/TraceRecorder.scala
+++ b/kamon-core/src/main/scala/kamon/trace/TraceRecorder.scala
@@ -16,7 +16,10 @@
 
 package kamon.trace
 
+import scala.language.experimental.macros
 import java.util.concurrent.atomic.AtomicLong
+import kamon.macros.InlineTraceContextMacro
+
 import scala.util.Try
 import java.net.InetAddress
 import akka.actor.ActorSystem
@@ -66,6 +69,8 @@ object TraceRecorder {
     try thunk finally setContext(oldContext)
   }
 
+  def withInlineTraceContextReplacement[T](traceCtx: Option[TraceContext])(thunk: ⇒ T): T = macro InlineTraceContextMacro.withInlineTraceContextImpl[T, Option[TraceContext]]
+
   def finish(metadata: Map[String, String] = Map.empty): Unit = currentContext.map(_.finish(metadata))
 
 }
diff --git a/kamon-datadog/src/test/scala/kamon/datadog/DatadogMetricSenderSpec.scala b/kamon-datadog/src/test/scala/kamon/datadog/DatadogMetricSenderSpec.scala
index 6a7191a1..bc692954 100644
--- a/kamon-datadog/src/test/scala/kamon/datadog/DatadogMetricSenderSpec.scala
+++ b/kamon-datadog/src/test/scala/kamon/datadog/DatadogMetricSenderSpec.scala
@@ -16,10 +16,10 @@
 
 package kamon.datadog
 
-import akka.testkit.{TestKitBase, TestProbe}
-import akka.actor.{Props, ActorRef, ActorSystem}
+import akka.testkit.{ TestKitBase, TestProbe }
+import akka.actor.{ Props, ActorRef, ActorSystem }
 import kamon.metrics.instruments.CounterRecorder
-import org.scalatest.{Matchers, WordSpecLike}
+import org.scalatest.{ Matchers, WordSpecLike }
 import kamon.metrics._
 import akka.io.Udp
 import org.HdrHistogram.HdrRecorder
@@ -70,7 +70,7 @@ class DatadogMetricSenderSpec extends TestKitBase with WordSpecLike with Matcher
       }
 
       val udp = setup(Map(testMetricName -> testRecorder.collect()))
-      udp.expectMsgType[Udp.Send]// let the first flush pass
+      udp.expectMsgType[Udp.Send] // let the first flush pass
 
       val Udp.Send(data, _, _) = udp.expectMsgType[Udp.Send]
       data.utf8String should be(s"kamon.actor.$testMetricName:$level|ms|#actor:user/kamon")
diff --git a/kamon-macros/src/main/scala/kamon/macros/InlineTraceContextMacro.scala b/kamon-macros/src/main/scala/kamon/macros/InlineTraceContextMacro.scala
new file mode 100644
index 00000000..ea8cc7b8
--- /dev/null
+++ b/kamon-macros/src/main/scala/kamon/macros/InlineTraceContextMacro.scala
@@ -0,0 +1,46 @@
+/*
+ * =========================================================================================
+ * Copyright © 2013 the kamon project <http://kamon.io/>
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+ * either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ * =========================================================================================
+ */
+
+package kamon.macros
+
+import scala.language.experimental.macros
+import scala.reflect.macros.Context
+
+object InlineTraceContextMacro {
+
+  // Macro to generate an inline version of kamon.trace.TraceRecorder.withTraceContext
+  def withInlineTraceContextImpl[T: c.WeakTypeTag, TC: c.WeakTypeTag](c: Context)(traceCtx: c.Expr[TC])(thunk: c.Expr[T]) = {
+    import c.universe._
+
+    val inlineThunk =
+      Block(
+        List(
+          ValDef(
+            Modifiers(), newTermName("oldContext"), TypeTree(),
+            Select(Ident(newTermName("TraceRecorder")), newTermName("currentContext"))),
+          Apply(
+            Select(Ident(newTermName("TraceRecorder")), newTermName("setContext")),
+            List(traceCtx.tree))),
+        Try(
+          thunk.tree,
+          List(),
+          Apply(
+            Select(Ident(newTermName("TraceRecorder")), newTermName("setContext")),
+            List(Ident(newTermName("oldContext"))))))
+
+    c.Expr[T](inlineThunk)
+  }
+}
diff --git a/project/Dependencies.scala b/project/Dependencies.scala
index 667f8dce..4a2f2276 100644
--- a/project/Dependencies.scala
+++ b/project/Dependencies.scala
@@ -35,7 +35,7 @@ object Dependencies {
   val slf4Api         = "org.slf4j"                 %   "slf4j-api"             % slf4jVersion
   val slf4nop         = "org.slf4j"                 %   "slf4j-nop"             % slf4jVersion
   val jsr166          = "io.gatling"                %   "jsr166e"               % "1.0"
-
+  val scalaCompiler   = "org.scala-lang"            %   "scala-compiler"        % Settings.ScalaVersion
 
   def compile   (deps: ModuleID*): Seq[ModuleID] = deps map (_ % "compile")
   def provided  (deps: ModuleID*): Seq[ModuleID] = deps map (_ % "provided")
diff --git a/project/Projects.scala b/project/Projects.scala
index 5d747e49..1043ae58 100644
--- a/project/Projects.scala
+++ b/project/Projects.scala
@@ -14,10 +14,13 @@ object Projects extends Build {
 
 
   lazy val kamonCore = Project("kamon-core", file("kamon-core"))
+    .dependsOn(kamonMacros)
     .settings(basicSettings: _*)
     .settings(formatSettings: _*)
     .settings(aspectJSettings: _*)
     .settings(
+      mappings in (Compile, packageBin) ++= mappings.in(kamonMacros, Compile, packageBin).value,
+      mappings in (Compile, packageSrc) ++= mappings.in(kamonMacros, Compile, packageSrc).value,
       libraryDependencies ++=
         compile(akkaActor, aspectJ, aspectjWeaver, hdrHistogram, jsr166) ++
         provided(logback) ++
@@ -81,13 +84,18 @@ object Projects extends Build {
     .settings(formatSettings: _*)
     .settings(libraryDependencies ++= compile(akkaActor) ++  test(scalatest, akkaTestKit, slf4Api, slf4nop))
     .dependsOn(kamonCore)
- 
+
   lazy val kamonDatadog = Project("kamon-datadog", file("kamon-datadog"))
     .settings(basicSettings: _*)
     .settings(formatSettings: _*)
     .settings(libraryDependencies ++= compile(akkaActor) ++  test(scalatest, akkaTestKit, slf4Api, slf4nop))
     .dependsOn(kamonCore)
 
+  lazy val kamonMacros = Project("kamon-macros", file("kamon-macros"))
+    .settings(basicSettings: _*)
+    .settings(formatSettings: _*)
+    .settings(noPublishing: _*)
+    .settings(libraryDependencies ++= compile(scalaCompiler))
 
   val noPublishing = Seq(publish := (), publishLocal := (), publishArtifact := false)
 }
diff --git a/project/Settings.scala b/project/Settings.scala
index 9abd5553..5805192e 100644
--- a/project/Settings.scala
+++ b/project/Settings.scala
@@ -8,8 +8,10 @@ import scalariform.formatter.preferences._
 
 object Settings {
 
-  lazy val basicSettings = seq(
-    scalaVersion  := "2.10.3",
+  val ScalaVersion = "2.10.3"
+
+  lazy val basicSettings = Seq(
+    scalaVersion  := ScalaVersion,
     resolvers    ++= Dependencies.resolutionRepos,
     fork in run   := true,
     scalacOptions := Seq(