schedule/labeler: Optimize locking to fix high-concurrency goroutine surge (#9857) (#9897)

close #9854

Signed-off-by: lhy1024 <admin@liudos.us>

Co-authored-by: lhy1024 <liuhanyang@pingcap.com>

Author: ti-chi-bot

pkg/schedule/labeler/labeler.go

@@ -156,81 +156,139 @@ func (l *RegionLabeler) GetSplitKeys(start, end []byte) [][]byte {
 	return l.rangeList.GetSplitKeys(start, end)
 }
 
+func filterExpiredLabels(rule *LabelRule, now time.Time) *LabelRule {
+	if rule == nil {
+		return nil
+	}
+
+	hasExpired := false
+	for i := range rule.Labels {
+		if rule.Labels[i].expireBefore(now) {
+			hasExpired = true
+			break
+		}
+	}
+	if !hasExpired {
+		return rule
+	}
+
+	labels := make([]RegionLabel, 0, len(rule.Labels))
+	for i := range rule.Labels {
+		if !rule.Labels[i].expireBefore(now) {
+			labels = append(labels, rule.Labels[i])
+		}
+	}
+	if len(labels) == 0 {
+		return nil
+	}
+	return &LabelRule{
+		ID:       rule.ID,
+		Index:    rule.Index,
+		RuleType: rule.RuleType,
+		Data:     rule.Data,
+		Labels:   labels,
+	}
+}
+
 // GetAllLabelRules returns all the rules.
 func (l *RegionLabeler) GetAllLabelRules() []*LabelRule {
-	l.checkAndClearExpiredLabels()
 	l.RLock()
 	defer l.RUnlock()
+
+	now := time.Now()
 	rules := make([]*LabelRule, 0, len(l.labelRules))
 	for _, rule := range l.labelRules {
-		rules = append(rules, rule)
+		if filteredRule := filterExpiredLabels(rule, now); filteredRule != nil {
+			rules = append(rules, filteredRule)
+		}
 	}
 	return rules
 }
 
 // GetLabelRules returns the rules that match the given ids.
 func (l *RegionLabeler) GetLabelRules(ids []string) ([]*LabelRule, error) {
+	l.RLock()
+	defer l.RUnlock()
+
 	now := time.Now()
 	rules := make([]*LabelRule, 0, len(ids))
 	for _, id := range ids {
-		if rule := l.getAndCheckRule(id, now); rule != nil {
-			rules = append(rules, rule)
+		if rule, ok := l.labelRules[id]; ok {
+			if filteredRule := filterExpiredLabels(rule, now); filteredRule != nil {
+				rules = append(rules, filteredRule)
+			}
 		}
 	}
 	return rules, nil
 }
 
 // GetLabelRule returns the Rule with the same ID.
 func (l *RegionLabeler) GetLabelRule(id string) *LabelRule {
-	return l.getAndCheckRule(id, time.Now())
-}
-
-func (l *RegionLabeler) getAndCheckRule(id string, now time.Time) *LabelRule {
-	l.Lock()
-	defer l.Unlock()
+	l.RLock()
+	defer l.RUnlock()
 	rule, ok := l.labelRules[id]
 	if !ok {
 		return nil
 	}
-	if !rule.checkAndRemoveExpireLabels(now) {
-		return rule
-	}
-	if len(rule.Labels) == 0 {
-		l.storage.DeleteRegionRule(id)
-		delete(l.labelRules, id)
-		return nil
-	}
-	l.storage.SaveRegionRule(id, rule)
-	return rule
+	return filterExpiredLabels(rule, time.Now())
 }
 
 // SetLabelRule inserts or updates a LabelRule.
 func (l *RegionLabeler) SetLabelRule(rule *LabelRule) error {
 	if err := rule.checkAndAdjust(); err != nil {
 		return err
 	}
+	if err := l.storage.SaveRegionRule(rule.ID, rule); err != nil {
+		return err
+	}
+
+	// only Lock for in-memory update
 	l.Lock()
 	defer l.Unlock()
+	l.labelRules[rule.ID] = rule
+	l.buildRangeList()
+	return nil
+}
+
+// SetLabelRuleLocked inserts or updates a LabelRule but not buildRangeList.
+// It updates the in-memory states and storage at the same time.
+// It should be used in watcher.
+func (l *RegionLabeler) SetLabelRuleLocked(rule *LabelRule) error {
+	if err := rule.checkAndAdjust(); err != nil {
+		return err
+	}
 	if err := l.storage.SaveRegionRule(rule.ID, rule); err != nil {
 		return err
 	}
 	l.labelRules[rule.ID] = rule
-	l.buildRangeList()
 	return nil
 }
 
 // DeleteLabelRule removes a LabelRule.
 func (l *RegionLabeler) DeleteLabelRule(id string) error {
+	if err := l.storage.DeleteRegionRule(id); err != nil {
+		return err
+	}
+
+	// only Lock for in-memory update
 	l.Lock()
 	defer l.Unlock()
 	if _, ok := l.labelRules[id]; !ok {
-		return errs.ErrRegionRuleNotFound.FastGenByArgs(id)
+		return nil
 	}
+	delete(l.labelRules, id)
+	l.buildRangeList()
+	return nil
+}
+
+// DeleteLabelRuleLocked removes a LabelRule but not buildRangeList.
+// It updates the in-memory states and storage at the same time.
+// It should be used in watcher.
+func (l *RegionLabeler) DeleteLabelRuleLocked(id string) error {
 	if err := l.storage.DeleteRegionRule(id); err != nil {
 		return err
 	}
 	delete(l.labelRules, id)
-	l.buildRangeList()
 	return nil
 }
 

pkg/schedule/labeler/labeler_test.go

@@ -312,18 +312,19 @@ func TestLabelerRuleTTL(t *testing.T) {
 
 	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/schedule/labeler/regionLabelExpireSub1Minute"))
 	// rule2 should be exist since `GetRegionLabels` won't clear it physically.
-	checkRuleInMemoryAndStoage(re, labeler, "rule2", true)
+	checkRuleInMemoryAndStorage(re, labeler, "rule2", true)
 	re.Nil(labeler.GetLabelRule("rule2"))
 	// rule2 should be physically clear.
-	checkRuleInMemoryAndStoage(re, labeler, "rule2", false)
+	labeler.checkAndClearExpiredLabels()
+	checkRuleInMemoryAndStorage(re, labeler, "rule2", false)
 
 	re.Equal("", labeler.GetRegionLabel(region, "k2"))
 
 	re.NotNil(labeler.GetLabelRule("rule3"))
 	re.NotNil(labeler.GetLabelRule("rule1"))
 }
 
-func checkRuleInMemoryAndStoage(re *require.Assertions, labeler *RegionLabeler, ruleID string, exist bool) {
+func checkRuleInMemoryAndStorage(re *require.Assertions, labeler *RegionLabeler, ruleID string, exist bool) {
 	re.Equal(exist, labeler.labelRules[ruleID] != nil)
 	existInStorage := false
 	labeler.storage.LoadRegionRules(func(k, v string) {