add metric reset to wrapped prometheus registry

Change-Id: I035863210e19184eb2e8861772e58d2d1de8ec5e

staging/src/k8s.io/component-base/metrics/BUILD

@@ -45,6 +45,7 @@ go_test(
         "desc_test.go",
         "gauge_test.go",
         "histogram_test.go",
+        "http_test.go",
         "opts_test.go",
         "registry_test.go",
         "summary_test.go",

staging/src/k8s.io/component-base/metrics/http.go

@@ -17,6 +17,7 @@ limitations under the License.
 package metrics
 
 import (
+	"io"
 	"net/http"
 
 	"github.com/prometheus/client_golang/prometheus/promhttp"
@@ -61,3 +62,16 @@ func (ho *HandlerOpts) toPromhttpHandlerOpts() promhttp.HandlerOpts {
 func HandlerFor(reg Gatherer, opts HandlerOpts) http.Handler {
 	return promhttp.HandlerFor(reg, opts.toPromhttpHandlerOpts())
 }
+
+// HandlerWithReset return an http.Handler with Reset
+func HandlerWithReset(reg KubeRegistry, opts HandlerOpts) http.Handler {
+	defaultHandler := promhttp.HandlerFor(reg, opts.toPromhttpHandlerOpts())
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method == http.MethodDelete {
+			reg.Reset()
+			io.WriteString(w, "metrics reset\n")
+			return
+		}
+		defaultHandler.ServeHTTP(w, r)
+	})
+}

staging/src/k8s.io/component-base/metrics/http_test.go

@@ -0,0 +1,69 @@
+/*
+Copyright 2020 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package metrics
+
+import (
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	apimachineryversion "k8s.io/apimachinery/pkg/version"
+)
+
+func TestResetHandler(t *testing.T) {
+	currentVersion := apimachineryversion.Info{
+		Major:      "1",
+		Minor:      "17",
+		GitVersion: "v1.17.1-alpha-1.12345",
+	}
+	registry := newKubeRegistry(currentVersion)
+	resetHandler := HandlerWithReset(registry, HandlerOpts{})
+	testCases := []struct {
+		desc         string
+		method       string
+		expectedBody string
+	}{
+		{
+			desc:         "Should return empty body on a get",
+			method:       http.MethodGet,
+			expectedBody: "",
+		},
+		{
+			desc:         "Should return 'metrics reset' in the body on a delete",
+			method:       http.MethodDelete,
+			expectedBody: "metrics reset\n",
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			req, err := http.NewRequest(tc.method, "http://sample.com/metrics", nil)
+			if err != nil {
+				t.Fatalf("Error creating http request")
+			}
+			rec := httptest.NewRecorder()
+			resetHandler.ServeHTTP(rec, req)
+			body, err := ioutil.ReadAll(rec.Result().Body)
+			if err != nil {
+				t.Fatalf("Error reading response body")
+			}
+			if string(body) != tc.expectedBody {
+				t.Errorf("Got '%s' as the response body, but want '%v'", body, tc.expectedBody)
+			}
+		})
+	}
+}

staging/src/k8s.io/component-base/metrics/legacyregistry/registry.go

@@ -46,6 +46,14 @@ func Handler() http.Handler {
 	return promhttp.InstrumentMetricHandler(prometheus.DefaultRegisterer, promhttp.HandlerFor(defaultRegistry, promhttp.HandlerOpts{}))
 }
 
+// HandlerWithReset returns an HTTP handler for the DefaultGatherer but invokes
+// registry reset if the http method is DELETE.
+func HandlerWithReset() http.Handler {
+	return promhttp.InstrumentMetricHandler(
+		prometheus.DefaultRegisterer,
+		metrics.HandlerWithReset(defaultRegistry, metrics.HandlerOpts{}))
+}
+
 // Register registers a collectable metric but uses the global registry
 func Register(c metrics.Registerable) error {
 	err := defaultRegistry.Register(c)
@@ -83,3 +91,8 @@ func CustomMustRegister(cs ...metrics.StableCollector) {
 		prometheus.MustRegister(c)
 	}
 }
+
+// Reset calls reset on the global registry
+func Reset() {
+	defaultRegistry.Reset()
+}

staging/src/k8s.io/component-base/metrics/registry.go

@@ -97,17 +97,30 @@ type Registerable interface {
 	FQName() string
 }
 
+type resettable interface {
+	Reset()
+}
+
 // KubeRegistry is an interface which implements a subset of prometheus.Registerer and
 // prometheus.Gatherer interfaces
 type KubeRegistry interface {
 	// Deprecated
 	RawMustRegister(...prometheus.Collector)
+	// CustomRegister is our internal variant of Prometheus registry.Register
 	CustomRegister(c StableCollector) error
+	// CustomMustRegister is our internal variant of Prometheus registry.MustRegister
 	CustomMustRegister(cs ...StableCollector)
+	// Register conforms to Prometheus registry.Register
 	Register(Registerable) error
+	// MustRegister conforms to Prometheus registry.MustRegister
 	MustRegister(...Registerable)
+	// Unregister conforms to Prometheus registry.Unregister
 	Unregister(collector Collector) bool
+	// Gather conforms to Prometheus gatherer.Gather
 	Gather() ([]*dto.MetricFamily, error)
+	// Reset invokes the Reset() function on all items in the registry
+	// which are added as resettables.
+	Reset()
 }
 
 // kubeRegistry is a wrapper around a prometheus registry-type object. Upon initialization
@@ -120,6 +133,8 @@ type kubeRegistry struct {
 	stableCollectors     []StableCollector       // stores all stable collector
 	hiddenCollectorsLock sync.RWMutex
 	stableCollectorsLock sync.RWMutex
+	resetLock            sync.RWMutex
+	resettables          []resettable
 }
 
 // Register registers a new Collector to be included in metrics
@@ -129,11 +144,11 @@ type kubeRegistry struct {
 // uniqueness criteria described in the documentation of metric.Desc.
 func (kr *kubeRegistry) Register(c Registerable) error {
 	if c.Create(&kr.version) {
+		defer kr.addResettable(c)
 		return kr.PromRegistry.Register(c)
 	}
 
 	kr.trackHiddenCollector(c)
-
 	return nil
 }
 
@@ -145,6 +160,7 @@ func (kr *kubeRegistry) MustRegister(cs ...Registerable) {
 	for _, c := range cs {
 		if c.Create(&kr.version) {
 			metrics = append(metrics, c)
+			kr.addResettable(c)
 		} else {
 			kr.trackHiddenCollector(c)
 		}
@@ -155,7 +171,7 @@ func (kr *kubeRegistry) MustRegister(cs ...Registerable) {
 // CustomRegister registers a new custom collector.
 func (kr *kubeRegistry) CustomRegister(c StableCollector) error {
 	kr.trackStableCollectors(c)
-
+	defer kr.addResettable(c)
 	if c.Create(&kr.version, c) {
 		return kr.PromRegistry.Register(c)
 	}
@@ -167,14 +183,13 @@ func (kr *kubeRegistry) CustomRegister(c StableCollector) error {
 // error.
 func (kr *kubeRegistry) CustomMustRegister(cs ...StableCollector) {
 	kr.trackStableCollectors(cs...)
-
 	collectors := make([]prometheus.Collector, 0, len(cs))
 	for _, c := range cs {
 		if c.Create(&kr.version, c) {
+			kr.addResettable(c)
 			collectors = append(collectors, c)
 		}
 	}
-
 	kr.PromRegistry.MustRegister(collectors...)
 }
 
@@ -185,6 +200,19 @@ func (kr *kubeRegistry) CustomMustRegister(cs ...StableCollector) {
 // Deprecated
 func (kr *kubeRegistry) RawMustRegister(cs ...prometheus.Collector) {
 	kr.PromRegistry.MustRegister(cs...)
+	for _, c := range cs {
+		kr.addResettable(c)
+	}
+}
+
+// addResettable will automatically add our metric to our reset
+// list if it satisfies the interface
+func (kr *kubeRegistry) addResettable(i interface{}) {
+	kr.resetLock.Lock()
+	defer kr.resetLock.Unlock()
+	if resettable, ok := i.(resettable); ok {
+		kr.resettables = append(kr.resettables, resettable)
+	}
 }
 
 // Unregister unregisters the Collector that equals the Collector passed
@@ -266,6 +294,15 @@ func (kr *kubeRegistry) enableHiddenStableCollectors() {
 	kr.CustomMustRegister(cs...)
 }
 
+// Reset invokes Reset on all metrics that are resettable.
+func (kr *kubeRegistry) Reset() {
+	kr.resetLock.RLock()
+	defer kr.resetLock.RUnlock()
+	for _, r := range kr.resettables {
+		r.Reset()
+	}
+}
+
 // BuildVersion is a helper function that can be easily mocked.
 var BuildVersion = version.Get
 
@@ -274,6 +311,7 @@ func newKubeRegistry(v apimachineryversion.Info) *kubeRegistry {
 		PromRegistry:     prometheus.NewRegistry(),
 		version:          parseVersion(v),
 		hiddenCollectors: make(map[string]Registerable),
+		resettables:      make([]resettable, 0),
 	}
 
 	registriesLock.Lock()
@@ -287,6 +325,5 @@ func newKubeRegistry(v apimachineryversion.Info) *kubeRegistry {
 // pre-registered.
 func NewKubeRegistry() KubeRegistry {
 	r := newKubeRegistry(BuildVersion())
-
 	return r
 }

staging/src/k8s.io/component-base/metrics/registry_test.go

@@ -471,3 +471,47 @@ func TestEnableHiddenStableCollector(t *testing.T) {
 		})
 	}
 }
+
+func TestRegistryReset(t *testing.T) {
+	currentVersion := apimachineryversion.Info{
+		Major:      "1",
+		Minor:      "17",
+		GitVersion: "v1.17.1-alpha-1.12345",
+	}
+	registry := newKubeRegistry(currentVersion)
+	resettableMetric := NewCounterVec(&CounterOpts{
+		Name: "reset_metric",
+		Help: "this metric can be reset",
+	}, []string{"label"})
+	// gauges cannot be reset
+	nonResettableMetric := NewGauge(&GaugeOpts{
+		Name: "not_reset_metric",
+		Help: "this metric cannot be reset",
+	})
+
+	registry.MustRegister(resettableMetric)
+	registry.MustRegister(nonResettableMetric)
+	resettableMetric.WithLabelValues("one").Inc()
+	resettableMetric.WithLabelValues("two").Inc()
+	resettableMetric.WithLabelValues("two").Inc()
+	nonResettableMetric.Inc()
+
+	nonResettableOutput := `
+        # HELP not_reset_metric [ALPHA] this metric cannot be reset
+        # TYPE not_reset_metric gauge
+        not_reset_metric 1
+`
+	resettableOutput := `
+        # HELP reset_metric [ALPHA] this metric can be reset
+        # TYPE reset_metric counter
+        reset_metric{label="one"} 1
+        reset_metric{label="two"} 2
+`
+	if err := testutil.GatherAndCompare(registry, strings.NewReader(nonResettableOutput+resettableOutput), "reset_metric", "not_reset_metric"); err != nil {
+		t.Fatal(err)
+	}
+	registry.Reset()
+	if err := testutil.GatherAndCompare(registry, strings.NewReader(nonResettableOutput), "reset_metric", "not_reset_metric"); err != nil {
+		t.Fatal(err)
+	}
+}