diff --git a/sessionctx/stmtctx/stmtctx.go b/sessionctx/stmtctx/stmtctx.go index bdf93278a59d5..100dbb79bee3f 100644 --- a/sessionctx/stmtctx/stmtctx.go +++ b/sessionctx/stmtctx/stmtctx.go @@ -173,12 +173,7 @@ type StatementContext struct { // stmtCache is used to store some statement-related values. stmtCache map[StmtCacheKey]interface{} - // resourceGroupTagWithRow cache for the current statement resource group tag (with `Row` label). - resourceGroupTagWithRow atomic.Value - // resourceGroupTagWithIndex cache for the current statement resource group tag (with `Index` label). - resourceGroupTagWithIndex atomic.Value - // resourceGroupTagWithUnknown cache for the current statement resource group tag (with `Unknown` label). - resourceGroupTagWithUnknown atomic.Value + // Map to store all CTE storages of current SQL. // Will clean up at the end of the execution. CTEStorageMap interface{} @@ -287,6 +282,9 @@ func (sc *StatementContext) GetPlanDigest() (normalized string, planDigest *pars // GetResourceGroupTagger returns the implementation of tikvrpc.ResourceGroupTagger related to self. func (sc *StatementContext) GetResourceGroupTagger() tikvrpc.ResourceGroupTagger { return func(req *tikvrpc.Request) { + if req == nil { + return + } req.ResourceGroupTag = sc.GetResourceGroupTagByLabel( resourcegrouptag.GetResourceGroupLabelByKey(resourcegrouptag.GetFirstKeyFromRequest(req))) } @@ -294,37 +292,14 @@ func (sc *StatementContext) GetResourceGroupTagger() tikvrpc.ResourceGroupTagger // GetResourceGroupTagByLabel gets the resource group of the statement based on the label. func (sc *StatementContext) GetResourceGroupTagByLabel(label tipb.ResourceGroupTagLabel) []byte { - switch label { - case tipb.ResourceGroupTagLabel_ResourceGroupTagLabelRow: - v := sc.resourceGroupTagWithRow.Load() - if v != nil { - return v.([]byte) - } - case tipb.ResourceGroupTagLabel_ResourceGroupTagLabelIndex: - v := sc.resourceGroupTagWithIndex.Load() - if v != nil { - return v.([]byte) - } - case tipb.ResourceGroupTagLabel_ResourceGroupTagLabelUnknown: - v := sc.resourceGroupTagWithUnknown.Load() - if v != nil { - return v.([]byte) - } + if sc == nil { + return nil } normalized, sqlDigest := sc.SQLDigest() if len(normalized) == 0 { return nil } - newTag := resourcegrouptag.EncodeResourceGroupTag(sqlDigest, sc.planDigest, label) - switch label { - case tipb.ResourceGroupTagLabel_ResourceGroupTagLabelRow: - sc.resourceGroupTagWithRow.Store(newTag) - case tipb.ResourceGroupTagLabel_ResourceGroupTagLabelIndex: - sc.resourceGroupTagWithIndex.Store(newTag) - case tipb.ResourceGroupTagLabel_ResourceGroupTagLabelUnknown: - sc.resourceGroupTagWithUnknown.Store(newTag) - } - return newTag + return resourcegrouptag.EncodeResourceGroupTag(sqlDigest, sc.planDigest, label) } // SetPlanDigest sets the normalized plan and plan digest.