client/resource_group: cache request source RU metrics

YuhaoZhang00 · YuhaoZhang00 · commit 1dd52fd31e23 · 2026-04-09T11:14:13.000+08:00
diff --git a/client/resource_group/controller/global_controller.go b/client/resource_group/controller/global_controller.go
@@ -337,6 +337,7 @@ func (c *ResourceGroupsController) Start(ctx context.Context) {
 			/* channels */
 			case <-c.loopCtx.Done():
 				metrics.ResourceGroupStatusGauge.Reset()
+				metrics.RequestSourceRUCounter.Reset()
 				return
 			case <-c.responseDeadlineCh:
 				c.run.inDegradedMode.Store(true)
@@ -583,6 +584,7 @@ func (c *ResourceGroupsController) cleanUpResourceGroup() {
 		gc.mu.Unlock()
 		if equalRU(latestConsumption, *gc.run.consumption) {
 			if gc.inactive || gc.tombstone.Load() {
+				gc.metrics.cleanupRequestSourceMetrics(resourceGroupName)
 				c.groupsController.Delete(resourceGroupName)
 				metrics.ResourceGroupStatusGauge.DeleteLabelValues(resourceGroupName, resourceGroupName)
 				return true
diff --git a/client/resource_group/controller/group_controller.go b/client/resource_group/controller/group_controller.go
@@ -106,6 +106,13 @@ type groupMetricsCollection struct {
 	tokenRequestCounter               prometheus.Counter
 	runningKVRequestCounter           prometheus.Gauge
 	consumeTokenHistogram             prometheus.Observer
+	sourceMetricsMu                   sync.RWMutex
+	sourceMetrics                     map[string]*requestSourceMetrics
+}
+
+type requestSourceMetrics struct {
+	rru prometheus.Counter
+	wru prometheus.Counter
 }
 
 func initMetrics(oldName, name string) *groupMetricsCollection {
@@ -122,6 +129,52 @@ func initMetrics(oldName, name string) *groupMetricsCollection {
 		tokenRequestCounter:               metrics.ResourceGroupTokenRequestCounter.WithLabelValues(oldName, name),
 		runningKVRequestCounter:           metrics.GroupRunningKVRequestCounter.WithLabelValues(name),
 		consumeTokenHistogram:             metrics.TokenConsumedHistogram.WithLabelValues(name),
+		sourceMetrics:                     make(map[string]*requestSourceMetrics),
+	}
+}
+
+func (mc *groupMetricsCollection) getOrCreateRequestSourceMetrics(resourceGroupName, requestSource string) *requestSourceMetrics {
+	mc.sourceMetricsMu.RLock()
+	sourceMetrics, ok := mc.sourceMetrics[requestSource]
+	mc.sourceMetricsMu.RUnlock()
+	if ok {
+		return sourceMetrics
+	}
+
+	mc.sourceMetricsMu.Lock()
+	defer mc.sourceMetricsMu.Unlock()
+	sourceMetrics, ok = mc.sourceMetrics[requestSource]
+	if ok {
+		return sourceMetrics
+	}
+	sourceMetrics = &requestSourceMetrics{
+		rru: metrics.RequestSourceRUCounter.WithLabelValues(resourceGroupName, requestSource, "rru"),
+		wru: metrics.RequestSourceRUCounter.WithLabelValues(resourceGroupName, requestSource, "wru"),
+	}
+	mc.sourceMetrics[requestSource] = sourceMetrics
+	return sourceMetrics
+}
+
+func (mc *groupMetricsCollection) addRequestSourceRU(resourceGroupName, requestSource string, consumption *rmpb.Consumption) {
+	if consumption == nil {
+		return
+	}
+	sourceMetrics := mc.getOrCreateRequestSourceMetrics(resourceGroupName, requestSource)
+	if consumption.RRU > 0 {
+		sourceMetrics.rru.Add(consumption.RRU)
+	}
+	if consumption.WRU > 0 {
+		sourceMetrics.wru.Add(consumption.WRU)
+	}
+}
+
+func (mc *groupMetricsCollection) cleanupRequestSourceMetrics(resourceGroupName string) {
+	mc.sourceMetricsMu.Lock()
+	defer mc.sourceMetricsMu.Unlock()
+	for requestSource := range mc.sourceMetrics {
+		metrics.RequestSourceRUCounter.DeleteLabelValues(resourceGroupName, requestSource, "rru")
+		metrics.RequestSourceRUCounter.DeleteLabelValues(resourceGroupName, requestSource, "wru")
+		delete(mc.sourceMetrics, requestSource)
 	}
 }
 
@@ -577,6 +630,8 @@ func (gc *groupCostController) onRequestWaitImpl(
 		waitDuration += d
 	}
 
+	gc.metrics.addRequestSourceRU(gc.name, info.RequestSource(), delta)
+
 	gc.mu.Lock()
 	// Calculate the penalty of the store
 	penalty = &rmpb.Consumption{}
@@ -622,6 +677,8 @@ func (gc *groupCostController) onResponseImpl(
 	add(gc.mu.globalCounter, count)
 	gc.mu.Unlock()
 
+	gc.metrics.addRequestSourceRU(gc.name, req.RequestSource(), delta)
+
 	return delta, nil
 }
 
@@ -663,6 +720,8 @@ func (gc *groupCostController) onResponseWaitImpl(
 	add(gc.mu.globalCounter, count)
 	gc.mu.Unlock()
 
+	gc.metrics.addRequestSourceRU(gc.name, req.RequestSource(), delta)
+
 	return delta, waitDuration, nil
 }
 
diff --git a/client/resource_group/controller/metrics/metrics.go b/client/resource_group/controller/metrics/metrics.go
@@ -49,6 +49,8 @@ var (
 	LowTokenRequestNotifyCounter *prometheus.CounterVec
 	// TokenConsumedHistogram comments placeholder
 	TokenConsumedHistogram *prometheus.HistogramVec
+	// RequestSourceRUCounter comments placeholder
+	RequestSourceRUCounter *prometheus.CounterVec
 	// FailedTokenRequestDuration comments placeholder, WithLabelValues is a heavy operation, define variable to avoid call it every time.
 	FailedTokenRequestDuration prometheus.Observer
 	// SuccessfulTokenRequestDuration comments placeholder, WithLabelValues is a heavy operation, define variable to avoid call it every time.
@@ -153,6 +155,15 @@ func initMetrics(constLabels prometheus.Labels) {
 			ConstLabels: constLabels,
 		}, []string{newResourceGroupNameLabel})
 
+	RequestSourceRUCounter = prometheus.NewCounterVec(
+		prometheus.CounterOpts{
+			Namespace:   namespace,
+			Subsystem:   requestSubsystem,
+			Name:        "ru_total",
+			Help:        "Counter of request RU consumption grouped by resource group and request source.",
+			ConstLabels: constLabels,
+		}, []string{newResourceGroupNameLabel, "request_source", errType})
+
 	// WithLabelValues is a heavy operation, define variable to avoid call it every time.
 	FailedTokenRequestDuration = TokenRequestDuration.WithLabelValues("fail")
 	SuccessfulTokenRequestDuration = TokenRequestDuration.WithLabelValues("success")
@@ -171,4 +182,5 @@ func InitAndRegisterMetrics(constLabels prometheus.Labels) {
 	prometheus.MustRegister(ResourceGroupTokenRequestCounter)
 	prometheus.MustRegister(LowTokenRequestNotifyCounter)
 	prometheus.MustRegister(TokenConsumedHistogram)
+	prometheus.MustRegister(RequestSourceRUCounter)
 }
diff --git a/client/resource_group/controller/model.go b/client/resource_group/controller/model.go
@@ -52,6 +52,7 @@ type RequestInfo interface {
 	StoreID() uint64
 	RequestSize() uint64
 	AccessLocationType() AccessLocationType
+	RequestSource() string
 }
 
 // ResponseInfo is the interface of the response information provider. A response should be
diff --git a/client/resource_group/controller/request_source_metrics_test.go b/client/resource_group/controller/request_source_metrics_test.go
@@ -0,0 +1,137 @@
+package controller
+
+import (
+	"context"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus"
+	dto "github.com/prometheus/client_model/go"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+
+	rmpb "github.com/pingcap/kvproto/pkg/resource_manager"
+
+	controllerMetrics "github.com/tikv/pd/client/resource_group/controller/metrics"
+)
+
+func counterValue(t *testing.T, metric prometheus.Counter) float64 {
+	t.Helper()
+	pb := &dto.Metric{}
+	require.NoError(t, metric.Write(pb))
+	return pb.GetCounter().GetValue()
+}
+
+func collectorMetricCount(collector prometheus.Collector) int {
+	ch := make(chan prometheus.Metric, 8)
+	go func() {
+		collector.Collect(ch)
+		close(ch)
+	}()
+	count := 0
+	for range ch {
+		count++
+	}
+	return count
+}
+
+func TestRequestSourceMetricsCachedByResourceGroup(t *testing.T) {
+	re := require.New(t)
+	gc := createTestGroupCostController(re)
+	req := &TestRequestInfo{
+		isWrite:       true,
+		writeBytes:    100,
+		numReplicas:   1,
+		storeID:       1,
+		accessType:    AccessUnknown,
+		requestSource: "internal_gc_test",
+	}
+	resp := &TestResponseInfo{
+		readBytes: 128,
+		succeed:   true,
+	}
+
+	beforeCount := collectorMetricCount(controllerMetrics.RequestSourceRUCounter)
+
+	reqConsumption, _, _, _, err := gc.onRequestWaitImpl(context.Background(), req)
+	re.NoError(err)
+	respConsumption, err := gc.onResponseImpl(req, resp)
+	re.NoError(err)
+	re.NotZero(reqConsumption.WRU)
+	re.NotZero(respConsumption.RRU)
+
+	gc.metrics.sourceMetricsMu.RLock()
+	sourceMetrics := gc.metrics.sourceMetrics[req.requestSource]
+	cacheSize := len(gc.metrics.sourceMetrics)
+	gc.metrics.sourceMetricsMu.RUnlock()
+
+	re.Equal(1, cacheSize)
+	re.NotNil(sourceMetrics)
+	re.Equal(reqConsumption.WRU, counterValue(t, sourceMetrics.wru))
+	re.Equal(respConsumption.RRU, counterValue(t, sourceMetrics.rru))
+	re.Equal(beforeCount+2, collectorMetricCount(controllerMetrics.RequestSourceRUCounter))
+
+	_, _, _, _, err = gc.onRequestWaitImpl(context.Background(), req)
+	re.NoError(err)
+	gc.metrics.sourceMetricsMu.RLock()
+	re.Equal(1, len(gc.metrics.sourceMetrics))
+	re.Same(sourceMetrics, gc.metrics.sourceMetrics[req.requestSource])
+	gc.metrics.sourceMetricsMu.RUnlock()
+
+	controllerMetrics.RequestSourceRUCounter.DeleteLabelValues(gc.name, req.requestSource, "rru")
+	controllerMetrics.RequestSourceRUCounter.DeleteLabelValues(gc.name, req.requestSource, "wru")
+}
+
+func TestCleanupResourceGroupRemovesRequestSourceMetrics(t *testing.T) {
+	re := require.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	mockProvider := newMockResourceGroupProvider()
+	controller, err := NewResourceGroupController(ctx, 1, mockProvider, nil, 0)
+	re.NoError(err)
+
+	group := &rmpb.ResourceGroup{
+		Name: "request-source-cleanup",
+		Mode: rmpb.GroupMode_RUMode,
+		RUSettings: &rmpb.GroupRequestUnitSettings{
+			RU: &rmpb.TokenBucket{
+				Settings: &rmpb.TokenLimitSettings{FillRate: 1000000},
+			},
+		},
+	}
+	mockProvider.On("GetResourceGroup", mock.Anything, group.Name, mock.Anything).Return(group, nil)
+
+	gc, err := controller.tryGetResourceGroupController(ctx, group.Name, false)
+	re.NoError(err)
+
+	req := &TestRequestInfo{
+		isWrite:       true,
+		writeBytes:    64,
+		numReplicas:   1,
+		storeID:       1,
+		requestSource: "internal_gc_cleanup",
+	}
+	resp := &TestResponseInfo{readBytes: 64, succeed: true}
+	beforeCount := collectorMetricCount(controllerMetrics.RequestSourceRUCounter)
+
+	_, _, _, _, err = gc.onRequestWaitImpl(context.Background(), req)
+	re.NoError(err)
+	_, err = gc.onResponseImpl(req, resp)
+	re.NoError(err)
+
+	gc.mu.Lock()
+	*gc.run.consumption = *gc.mu.consumption
+	gc.mu.Unlock()
+	gc.inactive = true
+
+	re.Greater(collectorMetricCount(controllerMetrics.RequestSourceRUCounter), beforeCount)
+
+	controller.cleanUpResourceGroup()
+
+	_, ok := controller.loadGroupController(group.Name)
+	re.False(ok)
+	gc.metrics.sourceMetricsMu.RLock()
+	re.Empty(gc.metrics.sourceMetrics)
+	gc.metrics.sourceMetricsMu.RUnlock()
+	re.Equal(beforeCount, collectorMetricCount(controllerMetrics.RequestSourceRUCounter))
+}
diff --git a/client/resource_group/controller/testutil.go b/client/resource_group/controller/testutil.go
@@ -22,11 +22,12 @@ import "time"
 
 // TestRequestInfo is used to test the request info interface.
 type TestRequestInfo struct {
-	isWrite     bool
-	writeBytes  uint64
-	numReplicas int64
-	storeID     uint64
-	accessType  AccessLocationType
+	isWrite       bool
+	writeBytes    uint64
+	numReplicas   int64
+	storeID       uint64
+	accessType    AccessLocationType
+	requestSource string
 }
 
 // NewTestRequestInfo creates a new TestRequestInfo.
@@ -70,6 +71,11 @@ func (tri *TestRequestInfo) AccessLocationType() AccessLocationType {
 	return tri.accessType
 }
 
+// RequestSource implements the RequestInfo interface.
+func (tri *TestRequestInfo) RequestSource() string {
+	return tri.requestSource
+}
+
 // TestResponseInfo is used to test the response info interface.
 type TestResponseInfo struct {
 	readBytes uint64

Original file line number	Diff line number	Diff line change
`@@ -52,6 +52,7 @@ type RequestInfo interface {`
`52`	`52`	`StoreID() uint64`
`53`	`53`	`RequestSize() uint64`
`54`	`54`	`AccessLocationType() AccessLocationType`
	`55`	`+ RequestSource() string`
`55`	`56`	`}`
`56`	`57`
`57`	`58`	`// ResponseInfo is the interface of the response information provider. A response should be`