server: fix the missing type of query total counter for resource groups above 2 (#52606)

close #52605

Author: nolouch

pkg/server/BUILD.bazel

@@ -200,6 +200,7 @@ go_test(
         "@com_github_pingcap_errors//:errors",
         "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_pingcap_kvproto//pkg/metapb",
+        "@com_github_prometheus_client_golang//prometheus/testutil",
         "@com_github_stretchr_testify//require",
         "@com_github_tikv_client_go_v2//error",
         "@com_github_tikv_client_go_v2//testutils",

pkg/server/conn.go

@@ -1206,7 +1206,7 @@ func (cc *clientConn) addMetrics(cmd byte, startTime time.Time, err error) {
 	if counter != nil {
 		counter.Inc()
 	} else {
-		label := strconv.Itoa(int(cmd))
+		label := server_metrics.CmdToString(cmd)
 		if err != nil {
 			metrics.QueryTotalCounter.WithLabelValues(label, "Error", resourceGroupName).Inc()
 		} else {

pkg/server/conn_test.go

@@ -30,11 +30,13 @@ import (
 	"testing"
 	"time"
 
+	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/tidb/pkg/config"
 	"github.com/pingcap/tidb/pkg/domain"
 	"github.com/pingcap/tidb/pkg/extension"
+	"github.com/pingcap/tidb/pkg/metrics"
 	"github.com/pingcap/tidb/pkg/parser/auth"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/server/internal"
@@ -54,6 +56,7 @@ import (
 	"github.com/pingcap/tidb/pkg/util/dbterror/exeerrors"
 	"github.com/pingcap/tidb/pkg/util/plancodec"
 	"github.com/pingcap/tidb/pkg/util/sqlkiller"
+	promtestutils "github.com/prometheus/client_golang/prometheus/testutil"
 	"github.com/stretchr/testify/require"
 	tikverr "github.com/tikv/client-go/v2/error"
 	"github.com/tikv/client-go/v2/testutils"
@@ -2085,3 +2088,40 @@ func TestCloseConn(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+func TestConnAddMetrics(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	re := require.New(t)
+	tk := testkit.NewTestKit(t, store)
+	cc := &clientConn{
+		alloc:      arena.NewAllocator(1024),
+		chunkAlloc: chunk.NewAllocator(),
+		pkt:        internal.NewPacketIOForTest(bufio.NewWriter(bytes.NewBuffer(nil))),
+	}
+	tk.MustExec("use test")
+	cc.SetCtx(&TiDBContext{Session: tk.Session(), stmts: make(map[int]*TiDBStatement)})
+
+	// default
+	cc.addMetrics(mysql.ComQuery, time.Now(), nil)
+	counter := metrics.QueryTotalCounter
+	v := promtestutils.ToFloat64(counter.WithLabelValues("Query", "OK", "default"))
+	require.Equal(t, 1.0, v)
+
+	// rg1
+	cc.getCtx().GetSessionVars().ResourceGroupName = "test_rg1"
+	cc.addMetrics(mysql.ComQuery, time.Now(), nil)
+	re.Equal(promtestutils.ToFloat64(counter.WithLabelValues("Query", "OK", "default")), 1.0)
+	re.Equal(promtestutils.ToFloat64(counter.WithLabelValues("Query", "OK", "test_rg1")), 1.0)
+	/// inc the counter again
+	cc.addMetrics(mysql.ComQuery, time.Now(), nil)
+	re.Equal(promtestutils.ToFloat64(counter.WithLabelValues("Query", "OK", "test_rg1")), 2.0)
+
+	// rg2
+	cc.getCtx().GetSessionVars().ResourceGroupName = "test_rg2"
+	// error
+	cc.addMetrics(mysql.ComQuery, time.Now(), errors.New("unknown error"))
+	re.Equal(promtestutils.ToFloat64(counter.WithLabelValues("Query", "Error", "test_rg2")), 1.0)
+	// ok
+	cc.addMetrics(mysql.ComStmtExecute, time.Now(), nil)
+	re.Equal(promtestutils.ToFloat64(counter.WithLabelValues("StmtExecute", "OK", "test_rg2")), 1.0)
+}

pkg/server/metrics/metrics.go

@@ -15,6 +15,8 @@
 package metrics
 
 import (
+	"strconv"
+
 	"github.com/pingcap/tidb/pkg/domain/resourcegroup"
 	"github.com/pingcap/tidb/pkg/metrics"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
@@ -46,6 +48,39 @@ func init() {
 	InitMetricsVars()
 }
 
+// CmdToString convert command type to string.
+func CmdToString(cmd byte) string {
+	switch cmd {
+	case mysql.ComSleep:
+		return "Sleep"
+	case mysql.ComQuit:
+		return "Quit"
+	case mysql.ComInitDB:
+		return "InitDB"
+	case mysql.ComQuery:
+		return "Query"
+	case mysql.ComPing:
+		return "Ping"
+	case mysql.ComFieldList:
+		return "FieldList"
+	case mysql.ComStmtPrepare:
+		return "StmtPrepare"
+	case mysql.ComStmtExecute:
+		return "StmtExecute"
+	case mysql.ComStmtFetch:
+		return "StmtFetch"
+	case mysql.ComStmtClose:
+		return "StmtClose"
+	case mysql.ComStmtSendLongData:
+		return "StmtSendLongData"
+	case mysql.ComStmtReset:
+		return "StmtReset"
+	case mysql.ComSetOption:
+		return "SetOption"
+	}
+	return strconv.Itoa(int(cmd))
+}
+
 // InitMetricsVars init server metrics vars.
 func InitMetricsVars() {
 	QueryTotalCountOk = []prometheus.Counter{