schedule/placement/fit: remove the recursion in pick peers for on rule (

#5269)

close #5268

Signed-off-by: shirly <AndreMouche@126.com>

Co-authored-by: Ti Chi Robot <ti-community-prow-bot@tidb.io>

server/schedule/placement/fit.go

@@ -16,6 +16,7 @@ package placement
 
 import (
 	"math"
+	"math/bits"
 	"sort"
 
 	"github.com/pingcap/kvproto/pkg/metapb"
@@ -227,34 +228,60 @@ func (w *fitWorker) fitRule(index int) bool {
 	if len(candidates) < count {
 		count = len(candidates)
 	}
-	return w.enumPeers(candidates, nil, index, count)
+
+	return w.fixRuleWithCandidates(candidates, index, count)
 }
 
-// Recursively traverses all feasible peer combinations.
-// For each combination, call `compareBest` to determine whether it is better
-// than the existing option.
+// Pick the most suitable peer combination for the rule with candidates.
 // Returns true if it replaces `bestFit` with a better alternative.
-func (w *fitWorker) enumPeers(candidates, selected []*fitPeer, index int, count int) bool {
-	if len(selected) == count {
-		// We collect enough peers. End recursive.
-		return w.compareBest(selected, index)
-	}
+func (w *fitWorker) fixRuleWithCandidates(candidates []*fitPeer, index int, count int) bool {
+	// map the candidates to binary numbers with len(candidates) bits,
+	// each bit can be 1 or 0, 1 means a picked candidate
+	// the binary numbers with `count` 1 means a choose for the current rule.
 
 	var better bool
-	// make sure the left number of candidates should be enough.
-	indexLimit := len(candidates) - (count - len(selected))
-	for i := 0; i <= indexLimit; i++ {
-		p := candidates[i]
-		p.selected = true
-		better = w.enumPeers(candidates[i+1:], append(selected, p), index, count) || better
-		p.selected = false
+	limit := uint(1<<len(candidates) - 1)
+	binaryInt := uint(1<<count - 1)
+	for ; binaryInt <= limit; binaryInt++ {
+		// there should be exactly `count` number in current binary number `binaryInt`
+		if bits.OnesCount(binaryInt) != count {
+			continue
+		}
+		selected := pickPeersFromBinaryInt(candidates, binaryInt)
+		better = w.compareBest(selected, index) || better
+		// reset the seleted items to false.
+		unSelectPeers(selected)
 		if w.exit {
 			break
 		}
 	}
 	return better
 }
 
+// pickPeersFromBinaryInt picks the candidates with the related index at the position of binary for the `binaryNumber`` is `1`.
+// binaryNumber = 5, which means the related binary is 101, it will returns {candidates[0],candidates[2]}
+// binaryNumber = 6, which means the related binary is 110, it will returns {candidates[1],candidates[2]}
+func pickPeersFromBinaryInt(candidates []*fitPeer, binaryNumber uint) []*fitPeer {
+	selected := make([]*fitPeer, 0)
+	for _, p := range candidates {
+		if binaryNumber&1 == 1 {
+			p.selected = true
+			selected = append(selected, p)
+		}
+		binaryNumber >>= 1
+		if binaryNumber == 0 {
+			break
+		}
+	}
+	return selected
+}
+
+func unSelectPeers(seleted []*fitPeer) {
+	for _, p := range seleted {
+		p.selected = false
+	}
+}
+
 // compareBest checks if the selected peers is better then previous best.
 // Returns true if it replaces `bestFit` with a better alternative.
 func (w *fitWorker) compareBest(selected []*fitPeer, index int) bool {

server/schedule/placement/fit_test.go

@@ -187,3 +187,35 @@ func TestIsolationScore(t *testing.T) {
 		testCase.checker(score1, score2)
 	}
 }
+
+func TestPickPeersFromBinaryInt(t *testing.T) {
+	re := require.New(t)
+	var candidates []*fitPeer
+	for id := uint64(1); id <= 10; id++ {
+		candidates = append(candidates, &fitPeer{
+			Peer: &metapb.Peer{Id: id},
+		})
+	}
+	testCases := []struct {
+		binary        string
+		expectedPeers []uint64
+	}{
+		{"0", []uint64{}},
+		{"1", []uint64{1}},
+		{"101", []uint64{1, 3}},
+		{"111", []uint64{1, 2, 3}},
+		{"1011", []uint64{1, 2, 4}},
+		{"100011", []uint64{1, 2, 6}},
+		{"1000001111", []uint64{1, 2, 3, 4, 10}},
+	}
+
+	for _, c := range testCases {
+		binaryNumber, err := strconv.ParseUint(c.binary, 2, 64)
+		re.NoError(err)
+		selected := pickPeersFromBinaryInt(candidates, uint(binaryNumber))
+		re.Len(selected, len(c.expectedPeers))
+		for id := 0; id < len(selected); id++ {
+			re.Equal(selected[id].Id, c.expectedPeers[id])
+		}
+	}
+}