allow enabling/disabling of metrics on a Span

2 files changed, 34 insertions, 4 deletions

diff --git a/kamon-core-tests/src/test/scala/kamon/trace/SpanMetrics.scala b/kamon-core-tests/src/test/scala/kamon/trace/SpanMetrics.scala
index 4425ac92..ea29d92e 100644
--- a/kamon-core-tests/src/test/scala/kamon/trace/SpanMetrics.scala
+++ b/kamon-core-tests/src/test/scala/kamon/trace/SpanMetrics.scala
@@ -42,6 +42,27 @@ class SpanMetrics extends WordSpecLike with Matchers with MetricInspection with
 
     }
 
+    "not be recorded when disableMetrics() is called on the Span" in {
+      val operation = "span-with-disabled-metrics"
+      buildSpan(operation)
+        .start()
+        .disableMetrics()
+        .finish()
+
+      Span.Metrics.ProcessingTime.valuesForTag("operation") shouldNot contain(operation)
+    }
+
+    "be recorded if metrics are enabled by calling enableMetrics() on the Span" in {
+      val operation = "span-with-re-enabled-metrics"
+      buildSpan(operation)
+        .start()
+        .disableMetrics()
+        .enableMetrics()
+        .finish()
+
+      Span.Metrics.ProcessingTime.valuesForTag("operation") should contain(operation)
+    }
+
     "record correctly error latency and count" in {
       val operation = "span-failure"
       val operationTag = "operation" -> operation
diff --git a/kamon-core/src/main/scala/kamon/trace/Span.scala b/kamon-core/src/main/scala/kamon/trace/Span.scala
index ceb1f56c..7410523b 100644
--- a/kamon-core/src/main/scala/kamon/trace/Span.scala
+++ b/kamon-core/src/main/scala/kamon/trace/Span.scala
@@ -51,7 +51,9 @@ sealed abstract class Span {
 
   def setOperationName(name: String): Span
 
-  def disableMetricsCollection(): Span
+  def enableMetrics(): Span
+
+  def disableMetrics(): Span
 
   def finish(finishTimestampMicros: Long): Unit
 
@@ -77,7 +79,8 @@ object Span {
     override def addError(error: String): Span = this
     override def addError(error: String, throwable: Throwable): Span = this
     override def setOperationName(name: String): Span = this
-    override def disableMetricsCollection(): Span = this
+    override def enableMetrics(): Span = this
+    override def disableMetrics(): Span = this
     override def finish(finishTimestampMicros: Long): Unit = {}
   }
 
@@ -166,7 +169,12 @@ object Span {
       this
     }
 
-    override def disableMetricsCollection(): Span = synchronized {
+    override def enableMetrics(): Span = synchronized {
+      collectMetrics = true
+      this
+    }
+
+    override def disableMetrics(): Span = synchronized {
       collectMetrics = false
       this
     }
@@ -231,7 +239,8 @@ object Span {
     override def addError(error: String): Span = this
     override def addError(error: String, throwable: Throwable): Span = this
     override def setOperationName(name: String): Span = this
-    override def disableMetricsCollection(): Span = this
+    override def enableMetrics(): Span = this
+    override def disableMetrics(): Span = this
     override def finish(finishTimestampMicros: Long): Unit = {}
   }