client/resource_group: keep request-source metrics state per group

Signed-off-by: Yuhao Zhang <yhzhang00@outlook.com>

Author: YuhaoZhang00

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()
@@ -388,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))
@@ -473,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 {
@@ -522,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
 	}
@@ -554,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
 	}
@@ -584,8 +621,9 @@ func (c *ResourceGroupsController) cleanUpResourceGroup() {
 		gc.mu.Unlock()
 		if equalRU(latestConsumption, *gc.run.consumption) {
 			if gc.inactive || gc.tombstone.Load() {
-				gc.metrics.cleanupRequestSourceMetrics(resourceGroupName)
+				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,16 +106,43 @@ type groupMetricsCollection struct {
 	tokenRequestCounter               prometheus.Counter
 	runningKVRequestCounter           prometheus.Gauge
 	consumeTokenHistogram             prometheus.Observer
-	sourceMetricsMu                   sync.RWMutex
-	sourceMetrics                     map[string]*requestSourceMetrics
+	sourceState                       *requestSourceMetricsState
 }
 
 type requestSourceMetrics struct {
 	rru prometheus.Counter
 	wru prometheus.Counter
 }
 
-func initMetrics(oldName, name string) *groupMetricsCollection {
+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"
@@ -129,37 +156,50 @@ 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),
+		sourceState:                       sourceState,
 	}
 }
 
-func (mc *groupMetricsCollection) getOrCreateRequestSourceMetrics(resourceGroupName, requestSource string) *requestSourceMetrics {
-	mc.sourceMetricsMu.RLock()
-	sourceMetrics, ok := mc.sourceMetrics[requestSource]
-	mc.sourceMetricsMu.RUnlock()
+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.sourceMetricsMu.Lock()
-	defer mc.sourceMetricsMu.Unlock()
-	sourceMetrics, ok = mc.sourceMetrics[requestSource]
+	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(resourceGroupName, requestSource, "rru"),
-		wru: metrics.RequestSourceRUCounter.WithLabelValues(resourceGroupName, requestSource, "wru"),
+		rru: metrics.RequestSourceRUCounter.WithLabelValues(mc.sourceState.resourceGroupName, requestSource, "rru"),
+		wru: metrics.RequestSourceRUCounter.WithLabelValues(mc.sourceState.resourceGroupName, requestSource, "wru"),
 	}
-	mc.sourceMetrics[requestSource] = sourceMetrics
+	mc.sourceState.items[requestSource] = sourceMetrics
 	return sourceMetrics
 }
 
-func (mc *groupMetricsCollection) addRequestSourceRU(resourceGroupName, requestSource string, consumption *rmpb.Consumption) {
+func (mc *groupMetricsCollection) addRequestSourceRU(requestSource string, consumption *rmpb.Consumption) {
 	if consumption == nil {
 		return
 	}
-	sourceMetrics := mc.getOrCreateRequestSourceMetrics(resourceGroupName, requestSource)
+	sourceMetrics := mc.getOrCreateRequestSourceMetrics(requestSource)
+	if sourceMetrics == nil {
+		return
+	}
 	if consumption.RRU > 0 {
 		sourceMetrics.rru.Add(consumption.RRU)
 	}
@@ -168,14 +208,8 @@ func (mc *groupMetricsCollection) addRequestSourceRU(resourceGroupName, requestS
 	}
 }
 
-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)
-	}
+func (mc *groupMetricsCollection) cleanupRequestSourceMetrics() {
+	mc.sourceState.cleanup()
 }
 
 type tokenCounter struct {
@@ -209,6 +243,7 @@ func newGroupCostController(
 	mainCfg *RUConfig,
 	lowRUNotifyChan chan notifyMsg,
 	tokenBucketUpdateChan chan *groupCostController,
+	sourceState *requestSourceMetricsState,
 ) (*groupCostController, error) {
 	switch group.Mode {
 	case rmpb.GroupMode_RUMode:
@@ -218,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,
@@ -630,7 +668,7 @@ func (gc *groupCostController) onRequestWaitImpl(
 		waitDuration += d
 	}
 
-	gc.metrics.addRequestSourceRU(gc.name, info.RequestSource(), delta)
+	gc.metrics.addRequestSourceRU(info.RequestSource(), delta)
 
 	gc.mu.Lock()
 	// Calculate the penalty of the store
@@ -677,7 +715,7 @@ func (gc *groupCostController) onResponseImpl(
 	add(gc.mu.globalCounter, count)
 	gc.mu.Unlock()
 
-	gc.metrics.addRequestSourceRU(gc.name, req.RequestSource(), delta)
+	gc.metrics.addRequestSourceRU(req.RequestSource(), delta)
 
 	return delta, nil
 }
@@ -720,7 +758,7 @@ func (gc *groupCostController) onResponseWaitImpl(
 	add(gc.mu.globalCounter, count)
 	gc.mu.Unlock()
 
-	gc.metrics.addRequestSourceRU(gc.name, req.RequestSource(), delta)
+	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,