TopSQL: add table id in resource tag (pingcap#55555)

ref pingcap#55540
Benjamin2037 · Sep 11, 2024 · 7eef7e4 · 7eef7e4
1 parent 617c1e6
commit 7eef7e4
Show file tree

Hide file tree

Showing 24 changed files with 193 additions and 60 deletions.
diff --git a/DEPS.bzl b/DEPS.bzl
@@ -5893,13 +5893,13 @@ def go_deps():
         name = "com_github_pingcap_tipb",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/pingcap/tipb",
-        sha256 = "6e910c9689f1a81bad2ae55be1746d456c317d696ff2687390d3fb30f7d05c6d",
-        strip_prefix = "github.com/pingcap/[email protected]20240703084358-e46e4632bd2b",
+        sha256 = "598bb728ddaf83863e7c79c8790cd2ae907c7ae762b13bfcf4164da5b9e2e4dc",
+        strip_prefix = "github.com/pingcap/[email protected]20240823074000-a40c2347786e",
         urls = [
-            "http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240703084358-e46e4632bd2b.zip",
-            "http://ats.apps.svc/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240703084358-e46e4632bd2b.zip",
-            "https://cache.hawkingrei.com/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240703084358-e46e4632bd2b.zip",
-            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240703084358-e46e4632bd2b.zip",
+            "http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240823074000-a40c2347786e.zip",
+            "http://ats.apps.svc/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240823074000-a40c2347786e.zip",
+            "https://cache.hawkingrei.com/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240823074000-a40c2347786e.zip",
+            "https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20240823074000-a40c2347786e.zip",
         ],
     )
     go_repository(

diff --git a/go.mod b/go.mod
@@ -89,7 +89,7 @@ require (
 	github.com/pingcap/log v1.1.1-0.20240314023424-862ccc32f18d
 	github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5
 	github.com/pingcap/tidb/pkg/parser v0.0.0-20211011031125-9b13dc409c5e
-	github.com/pingcap/tipb v0.0.0-20240703084358-e46e4632bd2b
+	github.com/pingcap/tipb v0.0.0-20240823074000-a40c2347786e
 	github.com/prometheus/client_golang v1.20.2
 	github.com/prometheus/client_model v0.6.1
 	github.com/prometheus/common v0.57.0

diff --git a/go.sum b/go.sum
@@ -676,8 +676,8 @@ github.com/pingcap/log v1.1.1-0.20240314023424-862ccc32f18d h1:y3EueKVfVykdpTyfU
 github.com/pingcap/log v1.1.1-0.20240314023424-862ccc32f18d/go.mod h1:ORfBOFp1eteu2odzsyaxI+b8TzJwgjwyQcGhI+9SfEA=
 github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5 h1:T4pXRhBflzDeAhmOQHNPRRogMYxP13V7BkYw3ZsoSfE=
 github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5/go.mod h1:rlimy0GcTvjiJqvD5mXTRr8O2eNZPBrcUgiWVYp9530=
-github.com/pingcap/tipb v0.0.0-20240703084358-e46e4632bd2b h1:tySAGYw21A3Xa8CcA9jBTfrgAB3+KQWyqyW7fUyokzk=
-github.com/pingcap/tipb v0.0.0-20240703084358-e46e4632bd2b/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
+github.com/pingcap/tipb v0.0.0-20240823074000-a40c2347786e h1:chR6iRwU9MUivOYNH+26UqrO8Y7t3ZltX+Jukv+f+iM=
+github.com/pingcap/tipb v0.0.0-20240823074000-a40c2347786e/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs=
 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
 github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=

diff --git a/pkg/ddl/BUILD.bazel b/pkg/ddl/BUILD.bazel
@@ -160,7 +160,6 @@ go_library(
         "//pkg/util/mock",
         "//pkg/util/ppcpuusage",
         "//pkg/util/ranger",
-        "//pkg/util/resourcegrouptag",
         "//pkg/util/rowDecoder",
         "//pkg/util/rowcodec",
         "//pkg/util/set",

diff --git a/pkg/ddl/ddl.go b/pkg/ddl/ddl.go
@@ -62,7 +62,6 @@ import (
 	"github.com/pingcap/tidb/pkg/util/gcutil"
 	"github.com/pingcap/tidb/pkg/util/generic"
 	"github.com/pingcap/tidb/pkg/util/intest"
-	"github.com/tikv/client-go/v2/tikvrpc"
 	clientv3 "go.etcd.io/etcd/client/v3"
 	atomicutil "go.uber.org/atomic"
 	"go.uber.org/zap"
@@ -421,7 +420,7 @@ func (dc *ddlCtx) setDDLSourceForDiagnosis(jobID int64, jobType model.ActionType
 	ctx.setDDLLabelForDiagnosis(jobType)
 }
 
-func (dc *ddlCtx) getResourceGroupTaggerForTopSQL(jobID int64) tikvrpc.ResourceGroupTagger {
+func (dc *ddlCtx) getResourceGroupTaggerForTopSQL(jobID int64) *kv.ResourceGroupTagBuilder {
 	dc.jobCtx.Lock()
 	defer dc.jobCtx.Unlock()
 	ctx, exists := dc.jobCtx.jobCtxMap[jobID]

diff --git a/pkg/ddl/job_worker.go b/pkg/ddl/job_worker.go
@@ -44,10 +44,8 @@ import (
 	tidbutil "github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/dbterror"
 	tidblogutil "github.com/pingcap/tidb/pkg/util/logutil"
-	"github.com/pingcap/tidb/pkg/util/resourcegrouptag"
 	"github.com/pingcap/tidb/pkg/util/topsql"
 	topsqlstate "github.com/pingcap/tidb/pkg/util/topsql/state"
-	"github.com/tikv/client-go/v2/tikvrpc"
 	kvutil "github.com/tikv/client-go/v2/util"
 	atomicutil "go.uber.org/atomic"
 	"go.uber.org/zap"
@@ -637,17 +635,13 @@ func (w *worker) checkBeforeCommit() error {
 	return nil
 }
 
-func (w *ReorgContext) getResourceGroupTaggerForTopSQL() tikvrpc.ResourceGroupTagger {
+func (w *ReorgContext) getResourceGroupTaggerForTopSQL() *kv.ResourceGroupTagBuilder {
 	if !topsqlstate.TopSQLEnabled() || w.cacheDigest == nil {
 		return nil
 	}
 
 	digest := w.cacheDigest
-	tagger := func(req *tikvrpc.Request) {
-		req.ResourceGroupTag = resourcegrouptag.EncodeResourceGroupTag(digest, nil,
-			resourcegrouptag.GetResourceGroupLabelByKey(resourcegrouptag.GetFirstKeyFromRequest(req)))
-	}
-	return tagger
+	return kv.NewResourceGroupTagBuilder().SetSQLDigest(digest)
 }
 
 func (w *ReorgContext) ddlJobSourceType() string {

diff --git a/pkg/distsql/BUILD.bazel b/pkg/distsql/BUILD.bazel
@@ -44,7 +44,6 @@ go_library(
         "@com_github_pingcap_tipb//go-tipb",
         "@com_github_tikv_client_go_v2//metrics",
         "@com_github_tikv_client_go_v2//tikv",
-        "@com_github_tikv_client_go_v2//tikvrpc",
         "@com_github_tikv_client_go_v2//tikvrpc/interceptor",
         "@com_github_tikv_client_go_v2//util",
         "@org_golang_google_grpc//metadata",
@@ -66,7 +65,7 @@ go_test(
     embed = [":distsql"],
     flaky = True,
     race = "on",
-    shard_count = 27,
+    shard_count = 28,
     deps = [
         "//pkg/distsql/context",
         "//pkg/errctx",
@@ -93,6 +92,7 @@ go_test(
         "//pkg/util/mock",
         "//pkg/util/paging",
         "//pkg/util/ranger",
+        "@com_github_pingcap_kvproto//pkg/coprocessor",
         "@com_github_pingcap_tipb//go-tipb",
         "@com_github_stretchr_testify//require",
         "@com_github_tikv_client_go_v2//kv",

diff --git a/pkg/distsql/context/BUILD.bazel b/pkg/distsql/context/BUILD.bazel
@@ -18,7 +18,6 @@ go_library(
         "//pkg/util/tiflash",
         "//pkg/util/topsql/stmtstats",
         "@com_github_tikv_client_go_v2//kv",
-        "@com_github_tikv_client_go_v2//tikvrpc",
     ],
 )
 

diff --git a/pkg/distsql/context/context.go b/pkg/distsql/context/context.go
@@ -29,7 +29,6 @@ import (
 	"github.com/pingcap/tidb/pkg/util/tiflash"
 	"github.com/pingcap/tidb/pkg/util/topsql/stmtstats"
 	tikvstore "github.com/tikv/client-go/v2/kv"
-	"github.com/tikv/client-go/v2/tikvrpc"
 )
 
 // DistSQLContext provides all information needed by using functions in `distsql`
@@ -68,7 +67,7 @@ type DistSQLContext struct {
 	NotFillCache                  bool
 	TaskID                        uint64
 	Priority                      mysql.PriorityEnum
-	ResourceGroupTagger           tikvrpc.ResourceGroupTagger
+	ResourceGroupTagger           *kv.ResourceGroupTagBuilder
 	EnablePaging                  bool
 	MinPagingSize                 int
 	MaxPagingSize                 int

diff --git a/pkg/distsql/request_builder.go b/pkg/distsql/request_builder.go
@@ -38,7 +38,6 @@ import (
 	"github.com/pingcap/tidb/pkg/util/memory"
 	"github.com/pingcap/tidb/pkg/util/ranger"
 	"github.com/pingcap/tipb/go-tipb"
-	"github.com/tikv/client-go/v2/tikvrpc"
 )
 
 // RequestBuilder is used to build a "kv.Request".
@@ -374,7 +373,7 @@ func (builder *RequestBuilder) SetFromInfoSchema(is infoschema.MetaOnlyInfoSchem
 }
 
 // SetResourceGroupTagger sets the request resource group tagger.
-func (builder *RequestBuilder) SetResourceGroupTagger(tagger tikvrpc.ResourceGroupTagger) *RequestBuilder {
+func (builder *RequestBuilder) SetResourceGroupTagger(tagger *kv.ResourceGroupTagBuilder) *RequestBuilder {
 	builder.Request.ResourceGroupTagger = tagger
 	return builder
 }

diff --git a/pkg/distsql/request_builder_test.go b/pkg/distsql/request_builder_test.go
@@ -18,7 +18,9 @@ import (
 	"math"
 	"testing"
 	"time"
+	"unsafe"
 
+	"github.com/pingcap/kvproto/pkg/coprocessor"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/meta/model"
 	"github.com/pingcap/tidb/pkg/resourcegroup"
@@ -32,6 +34,8 @@ import (
 	"github.com/pingcap/tidb/pkg/util/ranger"
 	"github.com/pingcap/tipb/go-tipb"
 	"github.com/stretchr/testify/require"
+	clikv "github.com/tikv/client-go/v2/kv"
+	"github.com/tikv/client-go/v2/tikvrpc"
 )
 
 type handleRange struct {
@@ -852,3 +856,42 @@ func TestBuildTableRangeCommonHandle(t *testing.T) {
 		{StartKey: tablecodec.EncodeRowKey(7, low), EndKey: tablecodec.EncodeRowKey(7, high)},
 	}, ranges)
 }
+
+func TestRequestBuilderHandle(t *testing.T) {
+	handles := []kv.Handle{kv.IntHandle(0), kv.IntHandle(2), kv.IntHandle(3), kv.IntHandle(4),
+		kv.IntHandle(5), kv.IntHandle(10), kv.IntHandle(11), kv.IntHandle(100)}
+
+	resourceTagBuilder := kv.NewResourceGroupTagBuilder()
+	tableID := int64(15)
+	actual, err := (&RequestBuilder{}).SetTableHandles(tableID, handles).
+		SetDAGRequest(&tipb.DAGRequest{}).
+		SetDesc(false).
+		SetKeepOrder(false).
+		SetFromSessionVars(DefaultDistSQLContext).
+		SetResourceGroupTagger(resourceTagBuilder).
+		Build()
+	require.NoError(t, err)
+	ranges := make([]*coprocessor.KeyRange, 0, actual.KeyRanges.TotalRangeNum())
+	actual.KeyRanges.ForEachPartition(
+		func(innerRanges []kv.KeyRange) {
+			for _, ran := range innerRanges {
+				ranges = append(ranges, (*coprocessor.KeyRange)(unsafe.Pointer(&ran)))
+			}
+		})
+
+	copReq := coprocessor.Request{
+		Tp:      actual.Tp,
+		StartTs: actual.StartTs,
+		Data:    actual.Data,
+		Ranges:  ranges,
+	}
+	var seed uint32
+	req := tikvrpc.NewReplicaReadRequest(tikvrpc.CmdCop, &copReq, clikv.ReplicaReadLeader, &seed)
+	actual.ResourceGroupTagger.Build(req)
+
+	// the request should have the resource group tag, and the tag should contain the table id
+	tag := &tipb.ResourceGroupTag{}
+	err = tag.Unmarshal(req.ResourceGroupTag)
+	require.NoError(t, err)
+	require.Equal(t, tag.TableId, tableID)
+}
diff --git a/pkg/executor/executor.go b/pkg/executor/executor.go
@@ -83,7 +83,6 @@ import (
 	"github.com/pingcap/tidb/pkg/util/logutil"
 	"github.com/pingcap/tidb/pkg/util/logutil/consistency"
 	"github.com/pingcap/tidb/pkg/util/memory"
-	"github.com/pingcap/tidb/pkg/util/resourcegrouptag"
 	"github.com/pingcap/tidb/pkg/util/sqlexec"
 	"github.com/pingcap/tidb/pkg/util/syncutil"
 	"github.com/pingcap/tidb/pkg/util/topsql"
@@ -1223,13 +1222,16 @@ func newLockCtx(sctx sessionctx.Context, lockWaitTime int64, numKeys int) (*tikv
 			return nil
 		}
 		if mutation := req.Mutations[0]; mutation != nil {
-			label := resourcegrouptag.GetResourceGroupLabelByKey(mutation.Key)
 			normalized, digest := seVars.StmtCtx.SQLDigest()
 			if len(normalized) == 0 {
 				return nil
 			}
 			_, planDigest := seVars.StmtCtx.GetPlanDigest()
-			return resourcegrouptag.EncodeResourceGroupTag(digest, planDigest, label)
+
+			return kv.NewResourceGroupTagBuilder().
+				SetPlanDigest(planDigest).
+				SetSQLDigest(digest).
+				EncodeTagWithKey(mutation.Key)
 		}
 		return nil
 	}

diff --git a/pkg/executor/resource_tag_test.go b/pkg/executor/resource_tag_test.go
@@ -57,7 +57,7 @@ func TestResourceGroupTag(t *testing.T) {
 
 	var sqlDigest, planDigest *parser.Digest
 	var tagLabel tipb.ResourceGroupTagLabel
-	checkFn := func() {}
+	checkFn := func(int64) {}
 	unistoreRPCClientSendHook := func(req *tikvrpc.Request) {
 		var startKey []byte
 		var ctx *kvrpcpb.Context
@@ -99,8 +99,10 @@ func TestResourceGroupTag(t *testing.T) {
 		require.NoError(t, err)
 		sqlDigest = parser.NewDigest(tag.SqlDigest)
 		planDigest = parser.NewDigest(tag.PlanDigest)
-		tagLabel = *tag.Label
-		checkFn()
+		if tag.Label != nil {
+			tagLabel = *tag.Label
+		}
+		checkFn(tag.TableId)
 	}
 	unistore.UnistoreRPCClientSendHook.Store(&unistoreRPCClientSendHook)
 
@@ -187,10 +189,11 @@ func TestResourceGroupTag(t *testing.T) {
 		_, expectSQLDigest := parser.NormalizeDigest(ca.sql)
 		var expectPlanDigest *parser.Digest
 		checkCnt := 0
-		checkFn = func() {
+		checkFn = func(tid int64) {
 			if ca.ignore {
 				return
 			}
+			require.Equal(t, tbInfo.Meta().ID, tid)
 			if expectPlanDigest == nil {
 				info := tk.Session().ShowProcess()
 				require.NotNil(t, info)
@@ -208,7 +211,6 @@ func TestResourceGroupTag(t *testing.T) {
 			require.True(t, ok)
 			checkCnt++
 		}
-
 		if strings.HasPrefix(ca.sql, "select") {
 			tk.MustQuery(ca.sql)
 		} else {

diff --git a/pkg/kv/BUILD.bazel b/pkg/kv/BUILD.bazel
@@ -25,6 +25,7 @@ go_library(
         "//pkg/config",
         "//pkg/errno",
         "//pkg/meta/model",
+        "//pkg/parser",
         "//pkg/parser/mysql",
         "//pkg/parser/terror",
         "//pkg/resourcegroup",
@@ -34,6 +35,7 @@ go_library(
         "//pkg/util/intest",
         "//pkg/util/logutil",
         "//pkg/util/memory",
+        "//pkg/util/resourcegrouptag",
         "//pkg/util/set",
         "//pkg/util/size",
         "//pkg/util/tiflash",