aggmetric: add clear method in ChildStorage

This patch introduces `clear` method in ChildStorage. It will clear all
children which are part of parent metric. This method will be used during
reinitialisation of aggmetric on cluster setting changes.

Epic: CRDB-43153
Part of: CRDB-48253
Release note: None

Author: aa-joshi

pkg/util/metric/aggmetric/BUILD.bazel

@@ -46,6 +46,7 @@ go_test(
         "@com_github_cockroachdb_crlib//testutils/require",
         "@com_github_prometheus_client_model//go",
         "@com_github_prometheus_common//expfmt",
+        "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
     ],
 )

pkg/util/metric/aggmetric/agg_metric.go

@@ -149,6 +149,16 @@ func (cs *childSet) get(labelVals ...string) (ChildMetric, bool) {
 	return cs.mu.children.Get(labelVals...)
 }
 
+// clear method removes all children from the childSet. It does not reset parent metric values.
+// Method should cautiously be used when childSet is reinitialised/updated. Today, it is
+// only used when cluster settings are updated to support app and db label values. For normal
+// operations, please use add, remove and get method to update the childSet.
+func (cs *childSet) clear() {
+	cs.mu.Lock()
+	defer cs.mu.Unlock()
+	cs.mu.children.Clear()
+}
+
 type MetricItem interface {
 	labelValuer
 }
@@ -185,6 +195,7 @@ type ChildrenStorage interface {
 	Del(key ChildMetric)
 	Do(f func(e interface{}))
 	GetChildMetric(e interface{}) ChildMetric
+	Clear()
 }
 
 var _ ChildrenStorage = &UnorderedCacheWrapper{}
@@ -229,6 +240,10 @@ func (ucw *UnorderedCacheWrapper) Do(f func(e interface{})) {
 	})
 }
 
+func (ucw *UnorderedCacheWrapper) Clear() {
+	ucw.cache.Clear()
+}
+
 type BtreeWrapper struct {
 	tree *btree.BTree
 }
@@ -267,6 +282,10 @@ func (b BtreeWrapper) GetChildMetric(e interface{}) ChildMetric {
 	return e.(ChildMetric)
 }
 
+func (b BtreeWrapper) Clear() {
+	b.tree.Clear(false)
+}
+
 func (lv *labelValuesSlice) Less(o btree.Item) bool {
 	ov := o.(labelValuer).labelValues()
 	if len(ov) != len(*lv) {

pkg/util/metric/aggmetric/agg_metric_test.go

@@ -71,22 +71,7 @@ func TestAggMetric(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 
 	r := metric.NewRegistry()
-	writePrometheusMetrics := func(t *testing.T) string {
-		var in bytes.Buffer
-		ex := metric.MakePrometheusExporter()
-		scrape := func(ex *metric.PrometheusExporter) {
-			ex.ScrapeRegistry(r, true /* includeChildMetrics */, true)
-		}
-		require.NoError(t, ex.ScrapeAndPrintAsText(&in, expfmt.FmtText, scrape))
-		var lines []string
-		for sc := bufio.NewScanner(&in); sc.Scan(); {
-			if !bytes.HasPrefix(sc.Bytes(), []byte{'#'}) {
-				lines = append(lines, sc.Text())
-			}
-		}
-		sort.Strings(lines)
-		return strings.Join(lines, "\n")
-	}
+	writePrometheusMetrics := WritePrometheusMetricsFunc(r)
 
 	c := NewCounter(metric.Metadata{
 		Name: "foo_counter",
@@ -296,3 +281,59 @@ func TestAggHistogramRotate(t *testing.T) {
 		// Go to beginning.
 	}
 }
+
+func TestAggMetricClear(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	r := metric.NewRegistry()
+	writePrometheusMetrics := WritePrometheusMetricsFunc(r)
+
+	c := NewCounter(metric.Metadata{
+		Name: "foo_counter",
+	}, "tenant_id")
+	r.AddMetric(c)
+
+	d := NewCounter(metric.Metadata{
+		Name: "bar_counter",
+	}, "tenant_id")
+	d.initWithCacheStorageType([]string{"tenant_id"})
+	r.AddMetric(d)
+
+	tenant2 := roachpb.MustMakeTenantID(2)
+	c1 := c.AddChild(tenant2.String())
+
+	t.Run("before clear", func(t *testing.T) {
+		c1.Inc(2)
+		d.Inc(2, "3")
+		testFile := "aggMetric_pre_clear.txt"
+		echotest.Require(t, writePrometheusMetrics(t), datapathutils.TestDataPath(t, testFile))
+	})
+
+	c.clear()
+	d.clear()
+
+	t.Run("post clear", func(t *testing.T) {
+		testFile := "aggMetric_post_clear.txt"
+		echotest.Require(t, writePrometheusMetrics(t), datapathutils.TestDataPath(t, testFile))
+	})
+}
+
+func WritePrometheusMetricsFunc(r *metric.Registry) func(t *testing.T) string {
+	writePrometheusMetrics := func(t *testing.T) string {
+		var in bytes.Buffer
+		ex := metric.MakePrometheusExporter()
+		scrape := func(ex *metric.PrometheusExporter) {
+			ex.ScrapeRegistry(r, true /* includeChildMetrics */, true)
+		}
+		require.NoError(t, ex.ScrapeAndPrintAsText(&in, expfmt.FmtText, scrape))
+		var lines []string
+		for sc := bufio.NewScanner(&in); sc.Scan(); {
+			if !bytes.HasPrefix(sc.Bytes(), []byte{'#'}) {
+				lines = append(lines, sc.Text())
+			}
+		}
+		sort.Strings(lines)
+		return strings.Join(lines, "\n")
+	}
+	return writePrometheusMetrics
+}

pkg/util/metric/aggmetric/gauge_test.go

@@ -20,13 +20,14 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/metric"
 	"github.com/cockroachdb/crlib/testutils/require"
 	"github.com/prometheus/common/expfmt"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestAggGaugeEviction(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	const cacheSize = 10
 	r := metric.NewRegistry()
-	writePrometheusMetrics := writePrometheusMetricsFunc(r)
+	writePrometheusMetrics := WritePrometheusMetricsFunc(r)
 
 	g := NewGauge(metric.Metadata{
 		Name: "foo_gauge",
@@ -101,11 +102,30 @@ func TestAggGaugeMethods(t *testing.T) {
 	echotest.Require(t, writePrometheusMetrics(t), datapathutils.TestDataPath(t, testFile))
 }
 
+func TestPanicForAggGaugeWithBtreeStorage(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	g := NewGauge(metric.Metadata{
+		Name: "foo_gauge",
+	}, "tenant_id")
+
+	assert.Panics(t, func() {
+		g.Inc(1, "1", "1")
+	}, "expected panic when Inc is invoked on AggGauge with BTreeStorage")
+
+	assert.Panics(t, func() {
+		g.Dec(1, "1", "1")
+	}, "expected panic when Dec is invoked on AggGauge with BTreeStorage")
+
+	assert.Panics(t, func() {
+		g.Update(1, "1", "1")
+	}, "expected panic when Update is invoked on AggGauge with BTreeStorage")
+}
+
 func TestAggGaugeFloat64Eviction(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	const cacheSize = 10
 	r := metric.NewRegistry()
-	writePrometheusMetrics := writePrometheusMetricsFunc(r)
+	writePrometheusMetrics := WritePrometheusMetricsFunc(r)
 
 	g := NewGaugeFloat64(metric.Metadata{
 		Name: "foo_gauge_float64",
@@ -136,22 +156,13 @@ func TestAggGaugeFloat64Eviction(t *testing.T) {
 	echotest.Require(t, writePrometheusMetrics(t), datapathutils.TestDataPath(t, testFile))
 }
 
-func writePrometheusMetricsFunc(r *metric.Registry) func(t *testing.T) string {
-	writePrometheusMetrics := func(t *testing.T) string {
-		var in bytes.Buffer
-		ex := metric.MakePrometheusExporter()
-		scrape := func(ex *metric.PrometheusExporter) {
-			ex.ScrapeRegistry(r, true /* includeChildMetrics */, true)
-		}
-		require.NoError(t, ex.ScrapeAndPrintAsText(&in, expfmt.FmtText, scrape))
-		var lines []string
-		for sc := bufio.NewScanner(&in); sc.Scan(); {
-			if !bytes.HasPrefix(sc.Bytes(), []byte{'#'}) {
-				lines = append(lines, sc.Text())
-			}
-		}
-		sort.Strings(lines)
-		return strings.Join(lines, "\n")
-	}
-	return writePrometheusMetrics
+func TestPanicForAggGaugeFloat64WithBtreeStorage(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	g := NewGaugeFloat64(metric.Metadata{
+		Name: "foo_gauge",
+	}, "tenant_id")
+
+	assert.Panics(t, func() {
+		g.Update(1, "1", "1")
+	}, "expected panic when Update is invoked on AggGaugeFloat64 with BTreeStorage")
 }

pkg/util/metric/aggmetric/testdata/aggMetric_post_clear.txt

@@ -0,0 +1,4 @@
+echo
+----
+bar_counter 2
+foo_counter 2

pkg/util/metric/aggmetric/testdata/aggMetric_pre_clear.txt

@@ -0,0 +1,6 @@
+echo
+----
+bar_counter 2
+bar_counter{tenant_id="3"} 2
+foo_counter 2
+foo_counter{tenant_id="2"} 2