Fix wrong resource group name for some requests (#788)

* fix wrong resource group name

Signed-off-by: Connor1996 <zbk602423539@gmail.com>

integration_tests/resource_group_test.go

@@ -0,0 +1,64 @@
+package tikv_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tikv/client-go/v2/tikv"
+	"github.com/tikv/client-go/v2/tikvrpc"
+)
+
+var _ tikv.Client = &resourceGroupNameMockClient{}
+
+type resourceGroupNameMockClient struct {
+	tikv.Client
+
+	t            *testing.T
+	expectedTag  string
+	requestCount int
+}
+
+func (c *resourceGroupNameMockClient) SendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error) {
+	if req.GetResourceControlContext().GetResourceGroupName() == c.expectedTag {
+		c.requestCount++
+	}
+	return c.Client.SendRequest(ctx, addr, req, timeout)
+}
+
+func TestResourceGroupName(t *testing.T) {
+	testTag := "test"
+	/* Get */
+	store := NewTestStore(t)
+	client := &resourceGroupNameMockClient{t: t, Client: store.GetTiKVClient(), expectedTag: testTag}
+	store.SetTiKVClient(client)
+	txn, err := store.Begin()
+	assert.NoError(t, err)
+	txn.SetResourceGroupName(testTag)
+	_, _ = txn.Get(context.Background(), []byte{})
+	assert.Equal(t, 1, client.requestCount)
+	assert.NoError(t, store.Close())
+
+	/* BatchGet */
+	store = NewTestStore(t)
+	client = &resourceGroupNameMockClient{t: t, Client: store.GetTiKVClient(), expectedTag: testTag}
+	store.SetTiKVClient(client)
+	txn, err = store.Begin()
+	assert.NoError(t, err)
+	txn.SetResourceGroupName(testTag)
+	_, _ = txn.BatchGet(context.Background(), [][]byte{[]byte("k")})
+	assert.Equal(t, 1, client.requestCount)
+	assert.NoError(t, store.Close())
+
+	/* Scan */
+	store = NewTestStore(t)
+	client = &resourceGroupNameMockClient{t: t, Client: store.GetTiKVClient(), expectedTag: testTag}
+	store.SetTiKVClient(client)
+	txn, err = store.Begin()
+	assert.NoError(t, err)
+	txn.SetResourceGroupName(testTag)
+	_, _ = txn.Iter([]byte("abc"), []byte("def"))
+	assert.Equal(t, 1, client.requestCount)
+	assert.NoError(t, store.Close())
+}

txnkv/txnsnapshot/scan.go

@@ -49,7 +49,6 @@ import (
 	"github.com/tikv/client-go/v2/tikvrpc"
 	"github.com/tikv/client-go/v2/tikvrpc/interceptor"
 	"github.com/tikv/client-go/v2/txnkv/txnlock"
-	"github.com/tikv/client-go/v2/util"
 	"go.uber.org/zap"
 )
 
@@ -248,7 +247,7 @@ func (s *Scanner) getData(bo *retry.Backoffer) error {
 			IsolationLevel:   s.snapshot.isolationLevel.ToPB(),
 			RequestSource:    s.snapshot.GetRequestSource(),
 			ResourceControlContext: &kvrpcpb.ResourceControlContext{
-				ResourceGroupName: util.ResourceGroupNameFromCtx(bo.GetCtx()),
+				ResourceGroupName: s.snapshot.mu.resourceGroupName,
 			},
 			BusyThresholdMs: uint32(s.snapshot.mu.busyThreshold.Milliseconds()),
 		})

txnkv/txnsnapshot/snapshot.go

@@ -398,7 +398,7 @@ func (s *KVSnapshot) batchGetSingleRegion(bo *retry.Backoffer, batch batchKeys,
 			IsolationLevel:   s.isolationLevel.ToPB(),
 			RequestSource:    s.GetRequestSource(),
 			ResourceControlContext: &kvrpcpb.ResourceControlContext{
-				ResourceGroupName: util.ResourceGroupNameFromCtx(bo.GetCtx()),
+				ResourceGroupName: s.mu.resourceGroupName,
 			},
 			BusyThresholdMs: uint32(s.mu.busyThreshold.Milliseconds()),
 		})
@@ -614,7 +614,7 @@ func (s *KVSnapshot) get(ctx context.Context, bo *retry.Backoffer, k []byte) ([]
 			IsolationLevel:   s.isolationLevel.ToPB(),
 			RequestSource:    s.GetRequestSource(),
 			ResourceControlContext: &kvrpcpb.ResourceControlContext{
-				ResourceGroupName: util.ResourceGroupNameFromCtx(bo.GetCtx()),
+				ResourceGroupName: s.mu.resourceGroupName,
 			},
 			BusyThresholdMs: uint32(s.mu.busyThreshold.Milliseconds()),
 		})