From 62cc95123c0079ed219a3efccb8836c3ded3667b Mon Sep 17 00:00:00 2001 From: Connor Date: Fri, 28 Apr 2023 19:03:56 +0800 Subject: [PATCH] Fix wrong resource group name for some requests (#788) * fix wrong resource group name Signed-off-by: Connor1996 --- integration_tests/resource_group_test.go | 64 ++++++++++++++++++++++++ txnkv/txnsnapshot/scan.go | 3 +- txnkv/txnsnapshot/snapshot.go | 4 +- 3 files changed, 67 insertions(+), 4 deletions(-) create mode 100644 integration_tests/resource_group_test.go diff --git a/integration_tests/resource_group_test.go b/integration_tests/resource_group_test.go new file mode 100644 index 00000000..c423d884 --- /dev/null +++ b/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()) +} diff --git a/txnkv/txnsnapshot/scan.go b/txnkv/txnsnapshot/scan.go index a20d1763..221cf591 100644 --- a/txnkv/txnsnapshot/scan.go +++ b/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()), }) diff --git a/txnkv/txnsnapshot/snapshot.go b/txnkv/txnsnapshot/snapshot.go index 12a0e544..a1d14d46 100644 --- a/txnkv/txnsnapshot/snapshot.go +++ b/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()), })