util-stats: propagate CounterishGaugeType

Problem: Regardess of having set `.withCounterishGauge()`, when building a gauge via the `MetricBuilder.gauge()` api, we always return a gauge schema. This makes the `.withCounterishGauge()`, in effect, non-functional. Solution: Rather than copying a `MetricBuilder` and resetting its `metricType`, simply verify that we're only calling the `.gauge()` `.counter()` or `.stat()` api on a `MetricBuilder` with the appropriate `metricType` JIRA Issues: CSL-11280, CSL-11336 Differential Revision: https://phabricator.twitter.biz/D749274
twitter · Sep 23, 2021 · 0bfe12d · 0bfe12d
1 parent 2d5817f
commit 0bfe12d
Show file tree

Hide file tree

Showing 6 changed files with 46 additions and 24 deletions.
diff --git a/util-jvm/src/main/scala/com/twitter/jvm/JvmStats.scala b/util-jvm/src/main/scala/com/twitter/jvm/JvmStats.scala
@@ -163,10 +163,10 @@ object JvmStats {
     }
 
     // note, these could be -1 if the collector doesn't have support for it.
-    val cycles = gcStats.addGauge("cycles") {
+    val cycles = gcStats.metricBuilder(GaugeType).withCounterishGauge.gauge("cycles") {
       gcPool.map(_.getCollectionCount).filter(_ > 0).sum.toFloat
     }
-    val msec = gcStats.addGauge("msec") {
+    val msec = gcStats.metricBuilder(GaugeType).withCounterishGauge.gauge("msec") {
       gcPool.map(_.getCollectionTime).filter(_ > 0).sum.toFloat
     }
 

diff --git a/util-jvm/src/test/scala/com/twitter/jvm/JvmStatsTest.scala b/util-jvm/src/test/scala/com/twitter/jvm/JvmStatsTest.scala
@@ -1,7 +1,9 @@
 package com.twitter.jvm
 
 import com.twitter.finagle.stats.InMemoryStatsReceiver
-import com.twitter.finagle.stats.exp.{ExpressionSchema, ExpressionSchemaKey, MetricExpression}
+import com.twitter.finagle.stats.exp.ExpressionSchema
+import com.twitter.finagle.stats.exp.ExpressionSchemaKey
+import com.twitter.finagle.stats.exp.MetricExpression
 import org.scalatest.funsuite.AnyFunSuite
 
 class JvmStatsTest extends AnyFunSuite {

diff --git a/util-stats/src/main/scala/com/twitter/finagle/stats/MetricBuilder.scala b/util-stats/src/main/scala/com/twitter/finagle/stats/MetricBuilder.scala
@@ -1,13 +1,11 @@
 package com.twitter.finagle.stats
 
-import com.twitter.finagle.stats.MetricBuilder.{
-  CounterType,
-  CounterishGaugeType,
-  GaugeType,
-  HistogramType,
-  MetricType,
-  UnlatchedCounter
-}
+import com.twitter.finagle.stats.MetricBuilder.CounterType
+import com.twitter.finagle.stats.MetricBuilder.CounterishGaugeType
+import com.twitter.finagle.stats.MetricBuilder.GaugeType
+import com.twitter.finagle.stats.MetricBuilder.HistogramType
+import com.twitter.finagle.stats.MetricBuilder.MetricType
+import com.twitter.finagle.stats.MetricBuilder.UnlatchedCounter
 import java.util.function.Supplier
 import scala.annotation.varargs
 
@@ -268,17 +266,17 @@ class MetricBuilder private (
    */
   @varargs
   final def counter(name: String*): Counter = {
-    val schema = withName(name: _*).counterSchema
-    this.statsReceiver.counter(schema)
+    this.statsReceiver.validateMetricType(this, CounterType)
+    this.statsReceiver.counter(this.withName(name: _*))
   }
 
   /**
    * Produce a gauge as described by the builder inside the underlying StatsReceiver.
    * @return the gauge created.
    */
   final def gauge(name: String*)(f: => Float): Gauge = {
-    val schema = withName(name: _*).gaugeSchema
-    this.statsReceiver.addGauge(schema)(f)
+    this.statsReceiver.validateMetricType(this, GaugeType)
+    this.statsReceiver.addGauge(this.withName(name: _*))(f)
   }
 
   /**
@@ -288,8 +286,8 @@ class MetricBuilder private (
    */
   @varargs
   final def gauge(f: Supplier[Float], name: String*): Gauge = {
-    val schema = withName(name: _*).gaugeSchema
-    this.statsReceiver.addGauge(schema)(f.get())
+    this.statsReceiver.validateMetricType(this, GaugeType)
+    this.statsReceiver.addGauge(this.withName(name: _*))(f.get())
   }
 
   /**
@@ -298,8 +296,8 @@ class MetricBuilder private (
    */
   @varargs
   final def histogram(name: String*): Stat = {
-    val schema = withName(name: _*).histogramSchema
-    this.statsReceiver.stat(schema)
+    this.statsReceiver.validateMetricType(this, HistogramType)
+    this.statsReceiver.stat(this.withName(name: _*))
   }
 
   def canEqual(other: Any): Boolean = other.isInstanceOf[MetricBuilder]

diff --git a/util-stats/src/main/scala/com/twitter/finagle/stats/StatsReceiver.scala b/util-stats/src/main/scala/com/twitter/finagle/stats/StatsReceiver.scala
@@ -404,7 +404,10 @@ trait StatsReceiver {
   private[stats] def counter0(name: String): Counter = counter(name)
   private[stats] def stat0(name: String): Stat = stat(name)
 
-  protected def validateMetricType(metricBuilder: MetricBuilder, expectedType: MetricType): Unit = {
+  private[stats] def validateMetricType(
+    metricBuilder: MetricBuilder,
+    expectedType: MetricType
+  ): Unit = {
     val typeMatch = expectedType match {
       case CounterType =>
         metricBuilder.metricType == CounterType || metricBuilder.metricType == UnlatchedCounter

diff --git a/util-stats/src/test/java/com/twitter/finagle/stats/MetricBuilderCompilationTest.java b/util-stats/src/test/java/com/twitter/finagle/stats/MetricBuilderCompilationTest.java
@@ -33,9 +33,12 @@ public void testWithMethods() {
   @Test
   public void testConstructingMetrics() {
     StatsReceiver sr = new InMemoryStatsReceiver();
-    MetricBuilder mb = sr.metricBuilder(MetricBuilder.CounterType$.MODULE$);
-    Gauge g = mb.gauge(() -> 3.0f, "my", "cool", "gauge");
-    Stat s = mb.histogram("my", "cool", "histogram");
-    Counter c = mb.counter("my", "cool", "counter");
+    MetricBuilder gmb = sr.metricBuilder(MetricBuilder.GaugeType$.MODULE$);
+    MetricBuilder smb = sr.metricBuilder(MetricBuilder.HistogramType$.MODULE$);
+    MetricBuilder cmb = sr.metricBuilder(MetricBuilder.CounterType$.MODULE$);
+
+    Gauge g = gmb.gauge(() -> 3.0f, "my", "cool", "gauge");
+    Stat s = smb.histogram("my", "cool", "histogram");
+    Counter c = cmb.counter("my", "cool", "counter");
   }
 }
diff --git a/util-stats/src/test/scala/com/twitter/finagle/stats/MetricBuilderTest.scala b/util-stats/src/test/scala/com/twitter/finagle/stats/MetricBuilderTest.scala
@@ -1,6 +1,8 @@
 package com.twitter.finagle.stats
 
 import com.twitter.finagle.stats.MetricBuilder.CounterType
+import com.twitter.finagle.stats.MetricBuilder.CounterishGaugeType
+import com.twitter.finagle.stats.MetricBuilder.GaugeType
 import org.scalatest.funsuite.AnyFunSuite
 
 class MetricBuilderTest extends AnyFunSuite {
@@ -16,4 +18,18 @@ class MetricBuilderTest extends AnyFunSuite {
     metadataScopeSeparator.setSeparator("_aTerriblyLongStringForSomeReason_")
     assert(mb.toString.contains("foo_aTerriblyLongStringForSomeReason_bar"))
   }
+
+  test("the .gauge() api builds CounterishGauge vs Gauge where appropriate") {
+    val mb = MetricBuilder(
+      name = Seq("c_gauge"),
+      metricType = GaugeType,
+      statsReceiver = new InMemoryStatsReceiver
+    ).withCounterishGauge
+
+    assert(mb.metricType == CounterishGaugeType)
+
+    val g = mb.gauge("c_gauge") { 1.0.toFloat }
+
+    assert(g.metadata.toMetricBuilder.get.metricType == CounterishGaugeType)
+  }
 }