From 608433358137277f4bccde1dbad797ecf3007b73 Mon Sep 17 00:00:00 2001 From: yibin87 Date: Mon, 16 Feb 2026 09:16:47 +0800 Subject: [PATCH 1/5] Cherry pick topsql adding network field to feature branch Signed-off-by: yibin87 --- pkg/executor/adapter.go | 6 +- pkg/server/conn.go | 11 +- pkg/sessionctx/variable/session.go | 6 + pkg/util/topsql/reporter/BUILD.bazel | 2 +- pkg/util/topsql/reporter/datamodel.go | 118 +++++++----- pkg/util/topsql/reporter/reporter.go | 51 +++++- pkg/util/topsql/reporter/reporter_test.go | 183 +++++++++++++++++++ pkg/util/topsql/stmtstats/BUILD.bazel | 2 +- pkg/util/topsql/stmtstats/aggregator_test.go | 4 +- pkg/util/topsql/stmtstats/stmtstats.go | 16 +- pkg/util/topsql/stmtstats/stmtstats_test.go | 65 +++++-- 11 files changed, 387 insertions(+), 77 deletions(-) diff --git a/pkg/executor/adapter.go b/pkg/executor/adapter.go index daff47639df30..d01359085d43a 100644 --- a/pkg/executor/adapter.go +++ b/pkg/executor/adapter.go @@ -2191,13 +2191,13 @@ func (a *ExecStmt) observeStmtBeginForTopSQL(ctx context.Context) context.Contex } // Always attach the SQL and plan info uses to catch the running SQL when Top SQL is enabled in execution. if stats != nil { - stats.OnExecutionBegin(sqlDigestByte, planDigestByte) + stats.OnExecutionBegin(sqlDigestByte, planDigestByte, vars.InPacketBytes.Load()) } return topsql.AttachSQLAndPlanInfo(ctx, sqlDigest, planDigest) } if stats != nil { - stats.OnExecutionBegin(sqlDigestByte, planDigestByte) + stats.OnExecutionBegin(sqlDigestByte, planDigestByte, vars.InPacketBytes.Load()) // This is a special logic prepared for TiKV's SQLExecCount. sc.KvExecCounter = stats.CreateKvExecCounter(sqlDigestByte, planDigestByte) } @@ -2222,7 +2222,7 @@ func (a *ExecStmt) observeStmtFinishedForTopSQL() { if stats := a.Ctx.GetStmtStats(); stats != nil && topsqlstate.TopSQLEnabled() { sqlDigest, planDigest := a.getSQLPlanDigest() execDuration := vars.GetTotalCostDuration() - stats.OnExecutionFinished(sqlDigest, planDigest, execDuration) + stats.OnExecutionFinished(sqlDigest, planDigest, execDuration, vars.OutPacketBytes.Load()) } } diff --git a/pkg/server/conn.go b/pkg/server/conn.go index 2fcc89982d60b..3d02d4ee3f7ac 100644 --- a/pkg/server/conn.go +++ b/pkg/server/conn.go @@ -486,7 +486,11 @@ func (cc *clientConn) readPacket() ([]byte, error) { if cc.getCtx() != nil { cc.pkt.SetMaxAllowedPacket(cc.ctx.GetSessionVars().MaxAllowedPacket) } - return cc.pkt.ReadPacket() + data, err := cc.pkt.ReadPacket() + if err == nil && cc.getCtx() != nil { + cc.ctx.GetSessionVars().InPacketBytes.Add(uint64(len(data))) + } + return data, err } func (cc *clientConn) writePacket(data []byte) error { @@ -495,6 +499,9 @@ func (cc *clientConn) writePacket(data []byte) error { failpoint.Return(nil) } }) + if cc.getCtx() != nil { + cc.ctx.GetSessionVars().OutPacketBytes.Add(uint64(len(data))) + } return cc.pkt.WritePacket(data) } @@ -1275,6 +1282,8 @@ func (cc *clientConn) dispatch(ctx context.Context, data []byte) error { defer func() { // reset killed for each request cc.ctx.GetSessionVars().SQLKiller.Reset() + cc.ctx.GetSessionVars().InPacketBytes.Store(0) + cc.ctx.GetSessionVars().OutPacketBytes.Store(0) }() t := time.Now() if (cc.ctx.Status() & mysql.ServerStatusInTrans) > 0 { diff --git a/pkg/sessionctx/variable/session.go b/pkg/sessionctx/variable/session.go index eb7c95fe87d41..73bc8f41ff9a8 100644 --- a/pkg/sessionctx/variable/session.go +++ b/pkg/sessionctx/variable/session.go @@ -1720,6 +1720,12 @@ type SessionVars struct { // InternalSQLScanUserTable indicates whether to use user table for internal SQL. it will be used by TTL scan InternalSQLScanUserTable bool + + // InPacketBytes records the total incoming packet bytes from clients for current session. + InPacketBytes atomic.Uint64 + + // OutPacketBytes records the total outcoming packet bytes to clients for current session. + OutPacketBytes atomic.Uint64 } // GetSessionVars implements the `SessionVarsProvider` interface. diff --git a/pkg/util/topsql/reporter/BUILD.bazel b/pkg/util/topsql/reporter/BUILD.bazel index 7f612bb2a412d..8b9301f22da7f 100644 --- a/pkg/util/topsql/reporter/BUILD.bazel +++ b/pkg/util/topsql/reporter/BUILD.bazel @@ -45,7 +45,7 @@ go_test( ], embed = [":reporter"], flaky = True, - shard_count = 36, + shard_count = 37, deps = [ "//pkg/config", "//pkg/testkit/testsetup", diff --git a/pkg/util/topsql/reporter/datamodel.go b/pkg/util/topsql/reporter/datamodel.go index 9aef688198fc4..81c98f99ec68f 100644 --- a/pkg/util/topsql/reporter/datamodel.go +++ b/pkg/util/topsql/reporter/datamodel.go @@ -86,12 +86,14 @@ func zeroTsItem() tsItem { // toProto converts the tsItem to the corresponding protobuf representation. func (i *tsItem) toProto() *tipb.TopSQLRecordItem { return &tipb.TopSQLRecordItem{ - TimestampSec: i.timestamp, - CpuTimeMs: i.cpuTimeMs, - StmtExecCount: i.stmtStats.ExecCount, - StmtKvExecCount: i.stmtStats.KvStatsItem.KvExecCount, - StmtDurationSumNs: i.stmtStats.SumDurationNs, - StmtDurationCount: i.stmtStats.DurationCount, + TimestampSec: i.timestamp, + CpuTimeMs: i.cpuTimeMs, + StmtExecCount: i.stmtStats.ExecCount, + StmtKvExecCount: i.stmtStats.KvStatsItem.KvExecCount, + StmtDurationSumNs: i.stmtStats.SumDurationNs, + StmtDurationCount: i.stmtStats.DurationCount, + StmtNetworkInBytes: i.stmtStats.NetworkInBytes, + StmtNetworkOutBytes: i.stmtStats.NetworkOutBytes, // Convert more indicators here. } } @@ -200,20 +202,25 @@ func (r *record) appendCPUTime(timestamp uint64, cpuTimeMs uint32) { // Before: // tsIndex: [10000 => 0] // tsItems: - // timestamp: [10000] - // cpuTimeMs: [0] - // stmtStats.ExecCount: [?] - // stmtStats.KvExecCount: [map{"?": ?}] - // stmtStats.DurationSum: [?] + // timestamp: [10000] + // cpuTimeMs: [0] + // stmtStats.ExecCount: [?] + // stmtStats.KvExecCount: [map{"?": ?}] + // stmtStats.DurationSum: [?] + // stmtStats.NetworkInBytes: [?] + // stmtStats.NetworkOutBytes: [?] // // After: // tsIndex: [10000 => 0] // tsItems: - // timestamp: [10000] - // cpuTimeMs: [123] - // stmtStats.ExecCount: [?] - // stmtStats.KvExecCount: [map{"?": ?}] - // stmtStats.DurationSum: [?] + // timestamp: [10000] + // cpuTimeMs: [123] + // stmtStats.ExecCount: [?] + // stmtStats.KvExecCount: [map{"?": ?}] + // stmtStats.DurationSum: [?] + // stmtStats.DurationSum: [?] + // stmtStats.NetworkInBytes: [?] + // stmtStats.NetworkOutBytes: [?] // r.tsItems[index].cpuTimeMs += cpuTimeMs } else { @@ -225,20 +232,24 @@ func (r *record) appendCPUTime(timestamp uint64, cpuTimeMs uint32) { // Before: // tsIndex: [] // tsItems: - // timestamp: [] - // cpuTimeMs: [] - // stmtStats.ExecCount: [] - // stmtStats.KvExecCount: [] - // stmtStats.DurationSum: [] + // timestamp: [] + // cpuTimeMs: [] + // stmtStats.ExecCount: [] + // stmtStats.KvExecCount: [] + // stmtStats.DurationSum: [] + // stmtStats.NetworkInBytes: [] + // stmtStats.NetworkOutBytes: [] // // After: // tsIndex: [10000 => 0] // tsItems: - // timestamp: [10000] - // cpuTimeMs: [123] - // stmtStats.ExecCount: [0] - // stmtStats.KvExecCount: [map{}] - // stmtStats.DurationSum: [0] + // timestamp: [10000] + // cpuTimeMs: [123] + // stmtStats.ExecCount: [0] + // stmtStats.KvExecCount: [map{}] + // stmtStats.DurationSum: [0] + // stmtStats.NetworkInBytes: [0] + // stmtStats.NetworkOutBytes: [0] // newItem := zeroTsItem() newItem.timestamp = timestamp @@ -258,25 +269,29 @@ func (r *record) appendStmtStatsItem(timestamp uint64, item stmtstats.StatementS // corresponding stmtStats has been set to 0 (or other values, // although impossible), so we merge it. // - // let timestamp = 10000, execCount = 123, kvExecCount = map{"1.1.1.1:1": 123}, durationSum = 456 - // + // let timestamp = 10000, execCount = 123, kvExecCount = map{"1.1.1.1:1": 123}, durationSum = 456, + // networkInBytes = 10, networkOutBytes = 20 // Before: // tsIndex: [10000 => 0] // tsItems: - // timestamp: [10000] - // cpuTimeMs: [?] - // stmtStats.ExecCount: [0] - // stmtStats.KvExecCount: [map{}] - // stmtStats.DurationSum: [0] + // timestamp: [10000] + // cpuTimeMs: [?] + // stmtStats.ExecCount: [0] + // stmtStats.KvExecCount: [map{}] + // stmtStats.DurationSum: [0] + // stmtStats.NetworkInBytes: [0] + // stmtStats.NetworkOutBytes: [0] // // After: // tsIndex: [10000 => 0] // tsItems: - // timestamp: [10000] - // cpuTimeMs: [?] - // stmtStats.ExecCount: [123] - // stmtStats.KvExecCount: [map{"1.1.1.1:1": 123}] - // stmtStats.DurationSum: [456] + // timestamp: [10000] + // cpuTimeMs: [?] + // stmtStats.ExecCount: [123] + // stmtStats.KvExecCount: [map{"1.1.1.1:1": 123}] + // stmtStats.DurationSum: [456] + // stmtStats.NetworkInBytes: [10] + // stmtStats.NetworkOutBytes: [20] // r.tsItems[index].stmtStats.Merge(&item) } else { @@ -284,24 +299,29 @@ func (r *record) appendStmtStatsItem(timestamp uint64, item stmtstats.StatementS // Other fields in tsItem except stmtStats will be initialized to 0. // // let timestamp = 10000, execCount = 123, kvExecCount = map{"1.1.1.1:1": 123}, durationSum = 456 + // networkInBytes = 10, networkOutBytes = 20 // // Before: // tsIndex: [] // tsItems: - // timestamp: [] - // cpuTimeMs: [] - // stmtStats.ExecCount: [] - // stmtStats.KvExecCount: [] - // stmtStats.DurationSum: [] + // timestamp: [] + // cpuTimeMs: [] + // stmtStats.ExecCount: [] + // stmtStats.KvExecCount: [] + // stmtStats.DurationSum: [] + // stmtStats.NetworkInBytes: [] + // stmtStats.NetworkOutBytes: [] // // After: // tsIndex: [10000 => 0] // tsItems: - // timestamp: [10000] - // cpuTimeMs: [0] - // stmtStats.ExecCount: [123] - // stmtStats.KvExecCount: [map{"1.1.1.1:1": 123}] - // stmtStats.DurationSum: [456] + // timestamp: [10000] + // cpuTimeMs: [0] + // stmtStats.ExecCount: [123] + // stmtStats.KvExecCount: [map{"1.1.1.1:1": 123}] + // stmtStats.DurationSum: [456] + // stmtStats.NetworkInBytes: [10] + // stmtStats.NetworkOutBytes: [20] // newItem := zeroTsItem() newItem.timestamp = timestamp @@ -543,7 +563,7 @@ func (c *collecting) getReportRecords() records { for _, v := range c.records { rs = append(rs, *v) } - if others != nil && others.totalCPUTimeMs > 0 { + if others != nil { rs = append(rs, *others) } return rs diff --git a/pkg/util/topsql/reporter/reporter.go b/pkg/util/topsql/reporter/reporter.go index 48d25b7495c3b..f9512dfb58752 100644 --- a/pkg/util/topsql/reporter/reporter.go +++ b/pkg/util/topsql/reporter/reporter.go @@ -25,6 +25,7 @@ import ( reporter_metrics "github.com/pingcap/tidb/pkg/util/topsql/reporter/metrics" topsqlstate "github.com/pingcap/tidb/pkg/util/topsql/state" "github.com/pingcap/tidb/pkg/util/topsql/stmtstats" + "github.com/wangjohn/quickselect" "go.uber.org/zap" ) @@ -233,20 +234,62 @@ func (tsr *RemoteTopSQLReporter) processCPUTimeData(timestamp uint64, data cpuRe func (tsr *RemoteTopSQLReporter) processStmtStatsData() { defer util.Recover("top-sql", "processStmtStatsData", nil, false) + maxLen := 0 + for _, data := range tsr.stmtStatsBuffer { + maxLen = max(maxLen, len(data)) + } + u64Slice := make([]uint64, 0, maxLen) + k := int(topsqlstate.GlobalState.MaxStatementCount.Load()) for timestamp, data := range tsr.stmtStatsBuffer { + kthNetworkBytes := findKthNetworkBytes(data, k, u64Slice) for digest, item := range data { sqlDigest, planDigest := []byte(digest.SQLDigest), []byte(digest.PlanDigest) - if tsr.collecting.hasEvicted(timestamp, sqlDigest, planDigest) { - // This timestamp+sql+plan has been evicted due to low CPUTime. + // Note, by filtering with the kthNetworkBytes, we get fewer than N records(N - 1 records, at most time). The actual picked records + // count is decided by the count of duplicated kthNetworkBytes，if kthNetworkBytes is unique, + // For performance reason, do not convert the whole map into a slice and pick exactly topN records. + if item.NetworkInBytes+item.NetworkOutBytes > kthNetworkBytes || !tsr.collecting.hasEvicted(timestamp, sqlDigest, planDigest) { + tsr.collecting.getOrCreateRecord(sqlDigest, planDigest).appendStmtStatsItem(timestamp, *item) + } else { tsr.collecting.appendOthersStmtStatsItem(timestamp, *item) - continue } - tsr.collecting.getOrCreateRecord(sqlDigest, planDigest).appendStmtStatsItem(timestamp, *item) } } tsr.stmtStatsBuffer = map[uint64]stmtstats.StatementStatsMap{} } +// The uint64Slice type attaches the QuickSelect interface to an array of uint64s. It +// implements Interface so that you can call QuickSelect(k) on any IntSlice. +type uint64Slice []uint64 + +func (t uint64Slice) Len() int { + return len(t) +} + +func (t uint64Slice) Less(i, j int) bool { + return t[i] > t[j] +} + +func (t uint64Slice) Swap(i, j int) { + t[i], t[j] = t[j], t[i] +} + +// findKthNetworkBytes finds the k-th largest network bytes in data using quickselect algorithm. +func findKthNetworkBytes(data stmtstats.StatementStatsMap, k int, u64Slice []uint64) uint64 { + var kthNetworkBytes uint64 + if len(data) > k { + u64Slice = u64Slice[:0] + for _, item := range data { + u64Slice = append(u64Slice, item.NetworkInBytes+item.NetworkOutBytes) + } + _ = quickselect.QuickSelect(uint64Slice(u64Slice), k) + kthNetworkBytes = u64Slice[0] + for i := range k { + kthNetworkBytes = min(kthNetworkBytes, u64Slice[i]) + } + } + return kthNetworkBytes +} + // takeDataAndSendToReportChan takes records data and then send to the report channel for reporting. func (tsr *RemoteTopSQLReporter) takeDataAndSendToReportChan() { // Send to report channel. When channel is full, data will be dropped. diff --git a/pkg/util/topsql/reporter/reporter_test.go b/pkg/util/topsql/reporter/reporter_test.go index 23708b196e3f1..797e41c05e868 100644 --- a/pkg/util/topsql/reporter/reporter_test.go +++ b/pkg/util/topsql/reporter/reporter_test.go @@ -15,6 +15,7 @@ package reporter import ( + "fmt" "sort" "strconv" "strings" @@ -520,6 +521,188 @@ func TestReporterWorker(t *testing.T) { assert.Equal(t, []byte("P1"), data.DataRecords[0].PlanDigest) } +func TestProcessStmtStatsData(t *testing.T) { + topsqlstate.GlobalState.ReportIntervalSeconds.Store(3) + topsqlstate.GlobalState.MaxStatementCount.Store(3) + + r := NewRemoteTopSQLReporter(mockPlanBinaryDecoderFunc, mockPlanBinaryCompressFunc) + r.Start() + defer r.Close() + + ch := make(chan *ReportData, 1) + ds := newMockDataSink(ch) + err := r.Register(ds) + assert.NoError(t, err) + + r.Collect(nil) + r.Collect([]collector.SQLCPUTimeRecord{ + { + SQLDigest: []byte("S1"), + PlanDigest: []byte("P1"), + CPUTimeMs: 1, + }, + { + SQLDigest: []byte("S2"), + PlanDigest: []byte("P2"), + CPUTimeMs: 2, + }, + }) + r.CollectStmtStatsMap(nil) + r.CollectStmtStatsMap(stmtstats.StatementStatsMap{ + stmtstats.SQLPlanDigest{ + SQLDigest: "S1", + PlanDigest: "P1", + }: &stmtstats.StatementStatsItem{ + ExecCount: 1, + NetworkInBytes: 1, + NetworkOutBytes: 0, + }, + stmtstats.SQLPlanDigest{ + SQLDigest: "S2", + PlanDigest: "P2", + }: &stmtstats.StatementStatsItem{ + ExecCount: 2, + NetworkInBytes: 0, + NetworkOutBytes: 2, + }, + stmtstats.SQLPlanDigest{ + SQLDigest: "S3", + PlanDigest: "P3", + }: &stmtstats.StatementStatsItem{ + ExecCount: 3, + NetworkInBytes: 1, + NetworkOutBytes: 2, + }, + }) + + var data *ReportData + select { + case data = <-ch: + case <-time.After(5 * time.Second): + require.Fail(t, "no data in ch") + } + + assert.Len(t, data.DataRecords, 3) + sort.Slice(data.DataRecords, func(i, j int) bool { + return data.DataRecords[i].Items[0].StmtNetworkInBytes+data.DataRecords[i].Items[0].StmtNetworkOutBytes < + data.DataRecords[j].Items[0].StmtNetworkInBytes+data.DataRecords[j].Items[0].StmtNetworkOutBytes + }) + // Check Si, Pi and corresponding item value i. + for i, record := range data.DataRecords { + j := i + 1 + assert.Equal(t, []byte(fmt.Sprintf("S%d", j)), record.SqlDigest) + assert.Equal(t, []byte(fmt.Sprintf("P%d", j)), record.PlanDigest) + assert.Equal(t, uint64(j), record.Items[0].StmtNetworkInBytes+record.Items[0].StmtNetworkOutBytes) + } + // S3, P3 has no recorded CPU time data, so the CpuTimeMs is 0. + assert.Equal(t, uint32(0), data.DataRecords[2].Items[0].CpuTimeMs) + + // Check cases with others and evicted + r.Collect(nil) + // S1/P1 and S4/P4 will be evicted since MaxStatementCount is 3 and its cputime is the lowest. + r.Collect([]collector.SQLCPUTimeRecord{ + { + SQLDigest: []byte("S1"), + PlanDigest: []byte("P1"), + CPUTimeMs: 1, + }, + { + SQLDigest: []byte("S2"), + PlanDigest: []byte("P2"), + CPUTimeMs: 2, + }, + { + SQLDigest: []byte("S3"), + PlanDigest: []byte("P3"), + CPUTimeMs: 3, + }, + { + SQLDigest: []byte("S4"), + PlanDigest: []byte("P4"), + CPUTimeMs: 0, + }, + { + SQLDigest: []byte("S7"), + PlanDigest: []byte("P7"), + CPUTimeMs: 7, + }, + }) + r.CollectStmtStatsMap(nil) + // S1/P1 and S7/P7 will be picked from network dimension, and added + // Also S2/P2 is not picked from network dimension, while it is not evicted in cputime data, so it won't be added to others. + // S4/P4 is not picked from network dimension, and evicted in cputime data, so it will be added to others. + r.CollectStmtStatsMap(stmtstats.StatementStatsMap{ + stmtstats.SQLPlanDigest{ + SQLDigest: "S1", + PlanDigest: "P1", + }: &stmtstats.StatementStatsItem{ + ExecCount: 10, + NetworkInBytes: 10, + NetworkOutBytes: 0, + }, + stmtstats.SQLPlanDigest{ + SQLDigest: "S2", + PlanDigest: "P2", + }: &stmtstats.StatementStatsItem{ + ExecCount: 2, + NetworkInBytes: 0, + NetworkOutBytes: 2, + }, + stmtstats.SQLPlanDigest{ + SQLDigest: "S4", + PlanDigest: "P4", + }: &stmtstats.StatementStatsItem{ + ExecCount: 4, + NetworkInBytes: 1, + NetworkOutBytes: 3, + }, + stmtstats.SQLPlanDigest{ + SQLDigest: "S6", + PlanDigest: "P6", + }: &stmtstats.StatementStatsItem{ + ExecCount: 6, + NetworkInBytes: 2, + NetworkOutBytes: 4, + }, + }) + + select { + case data = <-ch: + case <-time.After(5 * time.Second): + require.Fail(t, "no data in ch") + } + assert.Len(t, data.DataRecords, 6) + sort.Slice(data.DataRecords, func(i, j int) bool { + return (data.DataRecords[i].Items[0].StmtNetworkInBytes+data.DataRecords[i].Items[0].StmtNetworkOutBytes)*100+uint64(data.DataRecords[i].Items[0].CpuTimeMs) < + (data.DataRecords[j].Items[0].StmtNetworkInBytes+data.DataRecords[j].Items[0].StmtNetworkOutBytes)*100+uint64(data.DataRecords[j].Items[0].CpuTimeMs) + }) + // DataRecords should be: + // S3, P3, CpuTime=3, Network=0 + // S7, P7, CpuTime=7, Network=0 + // S2, P2, CpuTime=2, Network=2 + // Other, CpuTime=1, Network=4 + // S6, P6, CpuTime=0, Network=6 + // S1, P1, CpuTime=0, Network=10 + assert.Equal(t, []byte("S3"), data.DataRecords[0].SqlDigest) + assert.Equal(t, uint32(3), data.DataRecords[0].Items[0].CpuTimeMs) + assert.Equal(t, uint64(0), data.DataRecords[0].Items[0].StmtNetworkInBytes+data.DataRecords[0].Items[0].StmtNetworkOutBytes) + assert.Equal(t, []byte("S7"), data.DataRecords[1].SqlDigest) + assert.Equal(t, uint32(7), data.DataRecords[1].Items[0].CpuTimeMs) + assert.Equal(t, uint64(0), data.DataRecords[1].Items[0].StmtNetworkInBytes+data.DataRecords[0].Items[0].StmtNetworkOutBytes) + assert.Equal(t, []byte("S2"), data.DataRecords[2].SqlDigest) + assert.Equal(t, uint32(2), data.DataRecords[2].Items[0].CpuTimeMs) + assert.Equal(t, uint64(2), data.DataRecords[2].Items[0].StmtNetworkInBytes+data.DataRecords[2].Items[0].StmtNetworkOutBytes) + assert.Equal(t, []byte(nil), data.DataRecords[3].SqlDigest) + assert.Equal(t, uint32(1), data.DataRecords[3].Items[0].CpuTimeMs) + assert.Equal(t, uint64(4), data.DataRecords[3].Items[0].StmtNetworkInBytes+data.DataRecords[3].Items[0].StmtNetworkOutBytes) + assert.Equal(t, []byte("S6"), data.DataRecords[4].SqlDigest) + assert.Equal(t, uint32(0), data.DataRecords[4].Items[0].CpuTimeMs) + assert.Equal(t, uint64(6), data.DataRecords[4].Items[0].StmtNetworkInBytes+data.DataRecords[4].Items[0].StmtNetworkOutBytes) + assert.Equal(t, []byte("S1"), data.DataRecords[5].SqlDigest) + assert.Equal(t, uint32(0), data.DataRecords[5].Items[0].CpuTimeMs) + assert.Equal(t, uint64(10), data.DataRecords[5].Items[0].StmtNetworkInBytes+data.DataRecords[5].Items[0].StmtNetworkOutBytes) +} + func initializeCache(maxStatementsNum, interval int) (*RemoteTopSQLReporter, *mockDataSink2) { ts, ds := setupRemoteTopSQLReporter(maxStatementsNum, interval) populateCache(ts, 0, maxStatementsNum, 1) diff --git a/pkg/util/topsql/stmtstats/BUILD.bazel b/pkg/util/topsql/stmtstats/BUILD.bazel index ee5405204e0e7..f1132beaa2cb9 100644 --- a/pkg/util/topsql/stmtstats/BUILD.bazel +++ b/pkg/util/topsql/stmtstats/BUILD.bazel @@ -28,7 +28,7 @@ go_test( ], embed = [":stmtstats"], flaky = True, - shard_count = 11, + shard_count = 12, deps = [ "//pkg/testkit/testsetup", "//pkg/util/topsql/state", diff --git a/pkg/util/topsql/stmtstats/aggregator_test.go b/pkg/util/topsql/stmtstats/aggregator_test.go index 987deac63b329..04cdd94024321 100644 --- a/pkg/util/topsql/stmtstats/aggregator_test.go +++ b/pkg/util/topsql/stmtstats/aggregator_test.go @@ -58,8 +58,8 @@ func Test_aggregator_register_collect(t *testing.T) { finished: atomic.NewBool(false), } a.register(stats) - stats.OnExecutionBegin([]byte("SQL-1"), []byte("")) - stats.OnExecutionFinished([]byte("SQL-1"), []byte(""), time.Millisecond) + stats.OnExecutionBegin([]byte("SQL-1"), []byte(""), 0) + stats.OnExecutionFinished([]byte("SQL-1"), []byte(""), time.Millisecond, 0) total := StatementStatsMap{} a.registerCollector(newMockCollector(func(data StatementStatsMap) { total.Merge(data) diff --git a/pkg/util/topsql/stmtstats/stmtstats.go b/pkg/util/topsql/stmtstats/stmtstats.go index 1a08ace4b42a0..11297accc6a58 100644 --- a/pkg/util/topsql/stmtstats/stmtstats.go +++ b/pkg/util/topsql/stmtstats/stmtstats.go @@ -30,13 +30,13 @@ var _ StatementObserver = &StatementStats{} // corresponding locations, without paying attention to implementation details. type StatementObserver interface { // OnExecutionBegin should be called before statement execution. - OnExecutionBegin(sqlDigest, planDigest []byte) + OnExecutionBegin(sqlDigest, planDigest []byte, inNetworkBytes uint64) // OnExecutionFinished should be called after the statement is executed. // WARNING: Currently Only call StatementObserver API when TopSQL is enabled, // there is no guarantee that both OnExecutionBegin and OnExecutionFinished will be called for a SQL, // such as TopSQL is enabled during a SQL execution. - OnExecutionFinished(sqlDigest, planDigest []byte, execDuration time.Duration) + OnExecutionFinished(sqlDigest, planDigest []byte, execDuration time.Duration, outNetworkBytes uint64) } // StatementStats is a counter used locally in each session. @@ -60,17 +60,18 @@ func CreateStatementStats() *StatementStats { } // OnExecutionBegin implements StatementObserver.OnExecutionBegin. -func (s *StatementStats) OnExecutionBegin(sqlDigest, planDigest []byte) { +func (s *StatementStats) OnExecutionBegin(sqlDigest, planDigest []byte, inNetworkBytes uint64) { s.mu.Lock() defer s.mu.Unlock() item := s.GetOrCreateStatementStatsItem(sqlDigest, planDigest) item.ExecCount++ + item.NetworkInBytes += inNetworkBytes // Count more data here. } // OnExecutionFinished implements StatementObserver.OnExecutionFinished. -func (s *StatementStats) OnExecutionFinished(sqlDigest, planDigest []byte, execDuration time.Duration) { +func (s *StatementStats) OnExecutionFinished(sqlDigest, planDigest []byte, execDuration time.Duration, outNetworkBytes uint64) { ns := execDuration.Nanoseconds() if ns < 0 { return @@ -82,6 +83,7 @@ func (s *StatementStats) OnExecutionFinished(sqlDigest, planDigest []byte, execD item.SumDurationNs += uint64(ns) item.DurationCount++ + item.NetworkOutBytes += outNetworkBytes // Count more data here. } @@ -182,6 +184,10 @@ type StatementStatsItem struct { // DurationCount represents the number of SQL executions specially // used to calculate SQLDuration. DurationCount uint64 + // NetworkInBytes represents the total number of network input bytes from client. + NetworkInBytes uint64 + // NetworkOutBytes represents the total number of network input bytes to client. + NetworkOutBytes uint64 } // NewStatementStatsItem creates an empty StatementStatsItem. @@ -205,6 +211,8 @@ func (i *StatementStatsItem) Merge(other *StatementStatsItem) { i.ExecCount += other.ExecCount i.SumDurationNs += other.SumDurationNs i.DurationCount += other.DurationCount + i.NetworkInBytes += other.NetworkInBytes + i.NetworkOutBytes += other.NetworkOutBytes i.KvStatsItem.Merge(other.KvStatsItem) } diff --git a/pkg/util/topsql/stmtstats/stmtstats_test.go b/pkg/util/topsql/stmtstats/stmtstats_test.go index 7e3dd377d8a0a..508490f18b91a 100644 --- a/pkg/util/topsql/stmtstats/stmtstats_test.go +++ b/pkg/util/topsql/stmtstats/stmtstats_test.go @@ -85,18 +85,24 @@ func TestKvStatementStatsItem_Merge(t *testing.T) { func TestStatementsStatsItem_Merge(t *testing.T) { item1 := &StatementStatsItem{ - ExecCount: 1, - SumDurationNs: 100, - KvStatsItem: NewKvStatementStatsItem(), + ExecCount: 1, + SumDurationNs: 100, + KvStatsItem: NewKvStatementStatsItem(), + NetworkInBytes: 10, + NetworkOutBytes: 20, } item2 := &StatementStatsItem{ - ExecCount: 2, - SumDurationNs: 50, - KvStatsItem: NewKvStatementStatsItem(), + ExecCount: 2, + SumDurationNs: 50, + KvStatsItem: NewKvStatementStatsItem(), + NetworkInBytes: 50, + NetworkOutBytes: 60, } item1.Merge(item2) assert.Equal(t, uint64(3), item1.ExecCount) assert.Equal(t, uint64(150), item1.SumDurationNs) + assert.Equal(t, uint64(60), item1.NetworkInBytes) + assert.Equal(t, uint64(80), item1.NetworkOutBytes) } func TestStatementStatsMap_Merge(t *testing.T) { @@ -180,17 +186,17 @@ func TestExecCounter_AddExecCount_Take(t *testing.T) { m := stats.Take() assert.Len(t, m, 0) for range 1 { - stats.OnExecutionBegin([]byte("SQL-1"), []byte("")) + stats.OnExecutionBegin([]byte("SQL-1"), []byte(""), 0) } for range 2 { - stats.OnExecutionBegin([]byte("SQL-2"), []byte("")) - stats.OnExecutionFinished([]byte("SQL-2"), []byte(""), time.Second) + stats.OnExecutionBegin([]byte("SQL-2"), []byte(""), 0) + stats.OnExecutionFinished([]byte("SQL-2"), []byte(""), time.Second, 0) } for range 3 { - stats.OnExecutionBegin([]byte("SQL-3"), []byte("")) - stats.OnExecutionFinished([]byte("SQL-3"), []byte(""), time.Millisecond) + stats.OnExecutionBegin([]byte("SQL-3"), []byte(""), 0) + stats.OnExecutionFinished([]byte("SQL-3"), []byte(""), time.Millisecond, 0) } - stats.OnExecutionFinished([]byte("SQL-3"), []byte(""), -time.Millisecond) + stats.OnExecutionFinished([]byte("SQL-3"), []byte(""), -time.Millisecond, 0) m = stats.Take() assert.Len(t, m, 3) assert.Equal(t, uint64(1), m[SQLPlanDigest{SQLDigest: "SQL-1"}].ExecCount) @@ -202,3 +208,38 @@ func TestExecCounter_AddExecCount_Take(t *testing.T) { m = stats.Take() assert.Len(t, m, 0) } + +func TestNetworkBytesAccumulation(t *testing.T) { + stats := CreateStatementStats() + sqlDigest := []byte("SQL-1") + planDigest := []byte("PLAN-1") + + // Test NetworkInBytes accumulation in OnExecutionBegin + // Call OnExecutionBegin multiple times with different network input bytes + stats.OnExecutionBegin(sqlDigest, planDigest, 100) + stats.OnExecutionBegin(sqlDigest, planDigest, 200) + stats.OnExecutionBegin(sqlDigest, planDigest, 300) + + m := stats.Take() + assert.Len(t, m, 1) + key := SQLPlanDigest{SQLDigest: BinaryDigest(sqlDigest), PlanDigest: BinaryDigest(planDigest)} + item := m[key] + assert.NotNil(t, item) + // NetworkInBytes should be accumulated: 100 + 200 + 300 = 600 + assert.Equal(t, uint64(600), item.NetworkInBytes) + assert.Equal(t, uint64(3), item.ExecCount) + + // Test NetworkOutBytes accumulation in OnExecutionFinished + // Call OnExecutionFinished multiple times with different network output bytes + stats.OnExecutionFinished(sqlDigest, planDigest, time.Second, 50) + stats.OnExecutionFinished(sqlDigest, planDigest, time.Second, 150) + stats.OnExecutionFinished(sqlDigest, planDigest, time.Second, 250) + + m = stats.Take() + assert.Len(t, m, 1) + item = m[key] + assert.NotNil(t, item) + // NetworkOutBytes should be accumulated: 50 + 150 + 250 = 450 + assert.Equal(t, uint64(450), item.NetworkOutBytes) + assert.Equal(t, uint64(3), item.DurationCount) +} From 736bfe67c375a69f8fac8fb2a5e2b82f0066058e Mon Sep 17 00:00:00 2001 From: yibin87 Date: Mon, 2 Mar 2026 10:37:25 +0800 Subject: [PATCH 2/5] pkg/server/tests: stabilize TestAuditPluginRetrying with deterministic auto-commit retries - inject unistore failpoint in auto-commit branch - protect audit result collection with a mutex to avoid concurrent append races - switch concurrent UPDATE to and close each - test-only change in , no production logic touched Signed-off-by: yibin87 --- pkg/server/tests/commontest/tidb_test.go | 43 ++++++++++++++++++------ 1 file changed, 33 insertions(+), 10 deletions(-) diff --git a/pkg/server/tests/commontest/tidb_test.go b/pkg/server/tests/commontest/tidb_test.go index 78cf20d6a0f13..8921736ec8023 100644 --- a/pkg/server/tests/commontest/tidb_test.go +++ b/pkg/server/tests/commontest/tidb_test.go @@ -3517,6 +3517,24 @@ func TestAuditPluginRetrying(t *testing.T) { retrying bool } testResults := make([]normalTest, 0) + var testResultsMu sync.Mutex + appendResult := func(audit normalTest) { + testResultsMu.Lock() + defer testResultsMu.Unlock() + testResults = append(testResults, audit) + } + resetResults := func() { + testResultsMu.Lock() + defer testResultsMu.Unlock() + testResults = testResults[:0] + } + getResults := func() []normalTest { + testResultsMu.Lock() + defer testResultsMu.Unlock() + results := make([]normalTest, len(testResults)) + copy(results, testResults) + return results + } onGeneralEvent := func(ctx context.Context, sctx *variable.SessionVars, event plugin.GeneralEvent, cmd string) { // Only consider the Completed event @@ -3529,7 +3547,7 @@ func TestAuditPluginRetrying(t *testing.T) { audit.retrying = retrying.(bool) } audit.sql = sctx.StmtCtx.OriginalSQL - testResults = append(testResults, audit) + appendResult(audit) } plugin.LoadPluginForTest(t, onGeneralEvent) defer plugin.Shutdown(context.Background()) @@ -3546,7 +3564,9 @@ func TestAuditPluginRetrying(t *testing.T) { _, err = db.Exec("INSERT INTO auto_retry_test VALUES (1, 0)") require.NoError(t, err) - testResults = testResults[:0] + resetResults() + // Inject deterministic write conflicts to make statement retries observable in auto-commit mode. + testfailpoint.Enable(t, "github.com/pingcap/tidb/pkg/store/mockstore/unistore/tikv/pessimisticLockReturnWriteConflict", "10*return(true)") // a big enough concurrency to trigger retries concurrency := 500 var wg sync.WaitGroup @@ -3554,17 +3574,19 @@ func TestAuditPluginRetrying(t *testing.T) { wg.Add(1) conn, err := db.Conn(context.Background()) require.NoError(t, err) - go func() { + go func(conn *sql.Conn) { defer wg.Done() - _, err := conn.QueryContext(context.Background(), "UPDATE auto_retry_test SET val = val + 1 WHERE id = 1") + defer conn.Close() + _, err := conn.ExecContext(context.Background(), "UPDATE auto_retry_test SET val = val + 1 WHERE id = 1") require.NoError(t, err) - }() + }(conn) } wg.Wait() - require.Greater(t, len(testResults), concurrency) + results := getResults() + require.Greater(t, len(results), concurrency) nonRetryingCount := 0 - for _, res := range testResults { + for _, res := range results { if !res.retrying { nonRetryingCount++ } @@ -3597,7 +3619,7 @@ func TestAuditPluginRetrying(t *testing.T) { return conn } - testResults = testResults[:0] + resetResults() var wg sync.WaitGroup wg.Add(2) // Transaction 1 @@ -3633,9 +3655,10 @@ func TestAuditPluginRetrying(t *testing.T) { }() wg.Wait() + results := getResults() retryingCount := 0 nonRetryingCount := 0 - for _, res := range testResults { + for _, res := range results { if res.retrying { retryingCount++ } else { @@ -3665,7 +3688,7 @@ func TestAuditPluginRetrying(t *testing.T) { ts.RunTests(t, nil, func(dbt *testkit.DBTestKit) { db := dbt.GetDB() - testResults = testResults[:0] + resetResults() runExplicitTransactionRetry(db, true) }) } From 332baf7562d748cdd5a0cc9c432b30e311c5ca2c Mon Sep 17 00:00:00 2001 From: yibin87 Date: Fri, 22 May 2026 13:56:54 +0800 Subject: [PATCH 3/5] deps: bump tipb for TopSQL network fields Signed-off-by: yibin87 --- go.mod | 2 +- go.sum | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/go.mod b/go.mod index 24ae581b82439..02b102f0f7709 100644 --- a/go.mod +++ b/go.mod @@ -90,7 +90,7 @@ require ( github.com/pingcap/log v1.1.1-0.20250917021125-19901e015dc9 github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5 github.com/pingcap/tidb/pkg/parser v0.0.0-20211011031125-9b13dc409c5e - github.com/pingcap/tipb v0.0.0-20241022082558-0607513e7fa4 + github.com/pingcap/tipb v0.0.0-20260522023117-6ceadabb61de github.com/prometheus/client_golang v1.20.5 github.com/prometheus/client_model v0.6.1 github.com/prometheus/common v0.57.0 diff --git a/go.sum b/go.sum index 5e6100010600e..a9b7f711afb56 100644 --- a/go.sum +++ b/go.sum @@ -682,6 +682,8 @@ github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5 h1:T4pXRhBflzDeA github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5/go.mod h1:rlimy0GcTvjiJqvD5mXTRr8O2eNZPBrcUgiWVYp9530= github.com/pingcap/tipb v0.0.0-20241022082558-0607513e7fa4 h1:wvaUybJT0fUReCDcFtV3CEvMuI9iu+G7IW72tbSlil4= github.com/pingcap/tipb v0.0.0-20241022082558-0607513e7fa4/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs= +github.com/pingcap/tipb v0.0.0-20260522023117-6ceadabb61de h1:zZNGJcObLgFvuOF9CPviXomIN9QGAEDj9w9SUGp3s64= +github.com/pingcap/tipb v0.0.0-20260522023117-6ceadabb61de/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs= github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ= github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU= github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= From 5502135fa396b2f0ce4ec7e30f0220f287120b93 Mon Sep 17 00:00:00 2001 From: yibin87 Date: Fri, 22 May 2026 16:32:00 +0800 Subject: [PATCH 4/5] update bazel Signed-off-by: yibin87 --- DEPS.bzl | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/DEPS.bzl b/DEPS.bzl index ab65b4644f85e..afccfe04c7922 100644 --- a/DEPS.bzl +++ b/DEPS.bzl @@ -5854,13 +5854,13 @@ def go_deps(): name = "com_github_pingcap_tipb", build_file_proto_mode = "disable_global", importpath = "github.com/pingcap/tipb", - sha256 = "1b707429b5b938a05b250b5770be2a6aa243d6a4983d23b01bbca164e86b3e3c", - strip_prefix = "github.com/pingcap/tipb@v0.0.0-20241022082558-0607513e7fa4", + sha256 = "a1f1f8523c7a8f788e73fed28c2accee002f0c5c502fdbed257eb6a42810e784", + strip_prefix = "github.com/pingcap/tipb@v0.0.0-20260522023117-6ceadabb61de", urls = [ - "http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20241022082558-0607513e7fa4.zip", - "http://ats.apps.svc/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20241022082558-0607513e7fa4.zip", - "https://cache.hawkingrei.com/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20241022082558-0607513e7fa4.zip", - "https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20241022082558-0607513e7fa4.zip", + "http://bazel-cache.pingcap.net:8080/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20260522023117-6ceadabb61de.zip", + "http://ats.apps.svc/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20260522023117-6ceadabb61de.zip", + "https://cache.hawkingrei.com/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20260522023117-6ceadabb61de.zip", + "https://storage.googleapis.com/pingcapmirror/gomod/github.com/pingcap/tipb/com_github_pingcap_tipb-v0.0.0-20260522023117-6ceadabb61de.zip", ], ) go_repository( From d4e13ec4591d7491e6d7af8ee62613ba6e9a2733 Mon Sep 17 00:00:00 2001 From: yibin87 Date: Fri, 22 May 2026 17:08:04 +0800 Subject: [PATCH 5/5] Update go sum Signed-off-by: yibin87 --- go.sum | 2 -- 1 file changed, 2 deletions(-) diff --git a/go.sum b/go.sum index a9b7f711afb56..411a86854896e 100644 --- a/go.sum +++ b/go.sum @@ -680,8 +680,6 @@ github.com/pingcap/log v1.1.1-0.20250917021125-19901e015dc9 h1:qG9BSvlWFEE5otQGa github.com/pingcap/log v1.1.1-0.20250917021125-19901e015dc9/go.mod h1:ORfBOFp1eteu2odzsyaxI+b8TzJwgjwyQcGhI+9SfEA= github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5 h1:T4pXRhBflzDeAhmOQHNPRRogMYxP13V7BkYw3ZsoSfE= github.com/pingcap/sysutil v1.0.1-0.20240311050922-ae81ee01f3a5/go.mod h1:rlimy0GcTvjiJqvD5mXTRr8O2eNZPBrcUgiWVYp9530= -github.com/pingcap/tipb v0.0.0-20241022082558-0607513e7fa4 h1:wvaUybJT0fUReCDcFtV3CEvMuI9iu+G7IW72tbSlil4= -github.com/pingcap/tipb v0.0.0-20241022082558-0607513e7fa4/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs= github.com/pingcap/tipb v0.0.0-20260522023117-6ceadabb61de h1:zZNGJcObLgFvuOF9CPviXomIN9QGAEDj9w9SUGp3s64= github.com/pingcap/tipb v0.0.0-20260522023117-6ceadabb61de/go.mod h1:A7mrd7WHBl1o63LE2bIBGEJMTNWXqhgmYiOvMLxozfs= github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=