client/resource_group: fix race between cleanup and getOrCreate for request-source metrics state

Use LoadAndDelete in cleanupRequestSourceMetricsState so the map entry
is removed atomically before the state is closed. Any hot-path goroutine
still holding the old reference no-ops via the closed check, and the
next getOrCreateRequestSourceMetricsState allocates a fresh state
instead of returning a closed one. Addresses rleungx's review comment.

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

Author: YuhaoZhang00

client/resource_group/controller/global_controller.go

@@ -490,9 +490,8 @@ func (c *ResourceGroupsController) getOrCreateRequestSourceMetricsState(name str
 }
 
 func (c *ResourceGroupsController) cleanupRequestSourceMetricsState(name string) {
-	if state, ok := c.requestSourceStates.Load(name); ok {
-		state.(*requestSourceMetricsState).cleanup()
-		c.requestSourceStates.Delete(name)
+	if v, loaded := c.requestSourceStates.LoadAndDelete(name); loaded {
+		v.(*requestSourceMetricsState).cleanup()
 	}
 }
 
@@ -621,9 +620,8 @@ func (c *ResourceGroupsController) cleanUpResourceGroup() {
 		gc.mu.Unlock()
 		if equalRU(latestConsumption, *gc.run.consumption) {
 			if gc.inactive || gc.tombstone.Load() {
-				gc.metrics.cleanupRequestSourceMetrics()
+				c.cleanupRequestSourceMetricsState(resourceGroupName)
 				c.groupsController.Delete(resourceGroupName)
-				c.requestSourceStates.Delete(resourceGroupName)
 				metrics.ResourceGroupStatusGauge.DeleteLabelValues(resourceGroupName, resourceGroupName)
 				return true
 			}

client/resource_group/controller/group_controller.go

@@ -208,10 +208,6 @@ func (mc *groupMetricsCollection) addRequestSourceRU(requestSource string, consu
 	}
 }
 
-func (mc *groupMetricsCollection) cleanupRequestSourceMetrics() {
-	mc.sourceState.cleanup()
-}
-
 type tokenCounter struct {
 	fillRate uint64
 

client/resource_group/controller/request_source_metrics_test.go

@@ -396,3 +396,161 @@ func TestRevivedResourceGroupCleanupRemovesExistingRequestSourceMetrics(t *testi
 	re.False(ok)
 	re.Equal(beforeCount, collectorMetricCount(controllerMetrics.RequestSourceRUCounter))
 }
+
+// TestGetOrCreateAfterCleanupReturnsFreshState verifies the fix for the race
+// rleungx identified: after cleanupRequestSourceMetricsState runs,
+// getOrCreateRequestSourceMetricsState must return a fresh, non-closed state
+// so that a newly created gc can record metrics normally.
+func TestGetOrCreateAfterCleanupReturnsFreshState(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-create-after-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)
+
+	// Phase 1: create a gc, record some metrics, then clean up.
+	gc, err := controller.tryGetResourceGroupController(ctx, group.Name, false)
+	re.NoError(err)
+
+	req := &TestRequestInfo{
+		isWrite:       true,
+		writeBytes:    64,
+		numReplicas:   1,
+		storeID:       1,
+		requestSource: "internal_create_after_cleanup",
+	}
+	resp := &TestResponseInfo{readBytes: 64, succeed: true}
+
+	_, _, _, _, 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
+	controller.cleanUpResourceGroup()
+
+	_, loaded := controller.loadGroupController(group.Name)
+	re.False(loaded)
+
+	// Phase 2: getOrCreateRequestSourceMetricsState must return a fresh,
+	// non-closed state after cleanup has removed the old one.
+	state := controller.getOrCreateRequestSourceMetricsState(group.Name)
+	re.NotNil(state)
+
+	state.mu.RLock()
+	re.False(state.closed)
+	state.mu.RUnlock()
+
+	// Phase 3: create a new gc with this state and verify it can record metrics.
+	newGC, err := newGroupCostController(
+		group,
+		controller.ruConfig,
+		controller.lowTokenNotifyChan,
+		controller.tokenBucketUpdateChan,
+		state,
+	)
+	re.NoError(err)
+
+	beforeCount := collectorMetricCount(controllerMetrics.RequestSourceRUCounter)
+	_, _, _, _, err = newGC.onRequestWaitImpl(context.Background(), req)
+	re.NoError(err)
+	_, err = newGC.onResponseImpl(req, resp)
+	re.NoError(err)
+
+	sourceMetrics, cacheSize := requestSourceStateSnapshot(t, newGC, req.requestSource)
+	re.Equal(1, cacheSize)
+	re.NotNil(sourceMetrics)
+	re.Greater(counterValue(t, sourceMetrics.wru), float64(0))
+	re.Greater(counterValue(t, sourceMetrics.rru), float64(0))
+	re.Equal(beforeCount+2, collectorMetricCount(controllerMetrics.RequestSourceRUCounter))
+
+	// Cleanup for this test.
+	state.cleanup()
+}
+
+// TestCleanupThenRecreateViaFullPath exercises the full end-to-end path:
+// cleanup a group, then tryGetResourceGroupController re-creates it, and
+// the new gc records metrics successfully.
+func TestCleanupThenRecreateViaFullPath(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-full-recreate",
+		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)
+
+	req := &TestRequestInfo{
+		isWrite:       true,
+		writeBytes:    64,
+		numReplicas:   1,
+		storeID:       1,
+		requestSource: "internal_full_recreate",
+	}
+	resp := &TestResponseInfo{readBytes: 64, succeed: true}
+
+	// Create, use, and clean up the group.
+	gc, err := controller.tryGetResourceGroupController(ctx, group.Name, false)
+	re.NoError(err)
+	_, _, _, _, 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
+	controller.cleanUpResourceGroup()
+
+	_, loaded := controller.loadGroupController(group.Name)
+	re.False(loaded)
+
+	// Re-create through the normal path (simulates a new request arriving
+	// after the group was cleaned up).
+	gc2, err := controller.tryGetResourceGroupController(ctx, group.Name, false)
+	re.NoError(err)
+
+	beforeCount := collectorMetricCount(controllerMetrics.RequestSourceRUCounter)
+	_, _, _, _, err = gc2.onRequestWaitImpl(context.Background(), req)
+	re.NoError(err)
+	_, err = gc2.onResponseImpl(req, resp)
+	re.NoError(err)
+
+	sourceMetrics, cacheSize := requestSourceStateSnapshot(t, gc2, req.requestSource)
+	re.Equal(1, cacheSize)
+	re.NotNil(sourceMetrics)
+	re.Greater(counterValue(t, sourceMetrics.wru), float64(0))
+	re.Greater(counterValue(t, sourceMetrics.rru), float64(0))
+	re.Equal(beforeCount+2, collectorMetricCount(controllerMetrics.RequestSourceRUCounter))
+
+	// Cleanup for this test.
+	gc2.metrics.sourceState.cleanup()
+}