This is an automated cherry-pick of #63796

Signed-off-by: ti-chi-bot <[email protected]>

Author: yibin87

pkg/executor/adapter.go

@@ -2193,13 +2193,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)
 	}
@@ -2224,7 +2224,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())
 	}
 }
 

pkg/server/conn.go

@@ -487,7 +487,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 {
@@ -496,6 +500,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)
 }
 
@@ -1279,6 +1286,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 {

pkg/sessionctx/variable/session.go

@@ -1725,8 +1725,51 @@ type SessionVars struct {
 	// InternalSQLScanUserTable indicates whether to use user table for internal SQL. it will be used by TTL scan
 	InternalSQLScanUserTable bool
 
+<<<<<<< HEAD
 	// IndexLookUpPushDownPolicy indicates the policy of index look up push down.
 	IndexLookUpPushDownPolicy string
+=======
+	// MemArbitrator represents the properties to be controlled by the memory arbitrator.
+	MemArbitrator struct {
+		WaitAverse    MemArbitratorWaitAverseMode
+		QueryReserved int64
+	}
+
+	// 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
+}
+
+// ResetRelevantOptVarsAndFixes resets the relevant optimizer variables and fixes.
+func (s *SessionVars) ResetRelevantOptVarsAndFixes(record bool) {
+	s.RecordRelevantOptVarsAndFixes = record
+	s.RelevantOptVars = nil
+	s.RelevantOptFixes = nil
+}
+
+// RecordRelevantOptVar records the optimizer variable that is relevant to the current query.
+func (s *SessionVars) RecordRelevantOptVar(varName string) {
+	if !s.RecordRelevantOptVarsAndFixes {
+		return
+	}
+	if s.RelevantOptVars == nil {
+		s.RelevantOptVars = make(map[string]struct{})
+	}
+	s.RelevantOptVars[varName] = struct{}{}
+}
+
+// RecordRelevantOptFix records the optimizer fix that is relevant to the current query.
+func (s *SessionVars) RecordRelevantOptFix(fixID uint64) {
+	if !s.RecordRelevantOptVarsAndFixes {
+		return
+	}
+	if s.RelevantOptFixes == nil {
+		s.RelevantOptFixes = make(map[uint64]struct{})
+	}
+	s.RelevantOptFixes[fixID] = struct{}{}
+>>>>>>> e407943c073 (*: Collect client network info for topsql to pick topN sql with highest network traffic (#63796))
 }
 
 // GetSessionVars implements the `SessionVarsProvider` interface.

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",

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,50 +269,59 @@ 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 {
 		// For this timestamp, we have not appended any tsItem, so append it directly.
 		// 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

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.