*: reduce mem allocs caused by GetTextToLog (pingcap#54111)

ref pingcap#54047

pkg/executor/adapter.go

@@ -1288,21 +1288,23 @@ func (a *ExecStmt) logAudit() {
 
 // FormatSQL is used to format the original SQL, e.g. truncating long SQL, appending prepared arguments.
 func FormatSQL(sql string) stringutil.StringerFunc {
-	return func() string {
-		length := len(sql)
-		maxQueryLen := variable.QueryLogMaxLen.Load()
-		if maxQueryLen <= 0 {
-			return QueryReplacer.Replace(sql) // no limit
-		}
-		if int32(length) > maxQueryLen {
-			var result strings.Builder
-			result.Grow(int(maxQueryLen))
-			result.WriteString(sql[:maxQueryLen])
-			fmt.Fprintf(&result, "(len:%d)", length)
-			return QueryReplacer.Replace(result.String())
-		}
-		return QueryReplacer.Replace(sql)
+	return func() string { return formatSQL(sql) }
+}
+
+func formatSQL(sql string) string {
+	length := len(sql)
+	maxQueryLen := variable.QueryLogMaxLen.Load()
+	if maxQueryLen <= 0 {
+		return QueryReplacer.Replace(sql) // no limit
+	}
+	if int32(length) > maxQueryLen {
+		var result strings.Builder
+		result.Grow(int(maxQueryLen))
+		result.WriteString(sql[:maxQueryLen])
+		fmt.Fprintf(&result, "(len:%d)", length)
+		return QueryReplacer.Replace(result.String())
 	}
+	return QueryReplacer.Replace(sql)
 }
 
 func getPhaseDurationObserver(phase string, internal bool) prometheus.Observer {
@@ -1408,7 +1410,7 @@ func (a *ExecStmt) FinishExecuteStmt(txnTS uint64, err error, hasMoreResults boo
 			metrics.TiFlashQueryTotalCounter.WithLabelValues(metrics.ExecuteErrorToLabel(err), metrics.LblError).Inc()
 		}
 	}
-	sessVars.PrevStmt = FormatSQL(a.GetTextToLog(false))
+	a.updatePrevStmt()
 	a.recordLastQueryInfo(err)
 	a.observePhaseDurations(sessVars.InRestrictedSQL, execDetail.CommitDetail)
 	executeDuration := time.Since(sessVars.StartTime) - sessVars.DurationCompile
@@ -1935,7 +1937,7 @@ func (a *ExecStmt) SummaryStmt(succ bool) {
 	copTaskInfo := stmtCtx.CopTasksDetails()
 	memMax := sessVars.MemTracker.MaxConsumed()
 	diskMax := sessVars.DiskTracker.MaxConsumed()
-	sql := a.GetTextToLog(false)
+	sql := a.getLazyStmtText()
 	var stmtDetail execdetails.StmtExecDetails
 	stmtDetailRaw := a.GoCtx.Value(execdetails.StmtExecDetailKey)
 	if stmtDetailRaw != nil {
@@ -1973,7 +1975,7 @@ func (a *ExecStmt) SummaryStmt(succ bool) {
 
 	stmtExecInfo := &stmtsummary.StmtExecInfo{
 		SchemaName:          strings.ToLower(sessVars.CurrentDB),
-		OriginalSQL:         sql,
+		OriginalSQL:         &sql,
 		Charset:             charset,
 		Collation:           collation,
 		NormalizedSQL:       normalizedSQL,
@@ -2035,6 +2037,43 @@ func (a *ExecStmt) GetTextToLog(keepHint bool) string {
 	return sql
 }
 
+// getLazyText is equivalent to `a.GetTextToLog(false)`. Note that the s.Params is a shallow copy of
+// `sessVars.PlanCacheParams`, so you can only use the lazy text within the current stmt context.
+func (a *ExecStmt) getLazyStmtText() (s variable.LazyStmtText) {
+	sessVars := a.Ctx.GetSessionVars()
+	rmode := sessVars.EnableRedactLog
+	if rmode == errors.RedactLogEnable {
+		sql, _ := sessVars.StmtCtx.SQLDigest()
+		s.SetText(sql)
+	} else if sensitiveStmt, ok := a.StmtNode.(ast.SensitiveStmtNode); ok {
+		sql := sensitiveStmt.SecureText()
+		s.SetText(sql)
+	} else {
+		s.Redact = rmode
+		s.SQL = sessVars.StmtCtx.OriginalSQL
+		s.Params = *sessVars.PlanCacheParams
+	}
+	return
+}
+
+// updatePrevStmt is equivalent to `sessVars.PrevStmt = FormatSQL(a.GetTextToLog(false))`
+func (a *ExecStmt) updatePrevStmt() {
+	sessVars := a.Ctx.GetSessionVars()
+	if sessVars.PrevStmt == nil {
+		sessVars.PrevStmt = &variable.LazyStmtText{Format: formatSQL}
+	}
+	rmode := sessVars.EnableRedactLog
+	if rmode == errors.RedactLogEnable {
+		sql, _ := sessVars.StmtCtx.SQLDigest()
+		sessVars.PrevStmt.SetText(sql)
+	} else if sensitiveStmt, ok := a.StmtNode.(ast.SensitiveStmtNode); ok {
+		sql := sensitiveStmt.SecureText()
+		sessVars.PrevStmt.SetText(sql)
+	} else {
+		sessVars.PrevStmt.Update(rmode, sessVars.StmtCtx.OriginalSQL, sessVars.PlanCacheParams)
+	}
+}
+
 func (a *ExecStmt) observeStmtBeginForTopSQL(ctx context.Context) context.Context {
 	vars := a.Ctx.GetSessionVars()
 	sc := vars.StmtCtx

pkg/sessionctx/variable/BUILD.bazel

@@ -54,6 +54,7 @@ go_library(
         "//pkg/util/mathutil",
         "//pkg/util/memory",
         "//pkg/util/paging",
+        "//pkg/util/redact",
         "//pkg/util/replayer",
         "//pkg/util/rowcodec",
         "//pkg/util/servicescope",

pkg/sessionctx/variable/session.go

@@ -56,6 +56,7 @@ import (
 	"github.com/pingcap/tidb/pkg/util/kvcache"
 	"github.com/pingcap/tidb/pkg/util/mathutil"
 	"github.com/pingcap/tidb/pkg/util/memory"
+	"github.com/pingcap/tidb/pkg/util/redact"
 	"github.com/pingcap/tidb/pkg/util/replayer"
 	"github.com/pingcap/tidb/pkg/util/rowcodec"
 	"github.com/pingcap/tidb/pkg/util/sqlkiller"
@@ -1136,7 +1137,7 @@ type SessionVars struct {
 	DurationWaitTS time.Duration
 
 	// PrevStmt is used to store the previous executed statement in the current session.
-	PrevStmt fmt.Stringer
+	PrevStmt *LazyStmtText
 
 	// prevStmtDigest is used to store the digest of the previous statement in the current session.
 	prevStmtDigest string
@@ -1926,6 +1927,48 @@ func (p *PlanCacheParamList) AllParamValues() []types.Datum {
 	return p.paramValues
 }
 
+// LazyStmtText represents the sql text of a stmt that used in log. It's lazily evaluated to reduce the mem allocs.
+type LazyStmtText struct {
+	text   *string
+	SQL    string
+	Redact string
+	Params PlanCacheParamList
+	Format func(string) string
+}
+
+// SetText sets the text directly.
+func (s *LazyStmtText) SetText(text string) {
+	s.text = &text
+}
+
+// Update resets the lazy text and leads to re-eval for next `s.String()`. It copies params so it's safe to use
+// `SessionVars.PlanCacheParams` directly without worrying about the params get reset later.
+func (s *LazyStmtText) Update(redact string, sql string, params *PlanCacheParamList) {
+	s.text = nil
+	s.SQL = sql
+	s.Redact = redact
+	s.Params.Reset()
+	if params != nil {
+		s.Params.forNonPrepCache = params.forNonPrepCache
+		s.Params.paramValues = append(s.Params.paramValues, params.paramValues...)
+	}
+}
+
+// String implements fmt.Stringer.
+func (s *LazyStmtText) String() string {
+	if s == nil {
+		return ""
+	}
+	if s.text == nil {
+		text := redact.String(s.Redact, s.SQL+s.Params.String())
+		if s.Format != nil {
+			text = s.Format(text)
+		}
+		s.text = &text
+	}
+	return *s.text
+}
+
 // ConnectionInfo presents the connection information, which is mainly used by audit logs.
 type ConnectionInfo struct {
 	ConnectionID      uint64

pkg/util/stmtsummary/BUILD.bazel

@@ -51,6 +51,7 @@ go_test(
         "//pkg/util",
         "//pkg/util/execdetails",
         "//pkg/util/plancodec",
+        "//pkg/util/stringutil",
         "@com_github_pingcap_log//:log",
         "@com_github_stretchr_testify//require",
         "@com_github_tikv_client_go_v2//util",

pkg/util/stmtsummary/statement_summary.go

@@ -225,7 +225,7 @@ type stmtSummaryByDigestElement struct {
 // StmtExecInfo records execution information of each statement.
 type StmtExecInfo struct {
 	SchemaName          string
-	OriginalSQL         string
+	OriginalSQL         fmt.Stringer
 	Charset             string
 	Collation           string
 	NormalizedSQL       string
@@ -665,7 +665,7 @@ func newStmtSummaryByDigestElement(sei *StmtExecInfo, beginTime int64, intervalS
 	}
 	ssElement := &stmtSummaryByDigestElement{
 		beginTime: beginTime,
-		sampleSQL: formatSQL(sei.OriginalSQL),
+		sampleSQL: formatSQL(sei.OriginalSQL.String()),
 		charset:   sei.Charset,
 		collation: sei.Collation,
 		// PrevSQL is already truncated to cfg.Log.QueryLogMaxLen.

pkg/util/stmtsummary/statement_summary_test.go

@@ -30,6 +30,7 @@ import (
 	tidbutil "github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/execdetails"
 	"github.com/pingcap/tidb/pkg/util/plancodec"
+	"github.com/pingcap/tidb/pkg/util/stringutil"
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/client-go/v2/util"
 )
@@ -84,7 +85,7 @@ func TestAddStatement(t *testing.T) {
 	expectedSummaryElement := stmtSummaryByDigestElement{
 		beginTime:            now + 60,
 		endTime:              now + 1860,
-		sampleSQL:            stmtExecInfo1.OriginalSQL,
+		sampleSQL:            stmtExecInfo1.OriginalSQL.String(),
 		samplePlan:           samplePlan,
 		indexNames:           stmtExecInfo1.StmtCtx.IndexNames,
 		execCount:            1,
@@ -169,7 +170,7 @@ func TestAddStatement(t *testing.T) {
 	// greater than that of the first statement.
 	stmtExecInfo2 := &StmtExecInfo{
 		SchemaName:     "schema_name",
-		OriginalSQL:    "original_sql2",
+		OriginalSQL:    stringutil.StringerStr("original_sql2"),
 		NormalizedSQL:  "normalized_sql",
 		Digest:         "digest",
 		PlanDigest:     "plan_digest",
@@ -312,7 +313,7 @@ func TestAddStatement(t *testing.T) {
 	// less than that of the first statement.
 	stmtExecInfo3 := &StmtExecInfo{
 		SchemaName:     "schema_name",
-		OriginalSQL:    "original_sql3",
+		OriginalSQL:    stringutil.StringerStr("original_sql3"),
 		NormalizedSQL:  "normalized_sql",
 		Digest:         "digest",
 		PlanDigest:     "plan_digest",
@@ -618,7 +619,7 @@ func generateAnyExecInfo() *StmtExecInfo {
 
 	stmtExecInfo := &StmtExecInfo{
 		SchemaName:     "schema_name",
-		OriginalSQL:    "original_sql1",
+		OriginalSQL:    stringutil.StringerStr("original_sql1"),
 		NormalizedSQL:  "normalized_sql",
 		Digest:         "digest",
 		PlanDigest:     "plan_digest",
@@ -1023,7 +1024,7 @@ func TestMaxSQLLength(t *testing.T) {
 	str := strings.Repeat("a", length)
 
 	stmtExecInfo1 := generateAnyExecInfo()
-	stmtExecInfo1.OriginalSQL = str
+	stmtExecInfo1.OriginalSQL = stringutil.StringerStr(str)
 	stmtExecInfo1.NormalizedSQL = str
 	ssMap.AddStatement(stmtExecInfo1)
 
@@ -1123,7 +1124,7 @@ func TestDisableStmtSummary(t *testing.T) {
 	ssMap.beginTimeForCurInterval = now + 60
 
 	stmtExecInfo2 := stmtExecInfo1
-	stmtExecInfo2.OriginalSQL = "original_sql2"
+	stmtExecInfo2.OriginalSQL = stringutil.StringerStr("original_sql2")
 	stmtExecInfo2.NormalizedSQL = "normalized_sql2"
 	stmtExecInfo2.Digest = "digest2"
 	ssMap.AddStatement(stmtExecInfo2)
@@ -1197,7 +1198,7 @@ func TestGetMoreThanCntBindableStmt(t *testing.T) {
 	ssMap := newStmtSummaryByDigestMap()
 
 	stmtExecInfo1 := generateAnyExecInfo()
-	stmtExecInfo1.OriginalSQL = "insert 1"
+	stmtExecInfo1.OriginalSQL = stringutil.StringerStr("insert 1")
 	stmtExecInfo1.NormalizedSQL = "insert ?"
 	stmtExecInfo1.StmtCtx.StmtType = "Insert"
 	ssMap.AddStatement(stmtExecInfo1)

pkg/util/stmtsummary/v2/BUILD.bazel

@@ -26,6 +26,7 @@ go_library(
         "//pkg/util/plancodec",
         "//pkg/util/set",
         "//pkg/util/stmtsummary",
+        "//pkg/util/stringutil",
         "@com_github_pingcap_log//:log",
         "@com_github_tikv_client_go_v2//util",
         "@org_uber_go_atomic//:atomic",

pkg/util/stmtsummary/v2/record.go

@@ -27,6 +27,7 @@ import (
 	"github.com/pingcap/tidb/pkg/util/execdetails"
 	"github.com/pingcap/tidb/pkg/util/plancodec"
 	"github.com/pingcap/tidb/pkg/util/stmtsummary"
+	"github.com/pingcap/tidb/pkg/util/stringutil"
 	"github.com/tikv/client-go/v2/util"
 )
 
@@ -201,7 +202,7 @@ func NewStmtRecord(info *stmtsummary.StmtExecInfo) *StmtRecord {
 		NormalizedSQL: info.NormalizedSQL,
 		TableNames:    tableNames,
 		IsInternal:    info.IsInternal,
-		SampleSQL:     formatSQL(info.OriginalSQL),
+		SampleSQL:     formatSQL(info.OriginalSQL.String()),
 		Charset:       info.Charset,
 		Collation:     info.Collation,
 		// PrevSQL is already truncated to cfg.Log.QueryLogMaxLen.
@@ -610,7 +611,7 @@ func GenerateStmtExecInfo4Test(digest string) *stmtsummary.StmtExecInfo {
 
 	stmtExecInfo := &stmtsummary.StmtExecInfo{
 		SchemaName:     "schema_name",
-		OriginalSQL:    "original_sql1",
+		OriginalSQL:    stringutil.StringerStr("original_sql1"),
 		NormalizedSQL:  "normalized_sql",
 		Digest:         digest,
 		PlanDigest:     "plan_digest",

pkg/util/stmtsummary/v2/record_test.go

@@ -30,7 +30,7 @@ func TestStmtRecord(t *testing.T) {
 	require.Equal(t, info.NormalizedSQL, record1.NormalizedSQL)
 	require.Equal(t, "db1.tb1,db2.tb2", record1.TableNames)
 	require.Equal(t, info.IsInternal, record1.IsInternal)
-	require.Equal(t, formatSQL(info.OriginalSQL), record1.SampleSQL)
+	require.Equal(t, formatSQL(info.OriginalSQL.String()), record1.SampleSQL)
 	require.Equal(t, info.Charset, record1.Charset)
 	require.Equal(t, info.Collation, record1.Collation)
 	require.Equal(t, info.PrevSQL, record1.PrevSQL)