From 33480e8c8d8874811a5362767c81d40c9a057cb6 Mon Sep 17 00:00:00 2001 From: ShuNing Date: Wed, 3 Jan 2024 14:05:15 +0800 Subject: [PATCH] *: add last ru consumption for tidb_last_query_info (#49769) ref pingcap/tidb#49318 --- pkg/executor/adapter.go | 35 ++++++++++++++++++- pkg/session/session.go | 17 --------- pkg/session/test/variable/BUILD.bazel | 2 +- pkg/session/test/variable/variable_test.go | 22 ++++++++++++ .../sessionstates/session_states.go | 9 ++--- 5 files changed, 62 insertions(+), 23 deletions(-) diff --git a/pkg/executor/adapter.go b/pkg/executor/adapter.go index 24931b0cd5159..dd408d82020d7 100644 --- a/pkg/executor/adapter.go +++ b/pkg/executor/adapter.go @@ -49,6 +49,7 @@ import ( plannercore "github.com/pingcap/tidb/pkg/planner/core" "github.com/pingcap/tidb/pkg/plugin" "github.com/pingcap/tidb/pkg/sessionctx" + "github.com/pingcap/tidb/pkg/sessionctx/sessionstates" "github.com/pingcap/tidb/pkg/sessionctx/stmtctx" "github.com/pingcap/tidb/pkg/sessionctx/variable" "github.com/pingcap/tidb/pkg/sessiontxn" @@ -1404,7 +1405,7 @@ func (a *ExecStmt) FinishExecuteStmt(txnTS uint64, err error, hasMoreResults boo } } sessVars.PrevStmt = FormatSQL(a.GetTextToLog(false)) - + a.recordLastQueryInfo(err) a.observePhaseDurations(sessVars.InRestrictedSQL, execDetail.CommitDetail) executeDuration := time.Since(sessVars.StartTime) - sessVars.DurationCompile if sessVars.InRestrictedSQL { @@ -1449,6 +1450,38 @@ func (a *ExecStmt) FinishExecuteStmt(txnTS uint64, err error, hasMoreResults boo } } +func (a *ExecStmt) recordLastQueryInfo(err error) { + sessVars := a.Ctx.GetSessionVars() + // Record diagnostic information for DML statements + recordLastQuery := false + switch typ := a.StmtNode.(type) { + case *ast.ShowStmt: + recordLastQuery = typ.Tp != ast.ShowSessionStates + case *ast.ExecuteStmt, ast.DMLNode: + recordLastQuery = true + } + if recordLastQuery { + var lastRUConsumption float64 + if ruDetailRaw := a.GoCtx.Value(util.RUDetailsCtxKey); ruDetailRaw != nil { + ruDetail := ruDetailRaw.(*util.RUDetails) + lastRUConsumption = ruDetail.RRU() + ruDetail.WRU() + } + failpoint.Inject("mockRUConsumption", func(_ failpoint.Value) { + lastRUConsumption = float64(len(sessVars.StmtCtx.OriginalSQL)) + }) + // Keep the previous queryInfo for `show session_states` because the statement needs to encode it. + sessVars.LastQueryInfo = sessionstates.QueryInfo{ + TxnScope: sessVars.CheckAndGetTxnScope(), + StartTS: sessVars.TxnCtx.StartTS, + ForUpdateTS: sessVars.TxnCtx.GetForUpdateTS(), + LastRUConsumption: lastRUConsumption, + } + if err != nil { + sessVars.LastQueryInfo.ErrMsg = err.Error() + } + } +} + func (a *ExecStmt) checkPlanReplayerCapture(txnTS uint64) { if kv.GetInternalSourceType(a.GoCtx) == kv.InternalTxnStats { return diff --git a/pkg/session/session.go b/pkg/session/session.go index 5f5300aab2224..a2f44c7912fae 100644 --- a/pkg/session/session.go +++ b/pkg/session/session.go @@ -2362,23 +2362,6 @@ func runStmt(ctx context.Context, se *session, s sqlexec.Statement) (rs sqlexec. sessVars := se.sessionVars - // Record diagnostic information for DML statements - if stmt, ok := s.(*executor.ExecStmt).StmtNode.(ast.DMLNode); ok { - // Keep the previous queryInfo for `show session_states` because the statement needs to encode it. - if showStmt, ok := stmt.(*ast.ShowStmt); !ok || showStmt.Tp != ast.ShowSessionStates { - defer func() { - sessVars.LastQueryInfo = sessionstates.QueryInfo{ - TxnScope: sessVars.CheckAndGetTxnScope(), - StartTS: sessVars.TxnCtx.StartTS, - ForUpdateTS: sessVars.TxnCtx.GetForUpdateTS(), - } - if err != nil { - sessVars.LastQueryInfo.ErrMsg = err.Error() - } - }() - } - } - // Save origTxnCtx here to avoid it reset in the transaction retry. origTxnCtx := sessVars.TxnCtx err = se.checkTxnAborted(s) diff --git a/pkg/session/test/variable/BUILD.bazel b/pkg/session/test/variable/BUILD.bazel index 3f0a3dabf64f4..380a6dff7e4b8 100644 --- a/pkg/session/test/variable/BUILD.bazel +++ b/pkg/session/test/variable/BUILD.bazel @@ -8,7 +8,7 @@ go_test( "variable_test.go", ], flaky = True, - shard_count = 9, + shard_count = 10, deps = [ "//pkg/config", "//pkg/kv", diff --git a/pkg/session/test/variable/variable_test.go b/pkg/session/test/variable/variable_test.go index 776489678396c..84354537e445e 100644 --- a/pkg/session/test/variable/variable_test.go +++ b/pkg/session/test/variable/variable_test.go @@ -17,6 +17,7 @@ package variable import ( "context" "fmt" + "strings" "testing" "github.com/pingcap/failpoint" @@ -362,3 +363,24 @@ func TestIsolationRead(t *testing.T) { require.True(t, hasTiFlash) require.False(t, hasTiKV) } + +func TestLastQueryInfo(t *testing.T) { + require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/executor/mockRUConsumption", `return()`)) + defer func() { + require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/executor/mockRUConsumption")) + }() + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test") + tk.MustExec("drop table if exists t") + tk.MustExec("create table t(a int, b int, index idx(a))") + tk.MustExec(`prepare stmt1 from 'select * from t'`) + tk.MustExec("execute stmt1") + checkMatch := func(actual []string, expected []interface{}) bool { + return strings.Contains(actual[0], expected[0].(string)) + } + tk.MustQuery("select @@tidb_last_query_info;").CheckWithFunc(testkit.Rows(`"last_ru_consumption":15`), checkMatch) + tk.MustExec("select a from t where a = 1") + tk.MustQuery("select @@tidb_last_query_info;").CheckWithFunc(testkit.Rows(`"last_ru_consumption":27`), checkMatch) + tk.MustQuery("select @@tidb_last_query_info;").CheckWithFunc(testkit.Rows(`"last_ru_consumption":30`), checkMatch) +} diff --git a/pkg/sessionctx/sessionstates/session_states.go b/pkg/sessionctx/sessionstates/session_states.go index 833c0c3b9c09d..325fd1a520b79 100644 --- a/pkg/sessionctx/sessionstates/session_states.go +++ b/pkg/sessionctx/sessionstates/session_states.go @@ -49,10 +49,11 @@ type PreparedStmtInfo struct { // QueryInfo represents the information of last executed query. It's used to expose information for test purpose. type QueryInfo struct { - TxnScope string `json:"txn_scope"` - StartTS uint64 `json:"start_ts"` - ForUpdateTS uint64 `json:"for_update_ts"` - ErrMsg string `json:"error,omitempty"` + TxnScope string `json:"txn_scope"` + StartTS uint64 `json:"start_ts"` + ForUpdateTS uint64 `json:"for_update_ts"` + LastRUConsumption float64 `json:"last_ru_consumption"` + ErrMsg string `json:"error,omitempty"` } // LastDDLInfo represents the information of last DDL. It's used to expose information for test purpose.