storage: add a valid field to rule solver's candidates

Solve now returns all candidates instead of just the valid ones. To
get only the valid candidates, the new function onlyValid and new type
alias condidateList have also been added.

This change allows us to use solve for removeTarget. It also cleans
up the logic in removeTarget to more closely match the non-rule solver
version.

pkg/storage/allocator.go

@@ -220,16 +220,13 @@ func (a *Allocator) AllocateTarget(
 			return nil, errors.Errorf("%d matching stores are currently throttled", throttledStoreCount)
 		}
 
-		candidates, err := allocateRuleSolver.Solve(
+		candidates := allocateRuleSolver.Solve(
 			sl,
 			constraints,
 			existing,
 			a.storePool.getNodeLocalities(existing),
 		)
-		if err != nil {
-			return nil, err
-		}
-
+		candidates = candidates.onlyValid()
 		if len(candidates) == 0 {
 			return nil, &allocatorError{
 				required: constraints.Constraints,
@@ -289,51 +286,6 @@ func (a Allocator) RemoveTarget(
 		return roachpb.ReplicaDescriptor{}, errors.Errorf("must supply at least one replica to allocator.RemoveTarget()")
 	}
 
-	if a.options.UseRuleSolver {
-		// TODO(bram): #10275 Is this getStoreList call required? Compute candidate
-		// requires a store list, but we should be able to create one using only
-		// the stores that belong to the range.
-		// Use an invalid range ID as we don't care about a corrupt replicas since
-		// as we are removing a replica and not trying to add one.
-		sl, _, _ := a.storePool.getStoreList(roachpb.RangeID(0))
-
-		worstCandidate := candidate{constraint: math.Inf(0)}
-		var worstReplica roachpb.ReplicaDescriptor
-		for _, exist := range existing {
-			if exist.StoreID == leaseStoreID {
-				continue
-			}
-			desc, ok := a.storePool.getStoreDescriptor(exist.StoreID)
-			if !ok {
-				continue
-			}
-
-			currentCandidate, valid := removeRuleSolver.computeCandidate(desc, solveState{
-				constraints:            constraints,
-				sl:                     sl,
-				existing:               nil,
-				existingNodeLocalities: a.storePool.getNodeLocalities(existing),
-			})
-			// When a candidate is not valid, it means that it can be
-			// considered the worst existing replica.
-			if !valid {
-				return exist, nil
-			}
-
-			if currentCandidate.less(worstCandidate) {
-				worstCandidate = currentCandidate
-				worstReplica = exist
-			}
-
-		}
-
-		if !math.IsInf(worstCandidate.constraint, 0) {
-			return worstReplica, nil
-		}
-
-		return roachpb.ReplicaDescriptor{}, errors.New("could not select an appropriate replica to be removed")
-	}
-
 	// Retrieve store descriptors for the provided replicas from the StorePool.
 	descriptors := make([]roachpb.StoreDescriptor, 0, len(existing))
 	for _, exist := range existing {
@@ -344,15 +296,37 @@ func (a Allocator) RemoveTarget(
 			descriptors = append(descriptors, desc)
 		}
 	}
-
 	sl := makeStoreList(descriptors)
-	if bad := a.selectBad(sl); bad != nil {
+	var badStoreID roachpb.StoreID
+
+	if a.options.UseRuleSolver {
+		candidates := removeRuleSolver.Solve(
+			sl,
+			constraints,
+			existing,
+			a.storePool.getNodeLocalities(existing),
+		)
+
+		if len(candidates) != 0 {
+			// TODO(bram): There needs some randomness here and the logic from
+			// selectBad around rebalanceFromConvergesOnMean.
+			badStoreID = candidates[len(candidates)-1].store.StoreID
+		}
+	} else {
+		bad := a.selectBad(sl)
+		if bad != nil {
+			badStoreID = bad.StoreID
+		}
+	}
+
+	if badStoreID != 0 {
 		for _, exist := range existing {
-			if exist.StoreID == bad.StoreID {
+			if exist.StoreID == badStoreID {
 				return exist, nil
 			}
 		}
 	}
+
 	return roachpb.ReplicaDescriptor{}, errors.New("could not select an appropriate replica to be removed")
 }
 
@@ -429,14 +403,8 @@ func (a Allocator) RebalanceTarget(
 		existingStoreList := makeStoreList(existingDescs)
 		candidateStoreList := makeStoreList(candidateDescs)
 
-		existingCandidates, err := removeRuleSolver.Solve(existingStoreList, constraints, nil, nil)
-		if err != nil {
-			return nil, err
-		}
-		candidates, err := allocateRuleSolver.Solve(candidateStoreList, constraints, nil, nil)
-		if err != nil {
-			return nil, err
-		}
+		existingCandidates := removeRuleSolver.Solve(existingStoreList, constraints, nil, nil)
+		candidates := allocateRuleSolver.Solve(candidateStoreList, constraints, nil, nil)
 
 		// Find all candidates that are better than the worst existing store.
 		var worstCandidate candidate

pkg/storage/rule_solver.go

@@ -26,18 +26,37 @@ import (
 // candidate store for allocation.
 type candidate struct {
 	store      roachpb.StoreDescriptor
+	valid      bool
 	constraint float64 // Score used to pick the top candidates.
 	balance    float64 // Score used to choose between top candidates.
 }
 
 // less first compares constraint scores, then balance scores.
 func (c candidate) less(o candidate) bool {
+	if !o.valid {
+		return false
+	}
+	if !c.valid {
+		return true
+	}
 	if c.constraint != o.constraint {
 		return c.constraint < o.constraint
 	}
 	return c.balance < o.balance
 }
 
+type candidateList []candidate
+
+// onlyValid returns all the elements in a sorted candidate list that are valid.
+func (cl candidateList) onlyValid() candidateList {
+	for i := len(cl) - 1; i >= 0; i-- {
+		if cl[i].valid {
+			return cl[:i+1]
+		}
+	}
+	return candidateList{}
+}
+
 // solveState is used to pass solution state information into a rule.
 type solveState struct {
 	constraints            config.Constraints
@@ -71,49 +90,49 @@ var removeRuleSolver = ruleSolver{
 }
 
 // Solve runs the rules against the stores in the store list and returns all
-// passing stores and their scores ordered from best to worst score.
+// candidate stores and their scores ordered from best to worst score.
 func (rs ruleSolver) Solve(
 	sl StoreList,
 	c config.Constraints,
 	existing []roachpb.ReplicaDescriptor,
 	existingNodeLocalities map[roachpb.NodeID]roachpb.Locality,
-) ([]candidate, error) {
-	candidates := make([]candidate, 0, len(sl.stores))
+) candidateList {
+	candidates := make(candidateList, len(sl.stores), len(sl.stores))
 	state := solveState{
 		constraints:            c,
 		sl:                     sl,
 		existing:               existing,
 		existingNodeLocalities: existingNodeLocalities,
 	}
 
-	for _, store := range sl.stores {
-		if cand, ok := rs.computeCandidate(store, state); ok {
-			candidates = append(candidates, cand)
-		}
+	for i, store := range sl.stores {
+		candidates[i] = rs.computeCandidate(store, state)
 	}
 	sort.Sort(sort.Reverse(byScore(candidates)))
-	return candidates, nil
+	return candidates
 }
 
 // computeCandidate runs the rules against a candidate store using the provided
 // state and returns each candidate's score and if the candidate is valid.
-func (rs ruleSolver) computeCandidate(
-	store roachpb.StoreDescriptor, state solveState,
-) (candidate, bool) {
+func (rs ruleSolver) computeCandidate(store roachpb.StoreDescriptor, state solveState) candidate {
 	var totalConstraintScore, totalBalanceScore float64
 	for _, rule := range rs {
 		valid, constraintScore, balanceScore := rule(store, state)
 		if !valid {
-			return candidate{}, false
+			return candidate{
+				store: store,
+				valid: false,
+			}
 		}
 		totalConstraintScore += constraintScore
 		totalBalanceScore += balanceScore
 	}
 	return candidate{
 		store:      store,
+		valid:      true,
 		constraint: totalConstraintScore,
 		balance:    totalBalanceScore,
-	}, true
+	}
 }
 
 // ruleReplicasUniqueNodes returns true iff no existing replica is present on
@@ -201,7 +220,7 @@ func ruleCapacity(store roachpb.StoreDescriptor, state solveState) (bool, float6
 }
 
 // byScore implements sort.Interface to sort by scores.
-type byScore []candidate
+type byScore candidateList
 
 var _ sort.Interface = byScore(nil)