add resource group name in request context

go.mod

@@ -13,7 +13,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0
 	github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00
 	github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989
-	github.com/pingcap/kvproto v0.0.0-20221129023506-621ec37aac7a
+	github.com/pingcap/kvproto v0.0.0-20221227030452-22819f5b377a
 	github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.11.0

go.sum

@@ -155,8 +155,8 @@ github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00/go.mod h1:4qGtCB
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
 github.com/pingcap/kvproto v0.0.0-20221026112947-f8d61344b172/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
-github.com/pingcap/kvproto v0.0.0-20221129023506-621ec37aac7a h1:LzIZsQpXQlj8yF7+yvyOg680OaPq7bmPuDuszgXfHsw=
-github.com/pingcap/kvproto v0.0.0-20221129023506-621ec37aac7a/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
+github.com/pingcap/kvproto v0.0.0-20221227030452-22819f5b377a h1:4TiR0nGKmLmu2kTC22jOhUIplmv4GGUrUD9D2DTMms0=
+github.com/pingcap/kvproto v0.0.0-20221227030452-22819f5b377a/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
 github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81 h1:URLoJ61DmmY++Sa/yyPEQHG2s/ZBeV1FbIswHEMrdoY=
 github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
@@ -201,10 +201,6 @@ github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81P
 github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA=
 github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/tiancaiamao/gp v0.0.0-20221221095600-1a473d1f9b4b h1:4RNtqw1/tW67qP9fFgfQpTVd7DrfkaAWu4vsC18QmBo=
-github.com/tiancaiamao/gp v0.0.0-20221221095600-1a473d1f9b4b/go.mod h1:h4xBhSNtOeEosLJ4P7JyKXX7Cabg7AVkWCK5gV2vOrM=
-github.com/tiancaiamao/gp v0.0.0-20221230021323-917088386b9f h1:huMeRO2QPxK10eb19qevMkvLP0L1kGB9nybuZz2IvKk=
-github.com/tiancaiamao/gp v0.0.0-20221230021323-917088386b9f/go.mod h1:h4xBhSNtOeEosLJ4P7JyKXX7Cabg7AVkWCK5gV2vOrM=
 github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a h1:J/YdBZ46WKpXsxsW93SG+q0F8KI+yFrcIDT4c/RNoc4=
 github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a/go.mod h1:h4xBhSNtOeEosLJ4P7JyKXX7Cabg7AVkWCK5gV2vOrM=
 github.com/tikv/pd/client v0.0.0-20221031025758-80f0d8ca4d07 h1:ckPpxKcl75mO2N6a4cJXiZH43hvcHPpqc9dh1TmH1nc=

integration_tests/go.mod

@@ -6,7 +6,7 @@ require (
 	github.com/ninedraft/israce v0.0.3
 	github.com/pingcap/errors v0.11.5-0.20220729040631-518f63d66278
 	github.com/pingcap/failpoint v0.0.0-20220423142525-ae43b7f4e5c3
-	github.com/pingcap/kvproto v0.0.0-20221129023506-621ec37aac7a
+	github.com/pingcap/kvproto v0.0.0-20221227030452-22819f5b377a
 	github.com/pingcap/tidb v1.1.0-beta.0.20221101102559-97add26c8f84
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.8.0

integration_tests/go.sum

@@ -408,8 +408,8 @@ github.com/pingcap/fn v0.0.0-20200306044125-d5540d389059 h1:Pe2LbxRmbTfAoKJ65bZL
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
 github.com/pingcap/kvproto v0.0.0-20221026112947-f8d61344b172/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
-github.com/pingcap/kvproto v0.0.0-20221129023506-621ec37aac7a h1:LzIZsQpXQlj8yF7+yvyOg680OaPq7bmPuDuszgXfHsw=
-github.com/pingcap/kvproto v0.0.0-20221129023506-621ec37aac7a/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
+github.com/pingcap/kvproto v0.0.0-20221227030452-22819f5b377a h1:4TiR0nGKmLmu2kTC22jOhUIplmv4GGUrUD9D2DTMms0=
+github.com/pingcap/kvproto v0.0.0-20221227030452-22819f5b377a/go.mod h1:OYtxs0786qojVTmkVeufx93xe+jUgm56GUYRIKnmaGI=
 github.com/pingcap/log v0.0.0-20200511115504-543df19646ad/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81 h1:URLoJ61DmmY++Sa/yyPEQHG2s/ZBeV1FbIswHEMrdoY=
 github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=

txnkv/transaction/2pc.go

@@ -133,6 +133,7 @@ type twoPhaseCommitter struct {
 	detail              unsafe.Pointer
 	txnSize             int
 	hasNoNeedCommitKeys bool
+	resourceGroupName   string
 
 	primaryKey  []byte
 	forUpdateTS uint64
@@ -705,6 +706,7 @@ func (c *twoPhaseCommitter) initKeysAndMutations(ctx context.Context) error {
 	c.syncLog = txn.syncLog
 	c.resourceGroupTag = txn.resourceGroupTag
 	c.resourceGroupTagger = txn.resourceGroupTagger
+	c.resourceGroupName = txn.resourceGroupName
 	c.setDetail(commitDetail)
 
 	return nil

txnkv/transaction/cleanup.go

@@ -68,6 +68,7 @@ func (actionCleanup) handleSingleBatch(c *twoPhaseCommitter, bo *retry.Backoffer
 		ResourceGroupTag:       c.resourceGroupTag,
 		MaxExecutionDurationMs: uint64(client.MaxWriteExecutionTime.Milliseconds()),
 		RequestSource:          c.txn.GetRequestSource(),
+		ResourceGroupName:      c.resourceGroupName,
 	})
 	if c.resourceGroupTag == nil && c.resourceGroupTagger != nil {
 		c.resourceGroupTagger(req)

txnkv/transaction/commit.go

@@ -78,6 +78,7 @@ func (actionCommit) handleSingleBatch(c *twoPhaseCommitter, bo *retry.Backoffer,
 		TxnSource:              c.txnSource,
 		MaxExecutionDurationMs: uint64(client.MaxWriteExecutionTime.Milliseconds()),
 		RequestSource:          c.txn.GetRequestSource(),
+		ResourceGroupName:      c.resourceGroupName,
 	})
 	if c.resourceGroupTag == nil && c.resourceGroupTagger != nil {
 		c.resourceGroupTagger(req)

txnkv/transaction/pessimistic.go

@@ -129,6 +129,7 @@ func (action actionPessimisticLock) handleSingleBatch(c *twoPhaseCommitter, bo *
 		ResourceGroupTag:       action.LockCtx.ResourceGroupTag,
 		MaxExecutionDurationMs: uint64(client.MaxWriteExecutionTime.Milliseconds()),
 		RequestSource:          c.txn.GetRequestSource(),
+		ResourceGroupName:      c.resourceGroupName,
 	})
 	if action.LockCtx.ResourceGroupTag == nil && action.LockCtx.ResourceGroupTagger != nil {
 		req.ResourceGroupTag = action.LockCtx.ResourceGroupTagger(req.Req.(*kvrpcpb.PessimisticLockRequest))

txnkv/transaction/prewrite.go

@@ -185,6 +185,7 @@ func (c *twoPhaseCommitter) buildPrewriteRequest(batch batchMutations, txnSize u
 		TxnSource:              c.txnSource,
 		MaxExecutionDurationMs: uint64(client.MaxWriteExecutionTime.Milliseconds()),
 		RequestSource:          c.txn.GetRequestSource(),
+		ResourceGroupName:      c.resourceGroupName,
 	})
 	if c.resourceGroupTag == nil && c.resourceGroupTagger != nil {
 		c.resourceGroupTagger(r)

txnkv/transaction/txn.go

@@ -160,6 +160,8 @@ type KVTxn struct {
 	interceptor    interceptor.RPCInterceptor
 	assertionLevel kvrpcpb.AssertionLevel
 	*util.RequestSource
+	// resourceGroupName is the name of tenent resource group.
+	resourceGroupName string
 
 	aggressiveLockingContext *aggressiveLockingContext
 	aggressiveLockingDirty   bool
@@ -298,6 +300,12 @@ func (txn *KVTxn) SetResourceGroupTagger(tagger tikvrpc.ResourceGroupTagger) {
 	txn.GetSnapshot().SetResourceGroupTagger(tagger)
 }
 
+// SetResourceGroupName set resource group name for both read and write.
+func (txn *KVTxn) SetResourceGroupName(name string) {
+	txn.resourceGroupName = name
+	txn.GetSnapshot().SetResourceGroupName(name)
+}
+
 // SetRPCInterceptor sets interceptor.RPCInterceptor for the transaction and its related snapshot.
 // interceptor.RPCInterceptor will be executed before each RPC request is initiated.
 // Note that SetRPCInterceptor will replace the previously set interceptor.

txnkv/txnsnapshot/scan.go

@@ -226,11 +226,12 @@ func (s *Scanner) getData(bo *retry.Backoffer) error {
 		}
 		sreq := &kvrpcpb.ScanRequest{
 			Context: &kvrpcpb.Context{
-				Priority:         s.snapshot.priority.ToPB(),
-				NotFillCache:     s.snapshot.notFillCache,
-				IsolationLevel:   s.snapshot.isolationLevel.ToPB(),
-				ResourceGroupTag: s.snapshot.mu.resourceGroupTag,
-				RequestSource:    s.snapshot.GetRequestSource(),
+				Priority:          s.snapshot.priority.ToPB(),
+				NotFillCache:      s.snapshot.notFillCache,
+				IsolationLevel:    s.snapshot.isolationLevel.ToPB(),
+				ResourceGroupTag:  s.snapshot.mu.resourceGroupTag,
+				RequestSource:     s.snapshot.GetRequestSource(),
+				ResourceGroupName: s.snapshot.mu.resourceGroupName,
 			},
 			StartKey:   s.nextStartKey,
 			EndKey:     reqEndKey,
@@ -246,12 +247,13 @@ func (s *Scanner) getData(bo *retry.Backoffer) error {
 		}
 		s.snapshot.mu.RLock()
 		req := tikvrpc.NewReplicaReadRequest(tikvrpc.CmdScan, sreq, s.snapshot.mu.replicaRead, &s.snapshot.replicaReadSeed, kvrpcpb.Context{
-			Priority:         s.snapshot.priority.ToPB(),
-			NotFillCache:     s.snapshot.notFillCache,
-			TaskId:           s.snapshot.mu.taskID,
-			ResourceGroupTag: s.snapshot.mu.resourceGroupTag,
-			IsolationLevel:   s.snapshot.isolationLevel.ToPB(),
-			RequestSource:    s.snapshot.GetRequestSource(),
+			Priority:          s.snapshot.priority.ToPB(),
+			NotFillCache:      s.snapshot.notFillCache,
+			TaskId:            s.snapshot.mu.taskID,
+			ResourceGroupTag:  s.snapshot.mu.resourceGroupTag,
+			IsolationLevel:    s.snapshot.isolationLevel.ToPB(),
+			RequestSource:     s.snapshot.GetRequestSource(),
+			ResourceGroupName: s.snapshot.mu.resourceGroupName,
 		})
 		if s.snapshot.mu.resourceGroupTag == nil && s.snapshot.mu.resourceGroupTagger != nil {
 			s.snapshot.mu.resourceGroupTagger(req)

txnkv/txnsnapshot/snapshot.go

@@ -146,6 +146,8 @@ type KVSnapshot struct {
 		resourceGroupTagger tikvrpc.ResourceGroupTagger
 		// interceptor is used to decorate the RPC request logic related to the snapshot.
 		interceptor interceptor.RPCInterceptor
+		// resourceGroupName is used to bind the request to specified resource group.
+		resourceGroupName string
 	}
 	sampleStep uint32
 	*util.RequestSource
@@ -375,12 +377,13 @@ func (s *KVSnapshot) batchGetSingleRegion(bo *retry.Backoffer, batch batchKeys,
 			Keys:    pending,
 			Version: s.version,
 		}, s.mu.replicaRead, &s.replicaReadSeed, kvrpcpb.Context{
-			Priority:         s.priority.ToPB(),
-			NotFillCache:     s.notFillCache,
-			TaskId:           s.mu.taskID,
-			ResourceGroupTag: s.mu.resourceGroupTag,
-			IsolationLevel:   s.isolationLevel.ToPB(),
-			RequestSource:    s.GetRequestSource(),
+			Priority:          s.priority.ToPB(),
+			NotFillCache:      s.notFillCache,
+			TaskId:            s.mu.taskID,
+			ResourceGroupTag:  s.mu.resourceGroupTag,
+			IsolationLevel:    s.isolationLevel.ToPB(),
+			RequestSource:     s.GetRequestSource(),
+			ResourceGroupName: s.mu.resourceGroupName,
 		})
 		if s.mu.resourceGroupTag == nil && s.mu.resourceGroupTagger != nil {
 			s.mu.resourceGroupTagger(req)
@@ -578,12 +581,13 @@ func (s *KVSnapshot) get(ctx context.Context, bo *retry.Backoffer, k []byte) ([]
 			Key:     k,
 			Version: s.version,
 		}, s.mu.replicaRead, &s.replicaReadSeed, kvrpcpb.Context{
-			Priority:         s.priority.ToPB(),
-			NotFillCache:     s.notFillCache,
-			TaskId:           s.mu.taskID,
-			ResourceGroupTag: s.mu.resourceGroupTag,
-			IsolationLevel:   s.isolationLevel.ToPB(),
-			RequestSource:    s.GetRequestSource(),
+			Priority:          s.priority.ToPB(),
+			NotFillCache:      s.notFillCache,
+			TaskId:            s.mu.taskID,
+			ResourceGroupTag:  s.mu.resourceGroupTag,
+			IsolationLevel:    s.isolationLevel.ToPB(),
+			RequestSource:     s.GetRequestSource(),
+			ResourceGroupName: s.mu.resourceGroupName,
 		})
 	if s.mu.resourceGroupTag == nil && s.mu.resourceGroupTagger != nil {
 		s.mu.resourceGroupTagger(req)
@@ -852,6 +856,13 @@ func (s *KVSnapshot) AddRPCInterceptor(it interceptor.RPCInterceptor) {
 	s.mu.interceptor = interceptor.ChainRPCInterceptors(s.mu.interceptor, it)
 }
 
+// SetResourceGroupName set resource group name of the kv request.
+func (s *KVSnapshot) SetResourceGroupName(name string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.mu.resourceGroupName = name
+}
+
 // SnapCacheHitCount gets the snapshot cache hit count. Only for test.
 func (s *KVSnapshot) SnapCacheHitCount() int {
 	return int(atomic.LoadInt64(&s.mu.hitCnt))