From 2a1f2dc98f514debea9f6712b9ff54e35c08cd2d Mon Sep 17 00:00:00 2001
From: Artem Danilov <artem_danilov@airbnb.com>
Date: Mon, 27 Jan 2025 00:21:59 -0800
Subject: [PATCH] client: fix pd client metrics registration (#8994)

ref tikv/pd#8678

Signed-off-by: artem_danilov <artem_danilov@airbnb.com>

Co-authored-by: artem_danilov <artem_danilov@airbnb.com>
Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>
---
 client/metrics/metrics.go                    | 24 ++++++++++++++++++++
 client/pkg/circuitbreaker/circuit_breaker.go | 10 ++++++--
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/client/metrics/metrics.go b/client/metrics/metrics.go
index 67268c826f5..da7637b19be 100644
--- a/client/metrics/metrics.go
+++ b/client/metrics/metrics.go
@@ -15,6 +15,7 @@
 package metrics
 
 import (
+	"sync"
 	"sync/atomic"
 
 	"github.com/prometheus/client_golang/prometheus"
@@ -26,6 +27,28 @@ var initialized int32
 func init() {
 	initMetrics(prometheus.Labels{})
 	initCmdDurations()
+	initRegisteredConsumers()
+}
+
+var consumersInitializers = struct {
+	sync.Mutex
+	value []func()
+}{}
+
+// RegisterConsumer registers a consumer to be initialized when the metrics are (re)initialized
+func RegisterConsumer(initConsumer func()) {
+	consumersInitializers.Lock()
+	defer consumersInitializers.Unlock()
+	consumersInitializers.value = append(consumersInitializers.value, initConsumer)
+	initConsumer()
+}
+
+func initRegisteredConsumers() {
+	consumersInitializers.Lock()
+	defer consumersInitializers.Unlock()
+	for _, initConsumer := range consumersInitializers.value {
+		initConsumer()
+	}
 }
 
 // InitAndRegisterMetrics initializes and registers the metrics manually.
@@ -34,6 +57,7 @@ func InitAndRegisterMetrics(constLabels prometheus.Labels) {
 		// init metrics with constLabels
 		initMetrics(constLabels)
 		initCmdDurations()
+		initRegisteredConsumers()
 		// register metrics
 		registerMetrics()
 	}
diff --git a/client/pkg/circuitbreaker/circuit_breaker.go b/client/pkg/circuitbreaker/circuit_breaker.go
index 7466124f2bc..4621dfe4bc0 100644
--- a/client/pkg/circuitbreaker/circuit_breaker.go
+++ b/client/pkg/circuitbreaker/circuit_breaker.go
@@ -110,12 +110,18 @@ func NewCircuitBreaker(name string, st Settings) *CircuitBreaker {
 	cb.config = &st
 	cb.state = cb.newState(time.Now(), StateClosed)
 
-	metricName := replacer.Replace(name)
+	m.RegisterConsumer(func() {
+		registerMetrics(cb)
+	})
+	return cb
+}
+
+func registerMetrics(cb *CircuitBreaker) {
+	metricName := replacer.Replace(cb.name)
 	cb.successCounter = m.CircuitBreakerCounters.WithLabelValues(metricName, "success")
 	cb.errorCounter = m.CircuitBreakerCounters.WithLabelValues(metricName, "error")
 	cb.overloadCounter = m.CircuitBreakerCounters.WithLabelValues(metricName, "overload")
 	cb.fastFailCounter = m.CircuitBreakerCounters.WithLabelValues(metricName, "fast_fail")
-	return cb
 }
 
 // ChangeSettings changes the CircuitBreaker settings.