Repository: tikv/pd (2025-07-02T06:20:27Z)

Issue:
TokenBucketsRequest doesn't have keyspace ID
## Bug Report

We try to extract keyspace ID here:
https://github.com/tikv/pd/blob/6002a2eaf02bb4f01039a316849bd73163fb0e2c/pkg/mcs/resourcemanager/server/grpc_service.go#L190

But the request doesn't have this field, which results in "" after `getKeyspaceNameByID`.



Discussion:

Patch:
diff --git a/client/resource_group/controller/controller.go b/client/resource_group/controller/controller.go
index be3094a1a2f..24134fa6fef 100644
--- a/client/resource_group/controller/controller.go
+++ b/client/resource_group/controller/controller.go
@@ -627,6 +627,9 @@ func (c *ResourceGroupsController) collectTokenBucketRequests(ctx context.Contex
 		gc := value.(*groupCostController)
 		request := gc.collectRequestAndConsumption(typ)
 		if request != nil {
+			request.KeyspaceId = &rmpb.KeyspaceIDValue{
+				Value: c.keyspaceID,
+			}
 			c.run.currentRequests = append(c.run.currentRequests, request)
 			gc.metrics.tokenRequestCounter.Inc()
 		}


Test patch:
diff --git a/client/resource_group/controller/controller_test.go b/client/resource_group/controller/controller_test.go
index cc8a6a620f9..7eb5da7e4c5 100644
--- a/client/resource_group/controller/controller_test.go
+++ b/client/resource_group/controller/controller_test.go
@@ -547,3 +547,55 @@ func TestGetResourceGroup(t *testing.T) {
 	re.Error(err)
 	re.Nil(gc02)
 }
+
+func TestTokenBucketsRequestWithKeyspaceID(t *testing.T) {
+	re := require.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	checkKeyspace := func(keyspaceID uint32) {
+		mockProvider := newMockResourceGroupProvider()
+		controller, err := NewResourceGroupController(ctx, 1, mockProvider, nil, keyspaceID)
+		re.NoError(err)
+		controller.Start(ctx)
+
+		testResourceGroup := &rmpb.ResourceGroup{
+			Name: "test-group",
+			Mode: rmpb.GroupMode_RUMode,
+			RUSettings: &rmpb.GroupRequestUnitSettings{
+				RU: &rmpb.TokenBucket{Settings: &rmpb.TokenLimitSettings{FillRate: 1000000}},
+			},
+		}
+		mockProvider.On("GetResourceGroup", mock.Anything, "test-group", mock.Anything).Return(testResourceGroup, nil)
+
+		gc, err := controller.tryGetResourceGroupController(ctx, "test-group", false)
+		re.NoError(err)
+		re.NotNil(gc)
+
+		requestReceived := make(chan bool, 1)
+
+		mockProvider.On("AcquireTokenBuckets", mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
+			request := args.Get(1).(*rmpb.TokenBucketsRequest)
+			re.Len(request.Requests, 1)
+			req := request.Requests[0]
+			re.NotNil(req.KeyspaceId)
+			re.Equal(keyspaceID, req.GetKeyspaceId().GetValue())
+			requestReceived <- true
+		}).Return([]*rmpb.TokenBucketResponse{}, nil)
+
+		// Trigger a low token report to ensure collectTokenBucketRequests is called
+		counter := gc.run.requestUnitTokens[rmpb.RequestUnitType_RU]
+		counter.limiter.mu.Lock()
+		counter.limiter.notify()
+		counter.limiter.mu.Unlock()
+
+		select {
+		case <-requestReceived:
+			// Success
+		case <-time.After(5 * time.Second):
+			t.Fatal("timed out waiting for AcquireTokenBuckets to be called")
+		}
+	}
+	checkKeyspace(constants.NullKeyspaceID)
+	checkKeyspace(1)
+}