replica_strategy:speed up and reduce the complexity of selectStore to O(n)

server/schedule/checker/replica_strategy.go

@@ -72,12 +72,12 @@ func (s *ReplicaStrategy) SelectStoreToAdd(coLocationStores []*core.StoreInfo, e
 	strictStateFilter := &filter.StoreStateFilter{ActionScope: s.checkerName, MoveRegion: true}
 	targetCandidate := filter.NewCandidates(s.cluster.GetStores()).
 		FilterTarget(s.cluster.GetOpts(), filters...).
-		Sort(isolationComparer).Reverse().Top(isolationComparer). // greater isolation score is better
-		Sort(filter.RegionScoreComparer(s.cluster.GetOpts()))     // less region score is better
+		KeepTheTopStores(isolationComparer, false) // greater isolation score is better
 	if targetCandidate.Len() == 0 {
 		return 0, false
 	}
-	target := targetCandidate.FilterTarget(s.cluster.GetOpts(), strictStateFilter).PickFirst() // the filter does not ignore temp states
+	target := targetCandidate.FilterTarget(s.cluster.GetOpts(), strictStateFilter).
+		PickTheTopStore(filter.RegionScoreComparer(s.cluster.GetOpts()), true) // less region score is better
 	if target == nil {
 		return 0, true // filter by temporary states
 	}
@@ -124,9 +124,8 @@ func (s *ReplicaStrategy) SelectStoreToRemove(coLocationStores []*core.StoreInfo
 	isolationComparer := filter.IsolationComparer(s.locationLabels, coLocationStores)
 	source := filter.NewCandidates(coLocationStores).
 		FilterSource(s.cluster.GetOpts(), &filter.StoreStateFilter{ActionScope: replicaCheckerName, MoveRegion: true}).
-		Sort(isolationComparer).Top(isolationComparer).
-		Sort(filter.RegionScoreComparer(s.cluster.GetOpts())).Reverse().
-		PickFirst()
+		KeepTheTopStores(isolationComparer, true).
+		PickTheTopStore(filter.RegionScoreComparer(s.cluster.GetOpts()), false)
 	if source == nil {
 		log.Debug("no removable store", zap.Uint64("region-id", s.region.GetID()))
 		return 0

server/schedule/filter/candidates.go

@@ -51,32 +51,49 @@ func (c *StoreCandidates) Sort(less StoreComparer) *StoreCandidates {
 	return c
 }
 
-// Reverse reverses the candidate store list.
-func (c *StoreCandidates) Reverse() *StoreCandidates {
-	for i := len(c.Stores)/2 - 1; i >= 0; i-- {
-		opp := len(c.Stores) - 1 - i
-		c.Stores[i], c.Stores[opp] = c.Stores[opp], c.Stores[i]
-	}
-	return c
-}
-
 // Shuffle reorders all candidates randomly.
 func (c *StoreCandidates) Shuffle() *StoreCandidates {
 	rand.Shuffle(len(c.Stores), func(i, j int) { c.Stores[i], c.Stores[j] = c.Stores[j], c.Stores[i] })
 	return c
 }
 
-// Top keeps all stores that have the same priority with the first store.
-// The store list should be sorted before calling Top.
-func (c *StoreCandidates) Top(less StoreComparer) *StoreCandidates {
-	var i int
-	for i < len(c.Stores) && less(c.Stores[0], c.Stores[i]) == 0 {
-		i++
+// KeepTheTopStores keeps the slice of the stores in the front order by asc.
+func (c *StoreCandidates) KeepTheTopStores(cmp StoreComparer, asc bool) *StoreCandidates {
+	if len(c.Stores) <= 1 {
+		return c
+	}
+	topIdx := 0
+	for idx := 1; idx < c.Len(); idx++ {
+		compare := cmp(c.Stores[topIdx], c.Stores[idx])
+		if compare == 0 {
+			topIdx++
+		} else if (compare > 0 && asc) || (!asc && compare < 0) {
+			topIdx = 0
+		} else {
+			continue
+		}
+		c.Stores[idx], c.Stores[topIdx] = c.Stores[topIdx], c.Stores[idx]
 	}
-	c.Stores = c.Stores[:i]
+	c.Stores = c.Stores[:topIdx+1]
 	return c
 }
 
+// PickTheTopStore returns the first store order by asc.
+// it returns the min item when asc is true, return the max item when asc is false.
+func (c *StoreCandidates) PickTheTopStore(cmp StoreComparer, asc bool) *core.StoreInfo {
+	if len(c.Stores) == 0 {
+		return nil
+	}
+	topIdx := 0
+	for idx := 1; idx < len(c.Stores); idx++ {
+		compare := cmp(c.Stores[topIdx], c.Stores[idx])
+		if (compare > 0 && asc) || (!asc && compare < 0) {
+			topIdx = idx
+		}
+	}
+	return c.Stores[topIdx]
+}
+
 // PickFirst returns the first store in candidate list.
 func (c *StoreCandidates) PickFirst() *core.StoreInfo {
 	if len(c.Stores) == 0 {

server/schedule/filter/candidates_test.go

@@ -15,9 +15,10 @@
 package filter
 
 import (
+	"testing"
+
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/stretchr/testify/require"
-	"testing"
 
 	"github.com/tikv/pd/server/config"
 	"github.com/tikv/pd/server/core"
@@ -71,13 +72,16 @@ func TestCandidates(t *testing.T) {
 	re.Nil(store)
 
 	cs = newTestCandidates(1, 3, 5, 7, 6, 2, 4)
+	minStore := cs.PickTheTopStore(idComparer, true)
+	re.Equal(uint64(1), minStore.GetID())
+	maxStore := cs.PickTheTopStore(idComparer, false)
+	re.Equal(uint64(7), maxStore.GetID())
+
 	cs.Sort(idComparer)
 	check(re, cs, 1, 2, 3, 4, 5, 6, 7)
 	store = cs.PickFirst()
 	re.Equal(uint64(1), store.GetID())
-	cs.Reverse()
-	check(re, cs, 7, 6, 5, 4, 3, 2, 1)
-	store = cs.PickFirst()
+	store = cs.PickTheTopStore(idComparer, false)
 	re.Equal(uint64(7), store.GetID())
 	cs.Shuffle()
 	cs.Sort(idComparer)
@@ -87,7 +91,7 @@ func TestCandidates(t *testing.T) {
 	re.Less(store.GetID(), uint64(8))
 
 	cs = newTestCandidates(10, 15, 23, 20, 33, 32, 31)
-	cs.Sort(idComparer).Reverse().Top(idComparer2)
+	cs.KeepTheTopStores(idComparer2, false)
 	check(re, cs, 33, 32, 31)
 }