Change usagestats expvars so they can be set twice (#5515)

* Change Edition so it can be set twice Signed-off-by: Michel Hollands <michel.hollands@grafana.com> * Change Target so it can be set twice Signed-off-by: Michel Hollands <michel.hollands@grafana.com> * Test Target can be set twice Signed-off-by: Michel Hollands <michel.hollands@grafana.com> * Mention Target can be set multiple times Signed-off-by: Michel Hollands <michel.hollands@grafana.com> * Add existence checks to other stats types and add tests Signed-off-by: Michel Hollands <michel.hollands@grafana.com> * Remove from target and edition Signed-off-by: Michel Hollands <michel.hollands@grafana.com>
grafana · Mar 4, 2022 · 9e0941b · 9e0941b
1 parent 84c97ce
commit 9e0941b
Show file tree

Hide file tree

Showing 2 changed files with 106 additions and 7 deletions.
diff --git a/pkg/usagestats/stats.go b/pkg/usagestats/stats.go
@@ -154,26 +154,50 @@ func memstats() interface{} {
 }
 
 // NewFloat returns a new Float stats object.
+// If a Float stats object with the same name already exists it is returned.
 func NewFloat(name string) *expvar.Float {
+	existing := expvar.Get(statsPrefix + name)
+	if existing != nil {
+		if f, ok := existing.(*expvar.Float); ok {
+			return f
+		}
+		panic(fmt.Sprintf("%v is set to a non-float value", name))
+	}
 	return expvar.NewFloat(statsPrefix + name)
 }
 
 // NewInt returns a new Int stats object.
+// If an Int stats object object with the same name already exists it is returned.
 func NewInt(name string) *expvar.Int {
+	existing := expvar.Get(statsPrefix + name)
+	if existing != nil {
+		if i, ok := existing.(*expvar.Int); ok {
+			return i
+		}
+		panic(fmt.Sprintf("%v is set to a non-int value", name))
+	}
 	return expvar.NewInt(statsPrefix + name)
 }
 
 // NewString returns a new String stats object.
+// If a String stats object with the same name already exists it is returned.
 func NewString(name string) *expvar.String {
+	existing := expvar.Get(statsPrefix + name)
+	if existing != nil {
+		if s, ok := existing.(*expvar.String); ok {
+			return s
+		}
+		panic(fmt.Sprintf("%v is set to a non-string value", name))
+	}
 	return expvar.NewString(statsPrefix + name)
 }
 
-// Target sets the target name.
+// Target sets the target name. This can be set multiple times.
 func Target(target string) {
 	NewString(targetKey).Set(target)
 }
 
-// Edition sets the edition name.
+// Edition sets the edition name. This can be set multiple times.
 func Edition(edition string) {
 	NewString(editionKey).Set(edition)
 }
@@ -199,6 +223,7 @@ type Statistics struct {
 // - count
 // - stddev
 // - stdvar
+// If a Statistics object with the same name already exists it is returned.
 func NewStatistics(name string) *Statistics {
 	s := &Statistics{
 		min:   atomic.NewFloat64(math.Inf(0)),
@@ -208,6 +233,13 @@ func NewStatistics(name string) *Statistics {
 		mean:  atomic.NewFloat64(0),
 		value: atomic.NewFloat64(0),
 	}
+	existing := expvar.Get(statsPrefix + name)
+	if existing != nil {
+		if s, ok := existing.(*Statistics); ok {
+			return s
+		}
+		panic(fmt.Sprintf("%v is set to a non-Statistics value", name))
+	}
 	expvar.Publish(statsPrefix+name, s)
 	return s
 }
@@ -285,12 +317,20 @@ type Counter struct {
 }
 
 // NewCounter returns a new Counter stats object.
+// If a Counter stats object with the same name already exists it is returned.
 func NewCounter(name string) *Counter {
 	c := &Counter{
 		total:     atomic.NewInt64(0),
 		rate:      atomic.NewFloat64(0),
 		resetTime: time.Now(),
 	}
+	existing := expvar.Get(statsPrefix + name)
+	if existing != nil {
+		if c, ok := existing.(*Counter); ok {
+			return c
+		}
+		panic(fmt.Sprintf("%v is set to a non-Counter value", name))
+	}
 	expvar.Publish(statsPrefix+name, c)
 	return c
 }
@@ -328,12 +368,20 @@ type WordCounter struct {
 }
 
 // NewWordCounter returns a new WordCounter stats object.
-// WordCounter object is thread-safe and count the amount of word recorded.
+// The WordCounter object is thread-safe and counts the number of words recorded.
+// If a WordCounter stats object with the same name already exists it is returned.
 func NewWordCounter(name string) *WordCounter {
 	c := &WordCounter{
 		count: atomic.NewInt64(0),
 		words: sync.Map{},
 	}
+	existing := expvar.Get(statsPrefix + name)
+	if existing != nil {
+		if w, ok := existing.(*WordCounter); ok {
+			return w
+		}
+		panic(fmt.Sprintf("%v is set to a non-WordCounter value", name))
+	}
 	expvar.Publish(statsPrefix+name, c)
 	return c
 }

diff --git a/pkg/usagestats/stats_test.go b/pkg/usagestats/stats_test.go
@@ -19,16 +19,25 @@ func Test_BuildReport(t *testing.T) {
 		CreatedAt: now,
 	}
 
+	Edition("non-OSS")
 	Edition("OSS")
+	Target("distributor")
 	Target("compactor")
+	NewString("compression").Set("snappy")
 	NewString("compression").Set("lz4")
+	NewInt("compression_ratio").Set(50)
 	NewInt("compression_ratio").Set(100)
 	NewFloat("size_mb").Set(100.1)
+	NewFloat("size_mb").Set(200.1)
 	NewCounter("lines_written").Inc(200)
 	s := NewStatistics("query_throughput")
+	s.Record(25)
+	s = NewStatistics("query_throughput")
 	s.Record(300)
 	s.Record(5)
 	w := NewWordCounter("active_tenants")
+	w.Add("buz")
+	w = NewWordCounter("active_tenants")
 	w.Add("foo")
 	w.Add("bar")
 	w.Add("foo")
@@ -43,13 +52,13 @@ func Test_BuildReport(t *testing.T) {
 	require.Equal(t, r.Metrics["num_goroutine"], runtime.NumGoroutine())
 	require.Equal(t, r.Metrics["compression"], "lz4")
 	require.Equal(t, r.Metrics["compression_ratio"], int64(100))
-	require.Equal(t, r.Metrics["size_mb"], 100.1)
+	require.Equal(t, r.Metrics["size_mb"], 200.1)
 	require.Equal(t, r.Metrics["lines_written"].(map[string]interface{})["total"], int64(200))
 	require.Equal(t, r.Metrics["query_throughput"].(map[string]interface{})["min"], float64(5))
 	require.Equal(t, r.Metrics["query_throughput"].(map[string]interface{})["max"], float64(300))
-	require.Equal(t, r.Metrics["query_throughput"].(map[string]interface{})["count"], int64(2))
-	require.Equal(t, r.Metrics["query_throughput"].(map[string]interface{})["avg"], float64(300+5)/2)
-	require.Equal(t, r.Metrics["active_tenants"], int64(2))
+	require.Equal(t, r.Metrics["query_throughput"].(map[string]interface{})["count"], int64(3))
+	require.Equal(t, r.Metrics["query_throughput"].(map[string]interface{})["avg"], float64(25+300+5)/3)
+	require.Equal(t, r.Metrics["active_tenants"], int64(3))
 
 	out, _ := jsoniter.MarshalIndent(r, "", " ")
 	t.Log(string(out))
@@ -95,3 +104,45 @@ func TestWordCounter(t *testing.T) {
 	}
 	require.Equal(t, int64(2), w.Value())
 }
+
+func TestPanics(t *testing.T) {
+	require.Panics(t, func() {
+		NewStatistics("panicstats")
+		NewWordCounter("panicstats")
+	})
+
+	require.Panics(t, func() {
+		NewWordCounter("panicwordcounter")
+		NewCounter("panicwordcounter")
+	})
+
+	require.Panics(t, func() {
+		NewCounter("paniccounter")
+		NewStatistics("paniccounter")
+	})
+
+	require.Panics(t, func() {
+		NewFloat("panicfloat")
+		NewInt("panicfloat")
+	})
+
+	require.Panics(t, func() {
+		NewInt("panicint")
+		NewString("panicint")
+	})
+
+	require.Panics(t, func() {
+		NewString("panicstring")
+		NewFloat("panicstring")
+	})
+
+	require.Panics(t, func() {
+		NewFloat(targetKey)
+		Target("new target")
+	})
+
+	require.Panics(t, func() {
+		NewFloat(editionKey)
+		Edition("new edition")
+	})
+}