*: use pprof profile to collect CPU time group by SQL and plan digest

config/config.go

@@ -667,7 +667,7 @@ var defaultConf = Config{
 		HistorySize:         24,
 	},
 	TopSQL: TopSQL{
-		Enable:          true,
+		Enable:          false,
 		RefreshInterval: 1,
 		MaxStmtCount:    5000,
 	},

executor/adapter.go

@@ -56,6 +56,7 @@ import (
 	"github.com/pingcap/tidb/util/sqlexec"
 	"github.com/pingcap/tidb/util/stmtsummary"
 	"github.com/pingcap/tidb/util/stringutil"
+	"github.com/pingcap/tidb/util/tracecpu"
 
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
@@ -213,6 +214,7 @@ func (a *ExecStmt) PointGet(ctx context.Context, is infoschema.InfoSchema) (*rec
 		defer span1.Finish()
 		ctx = opentracing.ContextWithSpan(ctx, span1)
 	}
+	ctx = a.wrapAndSetPlanLabelForTopSQL(ctx)
 	startTs := uint64(math.MaxUint64)
 	err := a.Ctx.InitTxnWithStartTS(startTs)
 	if err != nil {
@@ -288,6 +290,18 @@ func (a *ExecStmt) RebuildPlan(ctx context.Context) (int64, error) {
 	return a.InfoSchema.SchemaMetaVersion(), nil
 }
 
+func (a *ExecStmt) wrapAndSetPlanLabelForTopSQL(ctx context.Context) context.Context {
+	if a.Plan == nil || !config.TopSQLEnabled() {
+		return ctx
+	}
+	_, sqlDigest := a.Ctx.GetSessionVars().StmtCtx.SQLDigest()
+	normalizedPlan, planDigest := getPlanDigest(a.Ctx, a.Plan)
+	if len(planDigest) > 0 {
+		ctx = tracecpu.SetSQLAndPlanLabels(ctx, sqlDigest.String(), planDigest, normalizedPlan)
+	}
+	return ctx
+}
+
 // Exec builds an Executor from a plan. If the Executor doesn't return result,
 // like the INSERT, UPDATE statements, it executes in this function, if the Executor returns
 // result, execution is done after this function returns, in the returned sqlexec.RecordSet Next method.
@@ -357,8 +371,8 @@ func (a *ExecStmt) Exec(ctx context.Context) (_ sqlexec.RecordSet, err error) {
 	if err != nil {
 		return nil, err
 	}
-
-	getPlanDigest(a.Ctx, a.Plan)
+	// ExecuteExec will rewrite `a.Plan`, so set plan label should be executed after `a.buildExecutor`.
+	ctx = a.wrapAndSetPlanLabelForTopSQL(ctx)
 
 	if err = e.Open(ctx); err != nil {
 		terror.Call(e.Close)
@@ -951,7 +965,7 @@ func (a *ExecStmt) LogSlowQuery(txnTS uint64, succ bool, hasMoreResults bool) {
 	statsInfos := plannercore.GetStatsInfo(a.Plan)
 	memMax := sessVars.StmtCtx.MemTracker.MaxConsumed()
 	diskMax := sessVars.StmtCtx.DiskTracker.MaxConsumed()
-	planDigest := getPlanDigest(a.Ctx, a.Plan)
+	_, planDigest := getPlanDigest(a.Ctx, a.Plan)
 	slowItems := &variable.SlowQueryLogItems{
 		TxnTS:             txnTS,
 		SQL:               sql.String(),
@@ -1043,15 +1057,15 @@ func getPlanTree(sctx sessionctx.Context, p plannercore.Plan) string {
 }
 
 // getPlanDigest will try to get the select plan tree if the plan is select or the select plan of delete/update/insert statement.
-func getPlanDigest(sctx sessionctx.Context, p plannercore.Plan) string {
+func getPlanDigest(sctx sessionctx.Context, p plannercore.Plan) (string, string) {
 	sc := sctx.GetSessionVars().StmtCtx
-	_, planDigest := sc.GetPlanDigest()
-	if planDigest != nil {
-		return planDigest.String()
+	normalized, planDigest := sc.GetPlanDigest()
+	if len(normalized) > 0 && planDigest != nil {
+		return normalized, planDigest.String()
 	}
-	normalized, planDigest := plannercore.NormalizePlan(p)
+	normalized, planDigest = plannercore.NormalizePlan(p)
 	sc.SetPlanDigest(normalized, planDigest)
-	return planDigest.String()
+	return normalized, planDigest.String()
 }
 
 // getEncodedPlan gets the encoded plan, and generates the hint string if indicated.
@@ -1125,11 +1139,11 @@ func (a *ExecStmt) SummaryStmt(succ bool) {
 	var planDigestGen func() string
 	if a.Plan.TP() == plancodec.TypePointGet {
 		planDigestGen = func() string {
-			planDigest := getPlanDigest(a.Ctx, a.Plan)
+			_, planDigest := getPlanDigest(a.Ctx, a.Plan)
 			return planDigest
 		}
 	} else {
-		planDigest = getPlanDigest(a.Ctx, a.Plan)
+		_, planDigest = getPlanDigest(a.Ctx, a.Plan)
 	}
 
 	execDetail := stmtCtx.GetExecDetails()

executor/executor.go

@@ -66,6 +66,7 @@ import (
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/memory"
 	"github.com/pingcap/tidb/util/resourcegrouptag"
+	"github.com/pingcap/tidb/util/tracecpu"
 	"go.uber.org/zap"
 )
 
@@ -1652,9 +1653,15 @@ func ResetContextOfStmt(ctx sessionctx.Context, s ast.StmtNode) (err error) {
 		sc.MemTracker.SetActionOnExceed(action)
 	}
 	if execStmt, ok := s.(*ast.ExecuteStmt); ok {
-		s, err = planner.GetPreparedStmt(execStmt, vars)
+		prepareStmt, err := planner.GetPreparedStmt(execStmt, vars)
 		if err != nil {
-			return
+			return err
+		}
+		s = prepareStmt.PreparedAst.Stmt
+		sc.InitSQLDigest(prepareStmt.NormalizedSQL, prepareStmt.SQLDigest)
+		// For `execute stmt` SQL, should reset the SQL digest with the prepare SQL digest.
+		if config.TopSQLEnabled() && prepareStmt.SQLDigest != nil {
+			tracecpu.SetSQLLabels(context.Background(), prepareStmt.NormalizedSQL, prepareStmt.SQLDigest.String())
 		}
 	}
 	// execute missed stmtID uses empty sql

planner/optimize.go

@@ -44,7 +44,7 @@ import (
 )
 
 // GetPreparedStmt extract the prepared statement from the execute statement.
-func GetPreparedStmt(stmt *ast.ExecuteStmt, vars *variable.SessionVars) (ast.StmtNode, error) {
+func GetPreparedStmt(stmt *ast.ExecuteStmt, vars *variable.SessionVars) (*plannercore.CachedPrepareStmt, error) {
 	var ok bool
 	execID := stmt.ExecID
 	if stmt.Name != "" {
@@ -57,20 +57,20 @@ func GetPreparedStmt(stmt *ast.ExecuteStmt, vars *variable.SessionVars) (ast.Stm
 		if !ok {
 			return nil, errors.Errorf("invalid CachedPrepareStmt type")
 		}
-		return preparedObj.PreparedAst.Stmt, nil
+		return preparedObj, nil
 	}
 	return nil, plannercore.ErrStmtNotFound
 }
 
 // IsReadOnly check whether the ast.Node is a read only statement.
 func IsReadOnly(node ast.Node, vars *variable.SessionVars) bool {
 	if execStmt, isExecStmt := node.(*ast.ExecuteStmt); isExecStmt {
-		s, err := GetPreparedStmt(execStmt, vars)
+		prepareStmt, err := GetPreparedStmt(execStmt, vars)
 		if err != nil {
 			logutil.BgLogger().Warn("GetPreparedStmt failed", zap.Error(err))
 			return false
 		}
-		return ast.IsReadOnly(s)
+		return ast.IsReadOnly(prepareStmt.PreparedAst.Stmt)
 	}
 	return ast.IsReadOnly(node)
 }

server/conn.go

@@ -78,12 +78,14 @@ import (
 	storeerr "github.com/pingcap/tidb/store/driver/error"
 	"github.com/pingcap/tidb/store/tikv/util"
 	"github.com/pingcap/tidb/tablecodec"
+	util2 "github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/arena"
 	"github.com/pingcap/tidb/util/chunk"
 	"github.com/pingcap/tidb/util/execdetails"
 	"github.com/pingcap/tidb/util/hack"
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/memory"
+	"github.com/pingcap/tidb/util/tracecpu"
 	"github.com/prometheus/client_golang/prometheus"
 	"go.uber.org/zap"
 )
@@ -906,14 +908,6 @@ func (cc *clientConn) ShutdownOrNotify() bool {
 	return false
 }
 
-func queryStrForLog(query string) string {
-	const size = 4096
-	if len(query) > size {
-		return query[:size] + fmt.Sprintf("(len: %d)", len(query))
-	}
-	return query
-}
-
 func errStrForLog(err error, enableRedactLog bool) string {
 	if enableRedactLog {
 		// currently, only ErrParse is considered when enableRedactLog because it may contain sensitive information like
@@ -1025,7 +1019,13 @@ func (cc *clientConn) dispatch(ctx context.Context, data []byte) error {
 	cc.lastPacket = data
 	cmd := data[0]
 	data = data[1:]
-	if variable.EnablePProfSQLCPU.Load() {
+	if config.TopSQLEnabled() {
+		normalizedSQL, digest := getLastStmtInConn{cc}.pprofLabelNormalizedAndDigest()
+		if len(digest) > 0 {
+			defer tracecpu.ResetGoroutineLabelsWithOriginalCtx(ctx)
+			ctx = tracecpu.SetSQLLabels(ctx, normalizedSQL, digest)
+		}
+	} else if variable.EnablePProfSQLCPU.Load() {
 		label := getLastStmtInConn{cc}.PProfLabel()
 		if len(label) > 0 {
 			defer pprof.SetGoroutineLabels(ctx)
@@ -2125,10 +2125,10 @@ func (cc getLastStmtInConn) String() string {
 		if cc.ctx.GetSessionVars().EnableRedactLog {
 			sql = parser.Normalize(sql)
 		}
-		return queryStrForLog(sql)
+		return util2.QueryStrForLog(sql)
 	case mysql.ComStmtExecute, mysql.ComStmtFetch:
 		stmtID := binary.LittleEndian.Uint32(data[0:4])
-		return queryStrForLog(cc.preparedStmt2String(stmtID))
+		return util2.QueryStrForLog(cc.preparedStmt2String(stmtID))
 	case mysql.ComStmtClose, mysql.ComStmtReset:
 		stmtID := binary.LittleEndian.Uint32(data[0:4])
 		return mysql.Command2Str[cmd] + " " + strconv.Itoa(int(stmtID))
@@ -2156,11 +2156,43 @@ func (cc getLastStmtInConn) PProfLabel() string {
 	case mysql.ComStmtReset:
 		return "ResetStmt"
 	case mysql.ComQuery, mysql.ComStmtPrepare:
-		return parser.Normalize(queryStrForLog(string(hack.String(data))))
+		return parser.Normalize(util2.QueryStrForLog(string(hack.String(data))))
 	case mysql.ComStmtExecute, mysql.ComStmtFetch:
 		stmtID := binary.LittleEndian.Uint32(data[0:4])
-		return queryStrForLog(cc.preparedStmt2StringNoArgs(stmtID))
+		return util2.QueryStrForLog(cc.preparedStmt2StringNoArgs(stmtID))
 	default:
 		return ""
 	}
 }
+
+// PProfLabelNormalizedAndDigest return sql and sql_digest label used to tag pprof.
+func (cc getLastStmtInConn) pprofLabelNormalizedAndDigest() (string, string) {
+	if len(cc.lastPacket) == 0 {
+		return "", ""
+	}
+	cmd, data := cc.lastPacket[0], cc.lastPacket[1:]
+	switch cmd {
+	case mysql.ComInitDB:
+		return "UseDB", ""
+	case mysql.ComFieldList:
+		return "ListFields", ""
+	case mysql.ComStmtClose:
+		return "CloseStmt", ""
+	case mysql.ComStmtReset:
+		return "ResetStmt", ""
+	case mysql.ComQuery, mysql.ComStmtPrepare:
+		normalized, digest := parser.NormalizeDigest(string(hack.String(data)))
+		return normalized, digest.String()
+	case mysql.ComStmtExecute, mysql.ComStmtFetch:
+		stmtID := binary.LittleEndian.Uint32(data[0:4])
+		prepareObj, _ := cc.preparedStmtID2CachePreparedStmt(stmtID)
+		if prepareObj != nil && prepareObj.SQLDigest != nil {
+			return prepareObj.NormalizedSQL, prepareObj.SQLDigest.String()
+		}
+		str := cc.preparedStmt2StringNoArgs(stmtID)
+		normalized, digest := parser.NormalizeDigest(str)
+		return normalized, digest.String()
+	default:
+		return "", ""
+	}
+}

server/conn_stmt.go

@@ -680,14 +680,30 @@ func (cc *clientConn) preparedStmt2StringNoArgs(stmtID uint32) string {
 	if sv == nil {
 		return ""
 	}
+	preparedObj, invalid := cc.preparedStmtID2CachePreparedStmt(stmtID)
+	if invalid {
+		return "invalidate CachedPrepareStmt type, ID: " + strconv.FormatUint(uint64(stmtID), 10)
+	}
+	if preparedObj == nil {
+		return "prepared statement not found, ID: " + strconv.FormatUint(uint64(stmtID), 10)
+	}
+	return preparedObj.PreparedAst.Stmt.Text()
+}
+
+func (cc *clientConn) preparedStmtID2CachePreparedStmt(stmtID uint32) (_ *plannercore.CachedPrepareStmt, invalid bool) {
+	sv := cc.ctx.GetSessionVars()
+	if sv == nil {
+		return nil, false
+	}
 	preparedPointer, ok := sv.PreparedStmts[stmtID]
 	if !ok {
-		return "prepared statement not found, ID: " + strconv.FormatUint(uint64(stmtID), 10)
+		// not found
+		return nil, false
 	}
 	preparedObj, ok := preparedPointer.(*plannercore.CachedPrepareStmt)
 	if !ok {
-		return "invalidate CachedPrepareStmt type, ID: " + strconv.FormatUint(uint64(stmtID), 10)
+		// invalid cache. should never happen.
+		return nil, true
 	}
-	preparedAst := preparedObj.PreparedAst
-	return preparedAst.Stmt.Text()
+	return preparedObj, false
 }

server/http_status.go

@@ -43,6 +43,7 @@ import (
 	"github.com/pingcap/tidb/util"
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/printer"
+	"github.com/pingcap/tidb/util/tracecpu"
 	"github.com/pingcap/tidb/util/versioninfo"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 	"github.com/soheilhy/cmux"
@@ -184,7 +185,7 @@ func (s *Server) startHTTPServer() {
 
 	serverMux.HandleFunc("/debug/pprof/", pprof.Index)
 	serverMux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
-	serverMux.HandleFunc("/debug/pprof/profile", pprof.Profile)
+	serverMux.HandleFunc("/debug/pprof/profile", tracecpu.ProfileHTTPHandler)
 	serverMux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
 	serverMux.HandleFunc("/debug/pprof/trace", pprof.Trace)
 	serverMux.HandleFunc("/debug/gogc", func(w http.ResponseWriter, r *http.Request) {
@@ -251,7 +252,7 @@ func (s *Server) startHTTPServer() {
 			serveError(w, http.StatusInternalServerError, fmt.Sprintf("Create zipped %s fail: %v", "profile", err))
 			return
 		}
-		if err := rpprof.StartCPUProfile(fw); err != nil {
+		if err := tracecpu.StartCPUProfile(fw); err != nil {
 			serveError(w, http.StatusInternalServerError,
 				fmt.Sprintf("Could not enable CPU profiling: %s", err))
 			return
@@ -261,7 +262,11 @@ func (s *Server) startHTTPServer() {
 			sec = 10
 		}
 		sleepWithCtx(r.Context(), time.Duration(sec)*time.Second)
-		rpprof.StopCPUProfile()
+		err = tracecpu.StopCPUProfile()
+		if err != nil {
+			serveError(w, http.StatusInternalServerError, fmt.Sprintf("Create zipped %s fail: %v", "config", err))
+			return
+		}
 
 		// dump config
 		fw, err = zw.Create("config")

server/sql_info_fetcher.go

@@ -20,7 +20,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
-	"runtime/pprof"
 	"strconv"
 	"strings"
 	"time"
@@ -35,6 +34,7 @@ import (
 	"github.com/pingcap/tidb/session"
 	"github.com/pingcap/tidb/statistics/handle"
 	"github.com/pingcap/tidb/util/sqlexec"
+	"github.com/pingcap/tidb/util/tracecpu"
 )
 
 type sqlInfoFetcher struct {
@@ -275,13 +275,13 @@ func (sh *sqlInfoFetcher) getExplainAnalyze(ctx context.Context, sql string, res
 }
 
 func (sh *sqlInfoFetcher) catchCPUProfile(ctx context.Context, sec int, buf *bytes.Buffer, errChan chan<- error) {
-	if err := pprof.StartCPUProfile(buf); err != nil {
+	if err := tracecpu.StartCPUProfile(buf); err != nil {
 		errChan <- err
 		return
 	}
 	sleepWithCtx(ctx, time.Duration(sec)*time.Second)
-	pprof.StopCPUProfile()
-	errChan <- nil
+	err := tracecpu.StopCPUProfile()
+	errChan <- err
 }
 
 func (sh *sqlInfoFetcher) getStatsForTable(pair tableNamePair) (*handle.JSONTable, error) {