topsql: reduce data race of sql digest (#30296)

pingcap · Dec 1, 2021 · 4e61d16 · 4e61d16
1 parent 80d18b6
commit 4e61d16
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 38 deletions.
diff --git a/executor/executor.go b/executor/executor.go
@@ -983,7 +983,12 @@ func newLockCtx(seVars *variable.SessionVars, lockWaitTime int64) *tikvstore.Loc
  }
  if mutation := req.Mutations[0]; mutation != nil {
  label := resourcegrouptag.GetResourceGroupLabelByKey(mutation.Key)
- return seVars.StmtCtx.GetResourceGroupTagByLabel(label)
+ normalized, digest := seVars.StmtCtx.SQLDigest()
+ if len(normalized) == 0 {
+ return nil
+ }
+ _, planDigest := seVars.StmtCtx.GetPlanDigest()
+ return resourcegrouptag.EncodeResourceGroupTag(digest, planDigest, label)
  }
  return nil
  }

diff --git a/sessionctx/stmtctx/stmtctx.go b/sessionctx/stmtctx/stmtctx.go
@@ -30,7 +30,6 @@ import (
  "github.com/pingcap/tidb/util/memory"
  "github.com/pingcap/tidb/util/resourcegrouptag"
  "github.com/pingcap/tidb/util/tracing"
- "github.com/pingcap/tipb/go-tipb"
  "github.com/tikv/client-go/v2/tikvrpc"
  "github.com/tikv/client-go/v2/util"
  atomic2 "go.uber.org/atomic"
@@ -289,27 +288,20 @@ func (sc *StatementContext) GetPlanDigest() (normalized string, planDigest *pars
 
 // GetResourceGroupTagger returns the implementation of tikvrpc.ResourceGroupTagger related to self.
 func (sc *StatementContext) GetResourceGroupTagger() tikvrpc.ResourceGroupTagger {
+ normalized, digest := sc.SQLDigest()
+ planDigest := sc.planDigest
  return func(req *tikvrpc.Request) {
  if req == nil {
  return
  }
- req.ResourceGroupTag = sc.GetResourceGroupTagByLabel(
+ if len(normalized) == 0 {
+ return
+ }
+ req.ResourceGroupTag = resourcegrouptag.EncodeResourceGroupTag(digest, planDigest,
  resourcegrouptag.GetResourceGroupLabelByKey(resourcegrouptag.GetFirstKeyFromRequest(req)))
  }
 }
 
-// GetResourceGroupTagByLabel gets the resource group of the statement based on the label.
-func (sc *StatementContext) GetResourceGroupTagByLabel(label tipb.ResourceGroupTagLabel) []byte {
- if sc == nil {
- return nil
- }
- normalized, sqlDigest := sc.SQLDigest()
- if len(normalized) == 0 {
- return nil
- }
- return resourcegrouptag.EncodeResourceGroupTag(sqlDigest, sc.planDigest, label)
-}
-
 // SetPlanDigest sets the normalized plan and plan digest.
 func (sc *StatementContext) SetPlanDigest(normalized string, planDigest *parser.Digest) {
  if planDigest != nil {

diff --git a/sessionctx/stmtctx/stmtctx_test.go b/sessionctx/stmtctx/stmtctx_test.go
@@ -21,7 +21,6 @@ import (
 
  "github.com/pingcap/tidb/sessionctx/stmtctx"
  "github.com/pingcap/tidb/util/execdetails"
- "github.com/pingcap/tipb/go-tipb"
  "github.com/stretchr/testify/require"
  "github.com/tikv/client-go/v2/util"
 )
@@ -91,25 +90,3 @@ func TestStatementContextPushDownFLags(t *testing.T) {
  require.Equal(t, tt.out, got)
  }
 }
-
-func TestGetResourceGroupTagByLabel(t *testing.T) {
- ctx := stmtctx.StatementContext{OriginalSQL: "SELECT * FROM t"}
- tagRow := ctx.GetResourceGroupTagByLabel(tipb.ResourceGroupTagLabel_ResourceGroupTagLabelRow)
- tagIndex := ctx.GetResourceGroupTagByLabel(tipb.ResourceGroupTagLabel_ResourceGroupTagLabelIndex)
- tagUnknown := ctx.GetResourceGroupTagByLabel(tipb.ResourceGroupTagLabel_ResourceGroupTagLabelUnknown)
- tagRow2 := ctx.GetResourceGroupTagByLabel(tipb.ResourceGroupTagLabel_ResourceGroupTagLabelRow)
- tagIndex2 := ctx.GetResourceGroupTagByLabel(tipb.ResourceGroupTagLabel_ResourceGroupTagLabelIndex)
- tagUnknown2 := ctx.GetResourceGroupTagByLabel(tipb.ResourceGroupTagLabel_ResourceGroupTagLabelUnknown)
- require.NotEmpty(t, tagRow)
- require.NotEmpty(t, tagIndex)
- require.NotEmpty(t, tagUnknown)
- require.NotEmpty(t, tagRow2)
- require.NotEmpty(t, tagIndex2)
- require.NotEmpty(t, tagUnknown2)
- require.Equal(t, &tagRow, &tagRow2) // mem addr
- require.Equal(t, &tagIndex, &tagIndex2)
- require.Equal(t, &tagUnknown, &tagUnknown2)
- require.NotEqual(t, &tagRow, &tagIndex)
- require.NotEqual(t, &tagRow, &tagUnknown)
- require.NotEqual(t, &tagIndex, &tagUnknown)
-}