client/resource_group: cache request source RU metrics

client/resource_group/controller/global_controller.go

@@ -142,11 +142,12 @@ var _ ResourceGroupKVInterceptor = (*ResourceGroupsController)(nil)
 
 // ResourceGroupsController implements ResourceGroupKVInterceptor.
 type ResourceGroupsController struct {
-	clientUniqueID   uint64
-	provider         ResourceGroupProvider
-	groupsController sync.Map
-	ruConfig         *RUConfig
-	keyspaceID       uint32
+	clientUniqueID      uint64
+	provider            ResourceGroupProvider
+	groupsController    sync.Map
+	requestSourceStates sync.Map
+	ruConfig            *RUConfig
+	keyspaceID          uint32
 
 	loopCtx    context.Context
 	loopCancel func()
@@ -337,6 +338,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)
@@ -387,7 +389,13 @@ func (c *ResourceGroupsController) Start(ctx context.Context) {
 							continue
 						}
 						// If the resource group is marked as tombstone before, re-create the resource group controller.
-						newGC, err := newGroupCostController(group, c.ruConfig, c.lowTokenNotifyChan, c.tokenBucketUpdateChan)
+						newGC, err := newGroupCostController(
+							group,
+							c.ruConfig,
+							c.lowTokenNotifyChan,
+							c.tokenBucketUpdateChan,
+							c.getOrCreateRequestSourceMetricsState(name),
+						)
 						if err != nil {
 							log.Warn("[resource group controller] re-create resource group cost controller for tombstone failed",
 								zap.String("name", name), zap.Error(err))
@@ -472,6 +480,22 @@ func (c *ResourceGroupsController) loadOrStoreGroupController(name string, gc *g
 	return tmp.(*groupCostController), loaded
 }
 
+func (c *ResourceGroupsController) getOrCreateRequestSourceMetricsState(name string) *requestSourceMetricsState {
+	if state, ok := c.requestSourceStates.Load(name); ok {
+		return state.(*requestSourceMetricsState)
+	}
+	state := newRequestSourceMetricsState(name)
+	actual, _ := c.requestSourceStates.LoadOrStore(name, state)
+	return actual.(*requestSourceMetricsState)
+}
+
+func (c *ResourceGroupsController) cleanupRequestSourceMetricsState(name string) {
+	if state, ok := c.requestSourceStates.Load(name); ok {
+		state.(*requestSourceMetricsState).cleanup()
+		c.requestSourceStates.Delete(name)
+	}
+}
+
 // NewResourceGroupNotExistErr returns a new error that indicates the resource group does not exist.
 // It's exported for testing.
 func NewResourceGroupNotExistErr(name string) error {
@@ -521,7 +545,13 @@ func (c *ResourceGroupsController) tryGetResourceGroupController(
 		return gc, nil
 	}
 	// Initialize the resource group controller.
-	gc, err = newGroupCostController(group, c.ruConfig, c.lowTokenNotifyChan, c.tokenBucketUpdateChan)
+	gc, err = newGroupCostController(
+		group,
+		c.ruConfig,
+		c.lowTokenNotifyChan,
+		c.tokenBucketUpdateChan,
+		c.getOrCreateRequestSourceMetricsState(name),
+	)
 	if err != nil {
 		return nil, err
 	}
@@ -553,15 +583,23 @@ func (c *ResourceGroupsController) tombstoneGroupCostController(name string) {
 		log.Warn("[resource group controller] get default resource group meta for tombstone failed",
 			zap.String("name", name), zap.Error(err))
 		// Directly delete the resource group controller if the default group is not available.
+		c.cleanupRequestSourceMetricsState(name)
 		c.groupsController.Delete(name)
 		return
 	}
 	// Create a default resource group controller for the tombstone resource group independently.
-	gc, err := newGroupCostController(defaultGC.getMeta(), c.ruConfig, c.lowTokenNotifyChan, c.tokenBucketUpdateChan)
+	gc, err := newGroupCostController(
+		defaultGC.getMeta(),
+		c.ruConfig,
+		c.lowTokenNotifyChan,
+		c.tokenBucketUpdateChan,
+		c.getOrCreateRequestSourceMetricsState(name),
+	)
 	if err != nil {
 		log.Warn("[resource group controller] create default resource group cost controller for tombstone failed",
 			zap.String("name", name), zap.Error(err))
 		// Directly delete the resource group controller if the default group controller cannot be created.
+		c.cleanupRequestSourceMetricsState(name)
 		c.groupsController.Delete(name)
 		return
 	}
@@ -583,7 +621,9 @@ func (c *ResourceGroupsController) cleanUpResourceGroup() {
 		gc.mu.Unlock()
 		if equalRU(latestConsumption, *gc.run.consumption) {
 			if gc.inactive || gc.tombstone.Load() {
+				gc.metrics.cleanupRequestSourceMetrics()
 				c.groupsController.Delete(resourceGroupName)
+				c.requestSourceStates.Delete(resourceGroupName)
 				metrics.ResourceGroupStatusGauge.DeleteLabelValues(resourceGroupName, resourceGroupName)
 				return true
 			}

client/resource_group/controller/group_controller.go

@@ -106,9 +106,43 @@ type groupMetricsCollection struct {
 	tokenRequestCounter               prometheus.Counter
 	runningKVRequestCounter           prometheus.Gauge
 	consumeTokenHistogram             prometheus.Observer
+	sourceState                       *requestSourceMetricsState
 }
 
-func initMetrics(oldName, name string) *groupMetricsCollection {
+type requestSourceMetrics struct {
+	rru prometheus.Counter
+	wru prometheus.Counter
+}
+
+type requestSourceMetricsState struct {
+	resourceGroupName string
+	mu                sync.RWMutex
+	closed            bool
+	items             map[string]*requestSourceMetrics
+}
+
+func newRequestSourceMetricsState(resourceGroupName string) *requestSourceMetricsState {
+	return &requestSourceMetricsState{
+		resourceGroupName: resourceGroupName,
+		items:             make(map[string]*requestSourceMetrics),
+	}
+}
+
+func (s *requestSourceMetricsState) cleanup() {
+	if s == nil {
+		return
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.closed = true
+	for requestSource := range s.items {
+		metrics.RequestSourceRUCounter.DeleteLabelValues(s.resourceGroupName, requestSource, "rru")
+		metrics.RequestSourceRUCounter.DeleteLabelValues(s.resourceGroupName, requestSource, "wru")
+		delete(s.items, requestSource)
+	}
+}
+
+func initMetrics(oldName, name string, sourceState *requestSourceMetricsState) *groupMetricsCollection {
 	const (
 		otherType     = "others"
 		throttledType = "throttled"
@@ -122,7 +156,60 @@ func initMetrics(oldName, name string) *groupMetricsCollection {
 		tokenRequestCounter:               metrics.ResourceGroupTokenRequestCounter.WithLabelValues(oldName, name),
 		runningKVRequestCounter:           metrics.GroupRunningKVRequestCounter.WithLabelValues(name),
 		consumeTokenHistogram:             metrics.TokenConsumedHistogram.WithLabelValues(name),
+		sourceState:                       sourceState,
+	}
+}
+
+func (mc *groupMetricsCollection) getOrCreateRequestSourceMetrics(requestSource string) *requestSourceMetrics {
+	if mc.sourceState == nil {
+		return nil
+	}
+	mc.sourceState.mu.RLock()
+	sourceMetrics, ok := mc.sourceState.items[requestSource]
+	closed := mc.sourceState.closed
+	mc.sourceState.mu.RUnlock()
+	if ok {
+		return sourceMetrics
+	}
+	if closed {
+		return nil
 	}
+
+	mc.sourceState.mu.Lock()
+	defer mc.sourceState.mu.Unlock()
+	if mc.sourceState.closed {
+		return nil
+	}
+	sourceMetrics, ok = mc.sourceState.items[requestSource]
+	if ok {
+		return sourceMetrics
+	}
+	sourceMetrics = &requestSourceMetrics{
+		rru: metrics.RequestSourceRUCounter.WithLabelValues(mc.sourceState.resourceGroupName, requestSource, "rru"),
+		wru: metrics.RequestSourceRUCounter.WithLabelValues(mc.sourceState.resourceGroupName, requestSource, "wru"),
+	}
+	mc.sourceState.items[requestSource] = sourceMetrics
+	return sourceMetrics
+}
+
+func (mc *groupMetricsCollection) addRequestSourceRU(requestSource string, consumption *rmpb.Consumption) {
+	if consumption == nil {
+		return
+	}
+	sourceMetrics := mc.getOrCreateRequestSourceMetrics(requestSource)
+	if sourceMetrics == nil {
+		return
+	}
+	if consumption.RRU > 0 {
+		sourceMetrics.rru.Add(consumption.RRU)
+	}
+	if consumption.WRU > 0 {
+		sourceMetrics.wru.Add(consumption.WRU)
+	}
+}
+
+func (mc *groupMetricsCollection) cleanupRequestSourceMetrics() {
+	mc.sourceState.cleanup()
 }
 
 type tokenCounter struct {
@@ -156,6 +243,7 @@ func newGroupCostController(
 	mainCfg *RUConfig,
 	lowRUNotifyChan chan notifyMsg,
 	tokenBucketUpdateChan chan *groupCostController,
+	sourceState *requestSourceMetricsState,
 ) (*groupCostController, error) {
 	switch group.Mode {
 	case rmpb.GroupMode_RUMode:
@@ -165,7 +253,10 @@ func newGroupCostController(
 	default:
 		return nil, errs.ErrClientResourceGroupConfigUnavailable.FastGenByArgs("not supports the resource type")
 	}
-	ms := initMetrics(group.Name, group.Name)
+	if sourceState == nil {
+		sourceState = newRequestSourceMetricsState(group.Name)
+	}
+	ms := initMetrics(group.Name, group.Name, sourceState)
 	gc := &groupCostController{
 		meta:    group,
 		name:    group.Name,
@@ -577,6 +668,8 @@ func (gc *groupCostController) onRequestWaitImpl(
 		waitDuration += d
 	}
 
+	gc.metrics.addRequestSourceRU(info.RequestSource(), delta)
+
 	gc.mu.Lock()
 	// Calculate the penalty of the store
 	penalty = &rmpb.Consumption{}
@@ -622,6 +715,8 @@ func (gc *groupCostController) onResponseImpl(
 	add(gc.mu.globalCounter, count)
 	gc.mu.Unlock()
 
+	gc.metrics.addRequestSourceRU(req.RequestSource(), delta)
+
 	return delta, nil
 }
 
@@ -663,6 +758,8 @@ func (gc *groupCostController) onResponseWaitImpl(
 	add(gc.mu.globalCounter, count)
 	gc.mu.Unlock()
 
+	gc.metrics.addRequestSourceRU(req.RequestSource(), delta)
+
 	return delta, waitDuration, nil
 }
 

client/resource_group/controller/group_controller_test.go

@@ -45,7 +45,7 @@ func createTestGroupCostController(re *require.Assertions) *groupCostController
 	}
 	ch1 := make(chan notifyMsg)
 	ch2 := make(chan *groupCostController)
-	gc, err := newGroupCostController(group, DefaultRUConfig(), ch1, ch2)
+	gc, err := newGroupCostController(group, DefaultRUConfig(), ch1, ch2, nil)
 	re.NoError(err)
 	return gc
 }
@@ -239,7 +239,7 @@ func TestAcquireTokensSignalAwareWait(t *testing.T) {
 	cfg := DefaultRUConfig()
 	cfg.WaitRetryInterval = 5 * time.Second
 	cfg.WaitRetryTimes = 3
-	gc, err := newGroupCostController(group, cfg, notifyCh, make(chan *groupCostController, 1))
+	gc, err := newGroupCostController(group, cfg, notifyCh, make(chan *groupCostController, 1), nil)
 	re.NoError(err)
 
 	// Set fillRate=0 so reservation always fails with InfDuration,

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)
 }

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