checker, statistic: avoid leak in label statistic

Signed-off-by: lhy1024 <[email protected]>

pkg/cluster/cluster.go

@@ -68,7 +68,7 @@ func HandleOverlaps(ctx context.Context, c Cluster, overlaps []*core.RegionInfo)
 		if c.GetRegionStats() != nil {
 			c.GetRegionStats().ClearDefunctRegion(item.GetID())
 		}
-		c.GetLabelStats().ClearDefunctRegion(item.GetID())
+		c.GetLabelStats().MarkDefunctRegion(item.GetID())
 		c.GetRuleManager().InvalidCache(item.GetID())
 	}
 }

pkg/mcs/scheduling/server/cluster.go

@@ -379,15 +379,18 @@ func (c *Cluster) waitSchedulersInitialized() {
 	}
 }
 
-// TODO: implement the following methods
-
 // UpdateRegionsLabelLevelStats updates the status of the region label level by types.
 func (c *Cluster) UpdateRegionsLabelLevelStats(regions []*core.RegionInfo) {
 	for _, region := range regions {
 		c.labelStats.Observe(region, c.getStoresWithoutLabelLocked(region, core.EngineKey, core.EngineTiFlash), c.persistConfig.GetLocationLabels())
 	}
 }
 
+// ClearDefunctRegionsLabelLevelStats clears the defunct regions' label level stats.
+func (c *Cluster) ClearDefunctRegionsLabelLevelStats() {
+	c.labelStats.ClearDefunctRegions()
+}
+
 func (c *Cluster) getStoresWithoutLabelLocked(region *core.RegionInfo, key, value string) []*core.StoreInfo {
 	stores := make([]*core.StoreInfo, 0, len(region.GetPeers()))
 	for _, p := range region.GetPeers() {

pkg/mock/mockcluster/mockcluster.go

@@ -125,6 +125,9 @@ func (mc *Cluster) AllocID() (uint64, error) {
 // UpdateRegionsLabelLevelStats updates the label level stats for the regions.
 func (*Cluster) UpdateRegionsLabelLevelStats(_ []*core.RegionInfo) {}
 
+// ClearDefunctRegionsLabelLevelStats clears the defunct regions' label level stats.
+func (*Cluster) ClearDefunctRegionsLabelLevelStats() {}
+
 // LoadRegion puts region info without leader
 func (mc *Cluster) LoadRegion(regionID uint64, peerStoreIDs ...uint64) {
 	//  regions load from etcd will have no leader

pkg/schedule/checker/checker_controller.go

@@ -141,6 +141,8 @@ func (c *Controller) PatrolRegions() {
 			c.cluster.UpdateRegionsLabelLevelStats(regions)
 			// When the key is nil, it means that the scan is finished.
 			if len(key) == 0 {
+				// Clear the defunct regions label level statistics.
+				c.cluster.ClearDefunctRegionsLabelLevelStats()
 				// update the scan limit.
 				c.patrolRegionScanLimit = calculateScanLimit(c.cluster)
 				// update the metrics.

pkg/schedule/core/cluster_informer.go

@@ -51,6 +51,7 @@ type CheckerCluster interface {
 	GetCheckerConfig() sc.CheckerConfigProvider
 	GetStoreConfig() sc.StoreConfigProvider
 	UpdateRegionsLabelLevelStats(regions []*core.RegionInfo)
+	ClearDefunctRegionsLabelLevelStats()
 }
 
 // SharedCluster is an aggregate interface that wraps multiple interfaces

pkg/statistics/region_collection.go

@@ -365,13 +365,15 @@ type LabelStatistics struct {
 	syncutil.RWMutex
 	regionLabelStats map[uint64]string
 	labelCounter     map[string]int
+	defunctRegions   map[uint64]struct{}
 }
 
 // NewLabelStatistics creates a new LabelStatistics.
 func NewLabelStatistics() *LabelStatistics {
 	return &LabelStatistics{
 		regionLabelStats: make(map[uint64]string),
 		labelCounter:     make(map[string]int),
+		defunctRegions:   make(map[uint64]struct{}),
 	}
 }
 
@@ -405,14 +407,26 @@ func ResetLabelStatsMetrics() {
 	regionLabelLevelGauge.Reset()
 }
 
-// ClearDefunctRegion is used to handle the overlap region.
-func (l *LabelStatistics) ClearDefunctRegion(regionID uint64) {
+// MarkDefunctRegion is used to handle the overlap region.
+// It is used to mark the region as defunct and remove it from the label statistics later.
+func (l *LabelStatistics) MarkDefunctRegion(regionID uint64) {
 	l.Lock()
 	defer l.Unlock()
-	if label, ok := l.regionLabelStats[regionID]; ok {
-		l.labelCounter[label]--
-		delete(l.regionLabelStats, regionID)
+	l.defunctRegions[regionID] = struct{}{}
+}
+
+// ClearDefunctRegions is used to handle the overlap region.
+// It is used to remove the defunct regions from the label statistics.
+func (l *LabelStatistics) ClearDefunctRegions() {
+	l.Lock()
+	defer l.Unlock()
+	for regionID := range l.defunctRegions {
+		if label, ok := l.regionLabelStats[regionID]; ok {
+			l.labelCounter[label]--
+			delete(l.regionLabelStats, regionID)
+		}
 	}
+	l.defunctRegions = make(map[uint64]struct{})
 }
 
 // GetLabelCounter is only used for tests.

server/cluster/scheduling_controller.go

@@ -238,6 +238,11 @@ func (sc *schedulingController) UpdateRegionsLabelLevelStats(regions []*core.Reg
 	}
 }
 
+// ClearDefunctRegionsLabelLevelStats clears the status of the region label level by types.
+func (sc *schedulingController) ClearDefunctRegionsLabelLevelStats() {
+	sc.labelStats.ClearDefunctRegions()
+}
+
 func (sc *schedulingController) getStoresWithoutLabelLocked(region *core.RegionInfo, key, value string) []*core.StoreInfo {
 	stores := make([]*core.StoreInfo, 0, len(region.GetPeers()))
 	for _, p := range region.GetPeers() {