From c75ec1a24404081a8a769437d7b38542a437c59f Mon Sep 17 00:00:00 2001
From: RainbowMango <renhongcai@huawei.com>
Date: Tue, 24 Sep 2019 21:05:16 +0800
Subject: [PATCH] Provides register apis for custom collector

---
 .../metrics/legacyregistry/registry.go        | 19 ++++++++++++++
 .../k8s.io/component-base/metrics/registry.go | 25 +++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/staging/src/k8s.io/component-base/metrics/legacyregistry/registry.go b/staging/src/k8s.io/component-base/metrics/legacyregistry/registry.go
index 267706c4b25..7f9609f71bd 100644
--- a/staging/src/k8s.io/component-base/metrics/legacyregistry/registry.go
+++ b/staging/src/k8s.io/component-base/metrics/legacyregistry/registry.go
@@ -85,3 +85,22 @@ func RawRegister(c prometheus.Collector) error {
 	prometheus.Register(c)
 	return err
 }
+
+// CustomRegister registers a custom collector but uses the global registry.
+func CustomRegister(c metrics.StableCollector) error {
+	err := defaultRegistry.CustomRegister(c)
+
+	//TODO(RainbowMango): Maybe we can wrap this error by error wrapping.(Golang 1.13)
+	_ = prometheus.Register(c)
+
+	return err
+}
+
+// CustomMustRegister registers custom collectors but uses the global registry.
+func CustomMustRegister(cs ...metrics.StableCollector) {
+	defaultRegistry.CustomMustRegister(cs...)
+
+	for _, c := range cs {
+		prometheus.MustRegister(c)
+	}
+}
diff --git a/staging/src/k8s.io/component-base/metrics/registry.go b/staging/src/k8s.io/component-base/metrics/registry.go
index ff8d5ac02ba..35e878997db 100644
--- a/staging/src/k8s.io/component-base/metrics/registry.go
+++ b/staging/src/k8s.io/component-base/metrics/registry.go
@@ -78,6 +78,8 @@ type KubeRegistry interface {
 	RawRegister(prometheus.Collector) error
 	// Deprecated
 	RawMustRegister(...prometheus.Collector)
+	CustomRegister(c StableCollector) error
+	CustomMustRegister(cs ...StableCollector)
 	Register(Registerable) error
 	MustRegister(...Registerable)
 	Unregister(Registerable) bool
@@ -117,6 +119,29 @@ func (kr *kubeRegistry) MustRegister(cs ...Registerable) {
 	kr.PromRegistry.MustRegister(metrics...)
 }
 
+// CustomRegister registers a new custom collector.
+func (kr *kubeRegistry) CustomRegister(c StableCollector) error {
+	if c.Create(&kr.version, c) {
+		return kr.PromRegistry.Register(c)
+	}
+
+	return nil
+}
+
+// CustomMustRegister works like CustomRegister but registers any number of
+// StableCollectors and panics upon the first registration that causes an
+// error.
+func (kr *kubeRegistry) CustomMustRegister(cs ...StableCollector) {
+	collectors := make([]prometheus.Collector, 0, len(cs))
+	for _, c := range cs {
+		if c.Create(&kr.version, c) {
+			collectors = append(collectors, c)
+		}
+	}
+
+	kr.PromRegistry.MustRegister(collectors...)
+}
+
 // RawRegister takes a native prometheus.Collector and registers the collector
 // to the registry. This bypasses metrics safety checks, so should only be used
 // to register custom prometheus collectors.