ensure that kamon.trace.ask-pattern-tracing setting is honored, closes #4

5 files changed, 111 insertions, 48 deletions

diff --git a/kamon-core/src/main/resources/reference.conf b/kamon-core/src/main/resources/reference.conf
index 706132bd..3868b743 100644
--- a/kamon-core/src/main/resources/reference.conf
+++ b/kamon-core/src/main/resources/reference.conf
@@ -53,6 +53,9 @@ kamon {
   }
 
   trace {
+    # If ask-pattern-tracing is enabled, a WARN level log message will be generated if a future generated by the `ask`
+    # pattern fails with a `AskTimeoutException` and the log message will contain a stack trace captured at the moment
+    # the future was created.
     ask-pattern-tracing = off
   }
 
diff --git a/kamon-core/src/main/scala/akka/instrumentation/AskPatternTracing.scala b/kamon-core/src/main/scala/akka/instrumentation/AskPatternTracing.scala
index b5b23e61..08508930 100644
--- a/kamon-core/src/main/scala/akka/instrumentation/AskPatternTracing.scala
+++ b/kamon-core/src/main/scala/akka/instrumentation/AskPatternTracing.scala
@@ -1,18 +1,19 @@
-/* ===================================================
+/*
+ * =========================================================================================
  * 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
+ * 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
+ *  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.
- * ========================================================== */
+ * 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 akka.instrumentation
 
 import org.aspectj.lang.annotation.{ AfterReturning, Pointcut, Aspect }
@@ -20,6 +21,8 @@ import akka.event.Logging.Warning
 import scala.compat.Platform.EOL
 import akka.actor.ActorRefProvider
 import akka.pattern.{ AskTimeoutException, PromiseActorRef }
+import kamon.trace.Trace
+import kamon.Kamon
 
 @Aspect
 class AskPatternTracing {
@@ -27,23 +30,25 @@ class AskPatternTracing {
   class StackTraceCaptureException extends Throwable
 
   @Pointcut(value = "execution(* akka.pattern.PromiseActorRef$.apply(..)) && args(provider, *)", argNames = "provider")
-  def promiseActorRefApply(provider: ActorRefProvider): Unit = {
-    provider.settings.config.getBoolean("kamon.trace.ask-pattern-tracing")
-  }
+  def promiseActorRefApply(provider: ActorRefProvider): Unit = {}
 
   @AfterReturning(pointcut = "promiseActorRefApply(provider)", returning = "promiseActor")
   def hookAskTimeoutWarning(provider: ActorRefProvider, promiseActor: PromiseActorRef): Unit = {
-    val future = promiseActor.result.future
     val system = promiseActor.provider.guardian.underlying.system
-    implicit val ec = system.dispatcher
-    val stack = new StackTraceCaptureException
+    val traceExtension = Kamon(Trace)(system)
+
+    if (traceExtension.enableAskPatternTracing) {
+      val future = promiseActor.result.future
+      implicit val ec = system.dispatcher
+      val stack = new StackTraceCaptureException
 
-    future onFailure {
-      case timeout: AskTimeoutException ⇒
-        val stackString = stack.getStackTrace.drop(3).mkString("", EOL, EOL)
+      future onFailure {
+        case timeout: AskTimeoutException ⇒
+          val stackString = stack.getStackTrace.drop(3).mkString("", EOL, EOL)
 
-        system.eventStream.publish(Warning("AskPatternTracing", classOf[AskPatternTracing],
-          "Timeout triggered for ask pattern registered at: " + stackString))
+          system.eventStream.publish(Warning("AskPatternTracing", classOf[AskPatternTracing],
+            "Timeout triggered for ask pattern registered at: " + stackString))
+      }
     }
   }
 }
diff --git a/kamon-core/src/main/scala/kamon/instrumentation/ActorLoggingTracing.scala b/kamon-core/src/main/scala/kamon/instrumentation/ActorLoggingTracing.scala
index 954f351a..85d39d3e 100644
--- a/kamon-core/src/main/scala/kamon/instrumentation/ActorLoggingTracing.scala
+++ b/kamon-core/src/main/scala/kamon/instrumentation/ActorLoggingTracing.scala
@@ -1,21 +1,22 @@
-/* ===================================================
+/*
+ * =========================================================================================
  * 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
+ * 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
+ *  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.
- * ========================================================== */
+ * 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.instrumentation
 
-import org.aspectj.lang.annotation.{ Around, Pointcut, DeclareMixin, Aspect }
+import org.aspectj.lang.annotation._
 import org.aspectj.lang.ProceedingJoinPoint
 import kamon.trace.{ TraceContextAware, TraceRecorder }
 
@@ -25,6 +26,15 @@ class ActorLoggingTracing {
   @DeclareMixin("akka.event.Logging.LogEvent+")
   def mixinTraceContextAwareToLogEvent: TraceContextAware = TraceContextAware.default
 
+  @Pointcut("execution(akka.event.Logging.LogEvent+.new(..)) && this(event)")
+  def logEventCreation(event: TraceContextAware): Unit = {}
+
+  @After("logEventCreation(event)")
+  def captureTraceContext(event: TraceContextAware): Unit = {
+    // Force initialization of TraceContextAware
+    event.traceContext
+  }
+
   @Pointcut("execution(* akka.event.slf4j.Slf4jLogger.withMdc(..)) && args(logSource, logEvent, logStatement)")
   def withMdcInvocation(logSource: String, logEvent: TraceContextAware, logStatement: () ⇒ _): Unit = {}
 
diff --git a/kamon-core/src/main/scala/kamon/trace/TraceExtension.scala b/kamon-core/src/main/scala/kamon/trace/TraceExtension.scala
new file mode 100644
index 00000000..a59abc18
--- /dev/null
+++ b/kamon-core/src/main/scala/kamon/trace/TraceExtension.scala
@@ -0,0 +1,36 @@
+/*
+ * =========================================================================================
+ * 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.trace
+
+import akka.actor.{ ExtendedActorSystem, ExtensionIdProvider, ExtensionId }
+import akka.actor
+import kamon.util.GlobPathFilter
+import kamon.Kamon
+
+class TraceExtension(system: ExtendedActorSystem) extends Kamon.Extension {
+  val config = system.settings.config.getConfig("kamon.trace")
+  val enableAskPatternTracing = config.getBoolean("ask-pattern-tracing")
+}
+
+object Trace extends ExtensionId[TraceExtension] with ExtensionIdProvider {
+  def lookup(): ExtensionId[_ <: actor.Extension] = Trace
+  def createExtension(system: ExtendedActorSystem): TraceExtension = new TraceExtension(system)
+
+  case class MetricGroupFilter(includes: List[GlobPathFilter], excludes: List[GlobPathFilter]) {
+    def accept(name: String): Boolean = includes.exists(_.accept(name)) && !excludes.exists(_.accept(name))
+  }
+}
diff --git a/kamon-core/src/test/scala/kamon/trace/instrumentation/AskPatternTracingSpec.scala b/kamon-core/src/test/scala/kamon/trace/instrumentation/AskPatternTracingSpec.scala
index 04a5c0e3..fb886de6 100644
--- a/kamon-core/src/test/scala/kamon/trace/instrumentation/AskPatternTracingSpec.scala
+++ b/kamon-core/src/test/scala/kamon/trace/instrumentation/AskPatternTracingSpec.scala
@@ -1,21 +1,22 @@
-/* ===================================================
+/*
+ * =========================================================================================
  * 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
+ * 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
+ *  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.
- * ========================================================== */
+ * 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.trace.instrumentation
 
-import akka.testkit.TestKit
+import akka.testkit.TestKitBase
 import akka.actor.{ Props, Actor, ActorSystem }
 import org.scalatest.{ Matchers, WordSpecLike }
 import akka.event.Logging.Warning
@@ -23,9 +24,17 @@ import scala.concurrent.duration._
 import akka.pattern.ask
 import akka.util.Timeout
 import kamon.trace.{ TraceContextAware, TraceRecorder }
-import org.scalatest.OptionValues._
+import com.typesafe.config.ConfigFactory
 
-class AskPatternTracingSpec extends TestKit(ActorSystem("ask-pattern-tracing-spec")) with WordSpecLike with Matchers {
+class AskPatternTracingSpec extends TestKitBase with WordSpecLike with Matchers {
+  implicit lazy val system: ActorSystem = ActorSystem("ask-pattern-tracing-spec", ConfigFactory.parseString(
+    """
+      |kamon {
+      |  trace {
+      |    ask-pattern-tracing = on
+      |  }
+      |}
+    """.stripMargin))
 
   "the AskPatternTracing" should {
     "log a warning with a stack trace and TraceContext taken from the moment the ask was triggered" in {