scheduler: use store leaders sum instead of hot leaders sum

pkg/mock/mockcluster/mockcluster.go

@@ -812,6 +812,6 @@ func (mc *Cluster) CheckRegionLeaderRead(region *core.RegionInfo) []*statistics.
 
 // ObserveRegionsStats records the current stores stats from region stats.
 func (mc *Cluster) ObserveRegionsStats() {
-	storeIDs, writeBytesRates, writeKeysRates := mc.BasicCluster.GetStoresWriteRate()
-	mc.HotStat.ObserveRegionsStats(storeIDs, writeBytesRates, writeKeysRates)
+	storesWriteRate := mc.BasicCluster.GetStoresWriteRate()
+	mc.HotStat.ObserveRegionsStats(storesWriteRate)
 }

server/cluster/cluster.go

@@ -340,10 +340,10 @@ func (c *RaftCluster) runStatsBackgroundJobs() {
 			return
 		case <-ticker.C:
 			c.RLock()
-			storeIDs, writeBytesRates, writeKeysRates := c.core.GetStoresWriteRate()
+			storesWriteRate := c.core.GetStoresWriteRate()
 			c.RUnlock()
 			c.Lock()
-			c.hotStat.ObserveRegionsStats(storeIDs, writeBytesRates, writeKeysRates)
+			c.hotStat.ObserveRegionsStats(storesWriteRate)
 			c.Unlock()
 		}
 	}

server/core/basic_cluster.go

@@ -283,33 +283,15 @@ func (bc *BasicCluster) GetAverageRegionSize() int64 {
 	return bc.Regions.GetAverageRegionSize()
 }
 
-func (bc *BasicCluster) getWriteRate(
-	f func(storeID uint64) (bytesRate, keysRate float64),
-) (storeIDs []uint64, bytesRates, keysRates []float64) {
+// GetStoresWriteRate get total write rate of each store's regions.
+func (bc *BasicCluster) GetStoresWriteRate() []StoreWriteRate {
 	bc.RLock()
 	defer bc.RUnlock()
-	count := len(bc.Stores.stores)
-	storeIDs = make([]uint64, 0, count)
-	bytesRates = make([]float64, 0, count)
-	keysRates = make([]float64, 0, count)
+	storesStats := make([]StoreWriteRate, 0, len(bc.Stores.stores))
 	for _, store := range bc.Stores.stores {
-		id := store.GetID()
-		bytesRate, keysRate := f(id)
-		storeIDs = append(storeIDs, id)
-		bytesRates = append(bytesRates, bytesRate)
-		keysRates = append(keysRates, keysRate)
+		storesStats = append(storesStats, bc.Regions.GetStoreWriteRate(store.GetID()))
 	}
-	return
-}
-
-// GetStoresLeaderWriteRate get total write rate of each store's leaders.
-func (bc *BasicCluster) GetStoresLeaderWriteRate() (storeIDs []uint64, bytesRates, keysRates []float64) {
-	return bc.getWriteRate(bc.Regions.GetStoreLeaderWriteRate)
-}
-
-// GetStoresWriteRate get total write rate of each store's regions.
-func (bc *BasicCluster) GetStoresWriteRate() (storeIDs []uint64, bytesRates, keysRates []float64) {
-	return bc.getWriteRate(bc.Regions.GetStoreWriteRate)
+	return storesStats
 }
 
 // PutStore put a store.

server/core/region.go

@@ -906,23 +906,35 @@ func (r *RegionsInfo) GetStoreRegionSize(storeID uint64) int64 {
 	return r.GetStoreLeaderRegionSize(storeID) + r.GetStoreFollowerRegionSize(storeID) + r.GetStoreLearnerRegionSize(storeID)
 }
 
-// GetStoreLeaderWriteRate get total write rate of store's leaders
-func (r *RegionsInfo) GetStoreLeaderWriteRate(storeID uint64) (bytesRate, keysRate float64) {
-	return r.leaders[storeID].TotalWriteRate()
+// StoreWriteRate is the statistics of total write rate of store's regions.
+type StoreWriteRate struct {
+	StoreID         uint64
+	PeerBytesRate   float64
+	PeerKeysRate    float64
+	LeaderBytesRate float64
+	LeaderKeysRate  float64
 }
 
 // GetStoreWriteRate get total write rate of store's regions
-func (r *RegionsInfo) GetStoreWriteRate(storeID uint64) (bytesRate, keysRate float64) {
-	storeBytesRate, storeKeysRate := r.leaders[storeID].TotalWriteRate()
-	bytesRate += storeBytesRate
-	keysRate += storeKeysRate
-	storeBytesRate, storeKeysRate = r.followers[storeID].TotalWriteRate()
-	bytesRate += storeBytesRate
-	keysRate += storeKeysRate
+func (r *RegionsInfo) GetStoreWriteRate(storeID uint64) StoreWriteRate {
+	leaderBytesRate, leaderKeysRate := r.leaders[storeID].TotalWriteRate()
+	peerBytesRate, peerKeysRate := leaderBytesRate, leaderKeysRate
+
+	storeBytesRate, storeKeysRate := r.followers[storeID].TotalWriteRate()
+	peerBytesRate += storeBytesRate
+	peerKeysRate += storeKeysRate
+
 	storeBytesRate, storeKeysRate = r.learners[storeID].TotalWriteRate()
-	bytesRate += storeBytesRate
-	keysRate += storeKeysRate
-	return
+	peerBytesRate += storeBytesRate
+	peerKeysRate += storeKeysRate
+
+	return StoreWriteRate{
+		StoreID:         storeID,
+		PeerBytesRate:   peerBytesRate,
+		PeerKeysRate:    peerKeysRate,
+		LeaderBytesRate: leaderBytesRate,
+		LeaderKeysRate:  leaderKeysRate,
+	}
 }
 
 // GetMetaRegions gets a set of metapb.Region from regionMap

server/schedulers/utils.go

@@ -512,10 +512,12 @@ type storeCollector interface {
 	GetLoads(storeLoads, peerLoadSum []float64, rwTy rwType, kind core.ResourceKind) (loads []float64)
 }
 
-type tikvCollector struct{}
+type tikvCollector struct {
+	isTraceRegionFlow bool
+}
 
-func newTikvCollector() storeCollector {
-	return tikvCollector{}
+func newTikvCollector(isTraceRegionFlow bool) storeCollector {
+	return tikvCollector{isTraceRegionFlow: isTraceRegionFlow}
 }
 
 func (c tikvCollector) Engine() string {
@@ -540,9 +542,14 @@ func (c tikvCollector) GetLoads(storeLoads, peerLoadSum []float64, rwTy rwType,
 			// For write requests, Write{Bytes, Keys} is applied to all Peers at the same time,
 			// while the Leader and Follower are under different loads (usually the Leader consumes more CPU).
 			// Write{QPS} does not require such processing.
-			loads[statistics.ByteDim] = peerLoadSum[statistics.ByteDim]
-			loads[statistics.KeyDim] = peerLoadSum[statistics.KeyDim]
 			loads[statistics.QueryDim] = storeLoads[statistics.StoreWriteQuery]
+			if c.isTraceRegionFlow {
+				loads[statistics.ByteDim] = storeLoads[statistics.StoreLeadersWriteBytes]
+				loads[statistics.KeyDim] = storeLoads[statistics.StoreLeadersWriteKeys]
+			} else {
+				loads[statistics.ByteDim] = peerLoadSum[statistics.ByteDim]
+				loads[statistics.KeyDim] = peerLoadSum[statistics.KeyDim]
+			}
 		case core.RegionKind:
 			loads[statistics.ByteDim] = storeLoads[statistics.StoreWriteBytes]
 			loads[statistics.KeyDim] = storeLoads[statistics.StoreWriteKeys]
@@ -611,7 +618,7 @@ func summaryStoresLoad(
 		storesLoads,
 		storeHotPeers,
 		rwTy, kind,
-		newTikvCollector(),
+		newTikvCollector(isTraceRegionFlow),
 	)
 	tiflashLoadDetail := summaryStoresLoadByEngine(
 		storeInfos,

server/statistics/kind.go

@@ -63,6 +63,8 @@ const (
 
 	StoreRegionsWriteBytes // Same as StoreWriteBytes, but it is counted by RegionHeartbeat.
 	StoreRegionsWriteKeys  // Same as StoreWriteKeys, but it is counted by RegionHeartbeat.
+	StoreLeadersWriteBytes
+	StoreLeadersWriteKeys
 
 	StoreStatCount
 )
@@ -91,6 +93,10 @@ func (k StoreStatKind) String() string {
 		return "store_regions_write_bytes"
 	case StoreRegionsWriteKeys:
 		return "store_regions_write_keys"
+	case StoreLeadersWriteBytes:
+		return "store_leaders_write_bytes"
+	case StoreLeadersWriteKeys:
+		return "store_leaders_write_keys"
 	}
 
 	return "unknown StoreStatKind"

server/statistics/store.go

@@ -78,11 +78,11 @@ func (s *StoresStats) Observe(storeID uint64, stats *pdpb.StoreStats) {
 	rollingStoreStat.Observe(stats)
 }
 
-// ObserveRegionsStats records the current stores status from region stats.
-func (s *StoresStats) ObserveRegionsStats(storeIDs []uint64, writeBytesRates, writeKeysRates []float64) {
-	for i, storeID := range storeIDs {
-		rollingStoreStat := s.GetOrCreateRollingStoreStats(storeID)
-		rollingStoreStat.ObserveRegionsStats(writeBytesRates[i], writeKeysRates[i])
+// ObserveRegionsStats records the current stores stats from region stats.
+func (s *StoresStats) ObserveRegionsStats(storesWriteRate []core.StoreWriteRate) {
+	for _, writeRate := range storesWriteRate {
+		rollingStoreStat := s.GetOrCreateRollingStoreStats(writeRate.StoreID)
+		rollingStoreStat.ObserveRegionsStats(writeRate)
 	}
 }
 
@@ -93,10 +93,10 @@ func (s *StoresStats) Set(storeID uint64, stats *pdpb.StoreStats) {
 }
 
 // SetRegionsStats sets the store statistics from region stats (for test).
-func (s *StoresStats) SetRegionsStats(storeIDs []uint64, writeBytesRates, writeKeysRates []float64) {
-	for i, storeID := range storeIDs {
-		rollingStoreStat := s.GetOrCreateRollingStoreStats(storeID)
-		rollingStoreStat.SetRegionsStats(writeBytesRates[i], writeKeysRates[i])
+func (s *StoresStats) SetRegionsStats(storesWriteRate []core.StoreWriteRate) {
+	for _, writeRate := range storesWriteRate {
+		rollingStoreStat := s.GetOrCreateRollingStoreStats(writeRate.StoreID)
+		rollingStoreStat.SetRegionsStats(writeRate)
 	}
 }
 
@@ -158,6 +158,8 @@ func newRollingStoreStats() *RollingStoreStats {
 	// The data from regionStats is used in TiFlash, so higher tolerance is required
 	movingAvgs[StoreRegionsWriteBytes] = movingaverage.NewMedianFilter(RegionsStatsRollingWindowsSize)
 	movingAvgs[StoreRegionsWriteKeys] = movingaverage.NewMedianFilter(RegionsStatsRollingWindowsSize)
+	movingAvgs[StoreLeadersWriteBytes] = movingaverage.NewMedianFilter(RegionsStatsRollingWindowsSize)
+	movingAvgs[StoreLeadersWriteKeys] = movingaverage.NewMedianFilter(RegionsStatsRollingWindowsSize)
 
 	return &RollingStoreStats{
 		timeMedians: timeMedians,
@@ -203,11 +205,13 @@ func (r *RollingStoreStats) Observe(stats *pdpb.StoreStats) {
 }
 
 // ObserveRegionsStats records current statistics from region stats.
-func (r *RollingStoreStats) ObserveRegionsStats(writeBytesRate, writeKeysRate float64) {
+func (r *RollingStoreStats) ObserveRegionsStats(writeRate core.StoreWriteRate) {
 	r.Lock()
 	defer r.Unlock()
-	r.movingAvgs[StoreRegionsWriteBytes].Add(writeBytesRate)
-	r.movingAvgs[StoreRegionsWriteKeys].Add(writeKeysRate)
+	r.movingAvgs[StoreRegionsWriteBytes].Add(writeRate.PeerBytesRate)
+	r.movingAvgs[StoreRegionsWriteKeys].Add(writeRate.PeerKeysRate)
+	r.movingAvgs[StoreLeadersWriteBytes].Add(writeRate.LeaderBytesRate)
+	r.movingAvgs[StoreLeadersWriteKeys].Add(writeRate.LeaderKeysRate)
 }
 
 // Set sets the statistics (for test).
@@ -232,11 +236,13 @@ func (r *RollingStoreStats) Set(stats *pdpb.StoreStats) {
 }
 
 // SetRegionsStats sets the statistics from region stats (for test).
-func (r *RollingStoreStats) SetRegionsStats(writeBytesRate, writeKeysRate float64) {
+func (r *RollingStoreStats) SetRegionsStats(writeRate core.StoreWriteRate) {
 	r.Lock()
 	defer r.Unlock()
-	r.movingAvgs[StoreRegionsWriteBytes].Set(writeBytesRate)
-	r.movingAvgs[StoreRegionsWriteKeys].Set(writeKeysRate)
+	r.movingAvgs[StoreRegionsWriteBytes].Set(writeRate.PeerBytesRate)
+	r.movingAvgs[StoreRegionsWriteKeys].Set(writeRate.PeerKeysRate)
+	r.movingAvgs[StoreLeadersWriteBytes].Set(writeRate.LeaderBytesRate)
+	r.movingAvgs[StoreLeadersWriteKeys].Set(writeRate.LeaderKeysRate)
 }
 
 // GetLoad returns store's load.
@@ -246,7 +252,7 @@ func (r *RollingStoreStats) GetLoad(k StoreStatKind) float64 {
 	switch k {
 	case StoreReadBytes, StoreReadKeys, StoreReadQuery, StoreWriteBytes, StoreWriteKeys, StoreWriteQuery:
 		return r.timeMedians[k].Get()
-	case StoreCPUUsage, StoreDiskReadRate, StoreDiskWriteRate, StoreRegionsWriteBytes, StoreRegionsWriteKeys:
+	case StoreCPUUsage, StoreDiskReadRate, StoreDiskWriteRate, StoreRegionsWriteBytes, StoreRegionsWriteKeys, StoreLeadersWriteBytes, StoreLeadersWriteKeys:
 		return r.movingAvgs[k].Get()
 	}
 	return 0
@@ -259,7 +265,7 @@ func (r *RollingStoreStats) GetInstantLoad(k StoreStatKind) float64 {
 	switch k {
 	case StoreReadBytes, StoreReadKeys, StoreReadQuery, StoreWriteBytes, StoreWriteKeys, StoreWriteQuery:
 		return r.timeMedians[k].GetInstantaneous()
-	case StoreCPUUsage, StoreDiskReadRate, StoreDiskWriteRate, StoreRegionsWriteBytes, StoreRegionsWriteKeys:
+	case StoreCPUUsage, StoreDiskReadRate, StoreDiskWriteRate, StoreRegionsWriteBytes, StoreRegionsWriteKeys, StoreLeadersWriteBytes, StoreLeadersWriteKeys:
 		return r.movingAvgs[k].GetInstantaneous()
 	}
 	return 0

server/statistics/store_collection.go

@@ -125,6 +125,8 @@ func (s *storeStatistics) Observe(store *core.StoreInfo, stats *StoresStats) {
 	storeStatusGauge.WithLabelValues(storeAddress, id, "store_disk_write_rate").Set(storeFlowStats.GetLoad(StoreDiskWriteRate))
 	storeStatusGauge.WithLabelValues(storeAddress, id, "store_regions_write_rate_bytes").Set(storeFlowStats.GetLoad(StoreRegionsWriteBytes))
 	storeStatusGauge.WithLabelValues(storeAddress, id, "store_regions_write_rate_keys").Set(storeFlowStats.GetLoad(StoreRegionsWriteKeys))
+	storeStatusGauge.WithLabelValues(storeAddress, id, "store_leaders_write_rate_bytes").Set(storeFlowStats.GetLoad(StoreLeadersWriteBytes))
+	storeStatusGauge.WithLabelValues(storeAddress, id, "store_leaders_write_rate_keys").Set(storeFlowStats.GetLoad(StoreLeadersWriteKeys))
 
 	storeStatusGauge.WithLabelValues(storeAddress, id, "store_write_rate_bytes_instant").Set(storeFlowStats.GetInstantLoad(StoreWriteBytes))
 	storeStatusGauge.WithLabelValues(storeAddress, id, "store_read_rate_bytes_instant").Set(storeFlowStats.GetInstantLoad(StoreReadBytes))
@@ -134,6 +136,8 @@ func (s *storeStatistics) Observe(store *core.StoreInfo, stats *StoresStats) {
 	storeStatusGauge.WithLabelValues(storeAddress, id, "store_read_query_rate_instant").Set(storeFlowStats.GetInstantLoad(StoreReadQuery))
 	storeStatusGauge.WithLabelValues(storeAddress, id, "store_regions_write_rate_bytes_instant").Set(storeFlowStats.GetInstantLoad(StoreRegionsWriteBytes))
 	storeStatusGauge.WithLabelValues(storeAddress, id, "store_regions_write_rate_keys_instant").Set(storeFlowStats.GetInstantLoad(StoreRegionsWriteKeys))
+	storeStatusGauge.WithLabelValues(storeAddress, id, "store_leaders_write_rate_bytes_instant").Set(storeFlowStats.GetInstantLoad(StoreLeadersWriteBytes))
+	storeStatusGauge.WithLabelValues(storeAddress, id, "store_leaders_write_rate_keys_instant").Set(storeFlowStats.GetInstantLoad(StoreLeadersWriteKeys))
 }
 
 func (s *storeStatistics) Collect() {

tests/pdctl/hot/hot_test.go

@@ -96,7 +96,13 @@ func (s *hotTestSuite) TestHot(c *C) {
 	}
 
 	for i := statistics.RegionsStatsRollingWindowsSize; i > 0; i-- {
-		rc.GetStoresStats().ObserveRegionsStats([]uint64{2}, []float64{float64(bytesWritten)}, []float64{float64(keysWritten)})
+		rc.GetStoresStats().ObserveRegionsStats([]core.StoreWriteRate{{
+			2,
+			float64(bytesWritten),
+			float64(keysWritten),
+			float64(bytesWritten),
+			float64(keysWritten),
+		}})
 	}
 
 	args := []string{"-u", pdAddr, "hot", "store"}