From 66cf166772712267aac922c51ded90c7dec3bdd5 Mon Sep 17 00:00:00 2001
From: Ivan Topolnjak <ivantopo@gmail.com>
Date: Tue, 12 Dec 2017 00:40:15 +0100
Subject: allow disabling/enabling metrics collection from the SpanBuilder

---
 .../src/test/scala/kamon/trace/LocalSpanSpec.scala         |  2 ++
 .../src/test/scala/kamon/trace/SpanMetrics.scala           |  7 ++++++-
 kamon-core/src/main/scala/kamon/trace/Span.scala           |  8 ++++----
 kamon-core/src/main/scala/kamon/trace/Tracer.scala         | 14 +++++++++++++-
 4 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/kamon-core-tests/src/test/scala/kamon/trace/LocalSpanSpec.scala b/kamon-core-tests/src/test/scala/kamon/trace/LocalSpanSpec.scala
index 5d636d16..21ad93b3 100644
--- a/kamon-core-tests/src/test/scala/kamon/trace/LocalSpanSpec.scala
+++ b/kamon-core-tests/src/test/scala/kamon/trace/LocalSpanSpec.scala
@@ -32,6 +32,8 @@ class LocalSpanSpec extends WordSpec with Matchers with BeforeAndAfterAll with E
         Kamon.buildSpan("test-span")
           .withTag("test", "value")
           .withStartTimestamp(100)
+          .disableMetrics()
+          .enableMetrics()
           .start()
           .finish(200)
 
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 ea29d92e..2a04024c 100644
--- a/kamon-core-tests/src/test/scala/kamon/trace/SpanMetrics.scala
+++ b/kamon-core-tests/src/test/scala/kamon/trace/SpanMetrics.scala
@@ -42,13 +42,18 @@ class SpanMetrics extends WordSpecLike with Matchers with MetricInspection with
 
     }
 
-    "not be recorded when disableMetrics() is called on the Span" in {
+    "not be recorded when disableMetrics() is called on the SpanBuilder or the Span" in {
       val operation = "span-with-disabled-metrics"
       buildSpan(operation)
         .start()
         .disableMetrics()
         .finish()
 
+      buildSpan(operation)
+        .disableMetrics()
+        .start()
+        .finish()
+
       Span.Metrics.ProcessingTime.valuesForTag("operation") shouldNot contain(operation)
     }
 
diff --git a/kamon-core/src/main/scala/kamon/trace/Span.scala b/kamon-core/src/main/scala/kamon/trace/Span.scala
index 7410523b..d8d11589 100644
--- a/kamon-core/src/main/scala/kamon/trace/Span.scala
+++ b/kamon-core/src/main/scala/kamon/trace/Span.scala
@@ -93,9 +93,9 @@ object Span {
     * @param spanSink
     */
   final class Local(spanContext: SpanContext, parent: Option[Span], initialOperationName: String, initialSpanTags: Map[String, Span.TagValue],
-    initialMetricTags: Map[String, String], startTimestampMicros: Long, spanSink: SpanSink, scopeSpanMetrics: Boolean) extends Span {
+    initialMetricTags: Map[String, String], startTimestampMicros: Long, spanSink: SpanSink, trackMetrics: Boolean, scopeSpanMetrics: Boolean) extends Span {
 
-    private var collectMetrics: Boolean = true
+    private var collectMetrics: Boolean = trackMetrics
     private var open: Boolean = true
     private val sampled: Boolean = spanContext.samplingDecision == SamplingDecision.Sample
     private var hasError: Boolean = false
@@ -222,8 +222,8 @@ object Span {
   object Local {
     def apply(spanContext: SpanContext, parent: Option[Span], initialOperationName: String, initialSpanTags: Map[String, Span.TagValue],
         initialMetricTags: Map[String, String], startTimestampMicros: Long, spanSink: SpanSink,
-        scopeSpanMetrics: Boolean): Local =
-      new Local(spanContext, parent, initialOperationName, initialSpanTags, initialMetricTags, startTimestampMicros, spanSink, scopeSpanMetrics)
+        trackMetrics: Boolean, scopeSpanMetrics: Boolean): Local =
+      new Local(spanContext, parent, initialOperationName, initialSpanTags, initialMetricTags, startTimestampMicros, spanSink, trackMetrics, scopeSpanMetrics)
   }
 
 
diff --git a/kamon-core/src/main/scala/kamon/trace/Tracer.scala b/kamon-core/src/main/scala/kamon/trace/Tracer.scala
index 03c44a2c..addeddf6 100644
--- a/kamon-core/src/main/scala/kamon/trace/Tracer.scala
+++ b/kamon-core/src/main/scala/kamon/trace/Tracer.scala
@@ -96,6 +96,7 @@ object Tracer {
     private var initialSpanTags = Map.empty[String, Span.TagValue]
     private var initialMetricTags = Map.empty[String, String]
     private var useParentFromContext = true
+    private var trackMetrics = true
 
     def asChildOf(parent: Span): SpanBuilder = {
       if(parent != Span.Empty) this.parentSpan = parent
@@ -145,6 +146,17 @@ object Tracer {
       this
     }
 
+    def enableMetrics(): SpanBuilder = {
+      this.trackMetrics = true
+      this
+    }
+
+    def disableMetrics(): SpanBuilder = {
+      this.trackMetrics = false
+      this
+    }
+
+
     def start(): Span = {
       val startTimestampMicros = if(startTimestamp != 0L) startTimestamp else Clock.microTimestamp()
 
@@ -165,7 +177,7 @@ object Tracer {
       }
 
       tracer.tracerMetrics.createdSpans.increment()
-      Span.Local(spanContext, nonRemoteParent, initialOperationName, initialSpanTags, initialMetricTags, startTimestampMicros, spanSink, tracer.scopeSpanMetrics)
+      Span.Local(spanContext, nonRemoteParent, initialOperationName, initialSpanTags, initialMetricTags, startTimestampMicros, spanSink, trackMetrics, tracer.scopeSpanMetrics)
     }
 
     private def joinParentContext(parent: Span, samplingDecision: SamplingDecision): SpanContext =
-- 
cgit v1.2.3