From 324f16fd976e4d4fcb47446c83768c5109df5475 Mon Sep 17 00:00:00 2001 From: Bram Gruneir Date: Wed, 7 Dec 2016 11:58:08 -0500 Subject: [PATCH] storage: replace the rule solver Remove the concept of rules and replace it with clean direct functions. --- pkg/storage/allocator.go | 100 ++--- pkg/storage/allocator_test.go | 101 +++-- pkg/storage/rule_solver.go | 393 +++++++++++--------- pkg/storage/rule_solver_test.go | 628 +++++++++++++++++++------------- 4 files changed, 661 insertions(+), 561 deletions(-) diff --git a/pkg/storage/allocator.go b/pkg/storage/allocator.go index e812a524efdd..6833d4448a19 100644 --- a/pkg/storage/allocator.go +++ b/pkg/storage/allocator.go @@ -212,29 +212,31 @@ func (a *Allocator) AllocateTarget( rangeID roachpb.RangeID, relaxConstraints bool, ) (*roachpb.StoreDescriptor, error) { + sl, aliveStoreCount, throttledStoreCount := a.storePool.getStoreList(rangeID) if a.options.UseRuleSolver { - sl, _, throttledStoreCount := a.storePool.getStoreList(rangeID) - // When there are throttled stores that do match, we shouldn't send - // the replica to purgatory. - if throttledStoreCount > 0 { - return nil, errors.Errorf("%d matching stores are currently throttled", throttledStoreCount) - } - - candidates := allocateRuleSolver.Solve( + candidates := allocateCandidates( sl, constraints, existing, a.storePool.getNodeLocalities(existing), a.storePool.deterministic, ) - candidates = candidates.onlyValid() - if len(candidates) == 0 { - return nil, &allocatorError{ - required: constraints.Constraints, - } + if log.V(3) { + log.Infof(context.TODO(), "allocate candidates: %s", candidates) + } + if len(candidates) != 0 { + chosenCandidate := candidates.selectGood(a.randGen).store + return &chosenCandidate, nil + } + + // When there are throttled stores that do match, we shouldn't send + // the replica to purgatory. + if throttledStoreCount > 0 { + return nil, errors.Errorf("%d matching stores are currently throttled", throttledStoreCount) + } + return nil, &allocatorError{ + required: constraints.Constraints, } - chosenCandidate := candidates.selectGood(a.randGen).store - return &chosenCandidate, nil } existingNodes := make(nodeIDSet, len(existing)) @@ -242,8 +244,6 @@ func (a *Allocator) AllocateTarget( existingNodes[repl.NodeID] = struct{}{} } - sl, aliveStoreCount, throttledStoreCount := a.storePool.getStoreList(rangeID) - // Because more redundancy is better than less, if relaxConstraints, the // matching here is lenient, and tries to find a target by relaxing an // attribute constraint, from last attribute to first. @@ -273,11 +273,6 @@ func (a *Allocator) AllocateTarget( // that have greater than the average number of replicas. Failing that, it // falls back to selecting a random target from any of the existing // replicas. It also will exclude any replica that lives on leaseStoreID. -// -// TODO(mrtracy): removeTarget eventually needs to accept the attributes from -// the zone config associated with the provided replicas. This will allow it to -// make correct decisions in the case of ranges with heterogeneous replica -// requirements (i.e. multiple data centers). func (a Allocator) RemoveTarget( constraints config.Constraints, existing []roachpb.ReplicaDescriptor, @@ -301,13 +296,15 @@ func (a Allocator) RemoveTarget( var badStoreID roachpb.StoreID if a.options.UseRuleSolver { - candidates := removeRuleSolver.Solve( + candidates := removeCandidates( sl, constraints, - existing, a.storePool.getNodeLocalities(existing), a.storePool.deterministic, ) + if log.V(3) { + log.Infof(context.TODO(), "remove candidates: %s", candidates) + } if len(candidates) != 0 { badStoreID = candidates.selectBad(a.randGen).store.StoreID } @@ -358,13 +355,11 @@ func (a Allocator) RebalanceTarget( if !a.options.AllowRebalance { return nil, nil } + sl, _, _ := a.storePool.getStoreList(rangeID) if a.options.UseRuleSolver { - sl, _, _ := a.storePool.getStoreList(rangeID) - if log.V(3) { - log.Infof(context.TODO(), "rebalance-target (lease-holder=%d):\n%s", leaseStoreID, sl) - } - + // TODO(bram): ShouldRebalance should be part of rebalanceCandidates + // and decision made afterward, not it's own function. var shouldRebalance bool for _, repl := range existing { if leaseStoreID == repl.StoreID { @@ -380,50 +375,20 @@ func (a Allocator) RebalanceTarget( return nil, nil } - // Load the exiting storesIDs into a map so to eliminate having to loop - // through the existing descriptors more than once. - existingStoreIDs := make(map[roachpb.StoreID]struct{}) - for _, repl := range existing { - existingStoreIDs[repl.StoreID] = struct{}{} - } - - // Split the store list into existing and candidate stores lists so that - // we can call solve independently on both store lists. - var existingDescs []roachpb.StoreDescriptor - var candidateDescs []roachpb.StoreDescriptor - for _, desc := range sl.stores { - if _, ok := existingStoreIDs[desc.StoreID]; ok { - existingDescs = append(existingDescs, desc) - } else { - candidateDescs = append(candidateDescs, desc) - } - } - - existingStoreList := makeStoreList(existingDescs) - candidateStoreList := makeStoreList(candidateDescs) - - localities := a.storePool.getNodeLocalities(existing) - existingCandidates := rebalanceExisting.Solve( - existingStoreList, - constraints, - existing, - localities, - a.storePool.deterministic, - ) - candidates := rebalance.Solve( - candidateStoreList, + existingCandidates, candidates := rebalanceCandidates( + sl, constraints, existing, - localities, + a.storePool.getNodeLocalities(existing), a.storePool.deterministic, ) - /* - fmt.Printf("existing: %s\n", existingCandidates) - fmt.Printf("candidates: %s\n", candidates) - */ + if log.V(3) { + log.Infof(context.TODO(), "existing replicas: %s", existingCandidates) + log.Infof(context.TODO(), "candidates: %s", candidates) + } - // Find all candidates that are better than the worst existing. + // Find all candidates that are better than the worst existing replica. targets := candidates.betterThan(existingCandidates[len(existingCandidates)-1]) if len(targets) != 0 { target := targets.selectGood(a.randGen).store @@ -432,7 +397,6 @@ func (a Allocator) RebalanceTarget( return nil, nil } - sl, _, _ := a.storePool.getStoreList(rangeID) sl = sl.filter(constraints) if log.V(3) { log.Infof(context.TODO(), "rebalance-target (lease-holder=%d):\n%s", leaseStoreID, sl) diff --git a/pkg/storage/allocator_test.go b/pkg/storage/allocator_test.go index 45e251a8656c..34a7c1f261a6 100644 --- a/pkg/storage/allocator_test.go +++ b/pkg/storage/allocator_test.go @@ -715,7 +715,6 @@ func TestAllocatorRebalance(t *testing.T) { if err != nil { t.Fatal(err) } - fmt.Printf("%d: %s\n", i, result) if result == nil { i-- // loop until we find 10 candidates continue @@ -1925,55 +1924,55 @@ func Example_rebalancingWithRuleSolver() { // | gen | store 0 | store 1 | store 2 | store 3 | store 4 | store 5 | store 6 | store 7 | store 8 | store 9 | store 10 | store 11 | store 12 | store 13 | store 14 | store 15 | store 16 | store 17 | store 18 | store 19 | // +-----+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+ // | 0 | 1 48% | 0 0% | 0 0% | 1 51% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | - // | 2 | 2 35% | 0 0% | 0 0% | 2 24% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 1 11% | 0 0% | 1 14% | 0 0% | 1 13% | 0 0% | 0 0% | 0 0% | - // | 4 | 2 19% | 0 0% | 0 0% | 2 8% | 0 0% | 0 0% | 1 7% | 2 11% | 0 0% | 1 1% | 2 9% | 0 0% | 2 3% | 1 1% | 2 9% | 1 9% | 2 4% | 0 0% | 2 13% | 0 0% | - // | 6 | 2 10% | 3 5% | 3 6% | 2 1% | 3 7% | 3 3% | 3 3% | 2 7% | 2 1% | 2 3% | 2 7% | 4 5% | 3 2% | 2 2% | 2 2% | 2 4% | 2 0% | 3 10% | 2 8% | 2 1% | - // | 8 | 4 5% | 5 6% | 5 4% | 4 2% | 5 6% | 5 5% | 5 6% | 4 7% | 4 2% | 4 4% | 5 6% | 5 2% | 5 3% | 4 4% | 4 3% | 4 4% | 4 3% | 5 8% | 4 6% | 4 3% | - // | 10 | 6 5% | 7 6% | 7 3% | 6 3% | 7 5% | 7 5% | 7 5% | 6 7% | 6 3% | 6 5% | 7 6% | 7 3% | 7 3% | 6 4% | 6 4% | 6 4% | 6 2% | 7 8% | 6 5% | 6 4% | - // | 12 | 8 5% | 9 5% | 9 5% | 8 3% | 9 5% | 9 4% | 9 5% | 8 6% | 8 4% | 8 4% | 9 7% | 9 3% | 9 3% | 8 4% | 8 4% | 8 4% | 8 3% | 9 7% | 8 5% | 8 5% | - // | 14 | 10 4% | 11 6% | 11 4% | 10 3% | 11 5% | 11 4% | 11 6% | 10 6% | 10 5% | 10 4% | 11 6% | 11 4% | 11 3% | 10 4% | 10 3% | 10 4% | 10 3% | 11 7% | 10 5% | 10 4% | - // | 16 | 12 4% | 13 6% | 13 4% | 12 3% | 13 5% | 13 3% | 13 5% | 12 6% | 12 5% | 12 5% | 13 6% | 13 4% | 13 3% | 12 4% | 12 4% | 12 4% | 12 3% | 13 6% | 12 4% | 12 5% | - // | 18 | 14 4% | 15 6% | 15 4% | 14 3% | 15 5% | 15 4% | 15 5% | 14 6% | 14 5% | 14 4% | 15 6% | 15 4% | 15 3% | 14 5% | 14 4% | 14 5% | 14 4% | 15 6% | 14 4% | 14 4% | - // | 20 | 16 4% | 17 6% | 17 4% | 16 4% | 17 5% | 17 3% | 17 5% | 16 5% | 16 6% | 16 4% | 17 5% | 17 5% | 17 3% | 16 5% | 16 4% | 16 5% | 16 4% | 17 6% | 16 4% | 16 4% | - // | 22 | 18 4% | 19 6% | 19 4% | 18 4% | 19 5% | 19 4% | 19 5% | 18 5% | 18 5% | 18 4% | 19 5% | 19 5% | 19 3% | 18 5% | 18 4% | 18 4% | 18 4% | 19 6% | 18 4% | 18 4% | - // | 24 | 20 5% | 21 5% | 21 4% | 20 4% | 21 5% | 21 3% | 21 4% | 20 5% | 20 5% | 20 4% | 21 6% | 21 5% | 21 3% | 20 4% | 20 4% | 20 4% | 20 4% | 21 6% | 20 5% | 20 4% | - // | 26 | 22 5% | 23 5% | 23 4% | 22 4% | 23 5% | 23 3% | 23 5% | 22 5% | 22 4% | 22 4% | 23 5% | 23 5% | 23 3% | 22 4% | 22 4% | 22 5% | 22 4% | 23 5% | 22 5% | 22 4% | - // | 28 | 24 5% | 25 5% | 25 4% | 24 4% | 25 5% | 25 3% | 25 4% | 24 5% | 24 5% | 24 4% | 25 5% | 25 5% | 25 3% | 24 4% | 24 4% | 24 5% | 24 4% | 25 5% | 24 5% | 24 4% | - // | 30 | 26 5% | 27 5% | 27 4% | 26 4% | 27 5% | 27 4% | 27 4% | 26 5% | 26 5% | 26 4% | 27 5% | 27 5% | 27 3% | 26 4% | 26 4% | 26 5% | 26 4% | 27 5% | 26 5% | 26 5% | - // | 32 | 28 5% | 29 5% | 29 4% | 28 4% | 29 5% | 29 4% | 29 4% | 28 5% | 28 4% | 28 4% | 29 5% | 29 5% | 29 3% | 28 4% | 28 4% | 28 5% | 28 4% | 29 5% | 28 5% | 28 5% | - // | 34 | 30 5% | 31 4% | 31 4% | 30 4% | 31 5% | 31 4% | 31 4% | 30 4% | 30 5% | 30 4% | 31 5% | 31 5% | 31 3% | 30 4% | 30 4% | 30 4% | 30 4% | 31 5% | 30 5% | 30 5% | - // | 36 | 32 5% | 33 5% | 33 4% | 32 4% | 33 5% | 33 4% | 33 4% | 32 5% | 32 4% | 32 4% | 33 5% | 33 5% | 33 3% | 32 4% | 32 4% | 32 4% | 32 4% | 33 5% | 32 5% | 32 5% | - // | 38 | 34 5% | 35 5% | 35 5% | 34 4% | 35 5% | 35 4% | 35 4% | 34 5% | 34 5% | 34 4% | 35 5% | 35 5% | 35 3% | 34 4% | 34 4% | 34 4% | 34 4% | 35 5% | 34 5% | 34 5% | - // | 40 | 36 5% | 37 5% | 37 5% | 36 4% | 37 5% | 37 4% | 37 4% | 36 4% | 36 5% | 36 4% | 37 5% | 37 5% | 37 3% | 36 4% | 36 4% | 36 5% | 36 4% | 37 5% | 36 5% | 36 5% | - // | 42 | 38 5% | 39 5% | 39 5% | 38 4% | 39 5% | 39 4% | 39 4% | 38 4% | 38 5% | 38 4% | 39 5% | 39 5% | 39 3% | 38 4% | 38 5% | 38 5% | 38 4% | 39 5% | 38 5% | 38 5% | - // | 44 | 40 5% | 41 5% | 41 5% | 40 4% | 41 5% | 41 4% | 41 5% | 40 4% | 40 5% | 40 4% | 41 5% | 41 5% | 41 3% | 40 4% | 40 5% | 40 5% | 40 4% | 41 5% | 40 5% | 40 5% | - // | 46 | 42 5% | 43 5% | 43 5% | 42 4% | 43 5% | 43 4% | 43 5% | 42 4% | 42 5% | 42 4% | 43 5% | 43 5% | 43 4% | 42 4% | 42 5% | 42 5% | 42 4% | 43 5% | 42 4% | 42 5% | - // | 48 | 44 5% | 45 5% | 45 5% | 44 4% | 45 5% | 45 4% | 45 5% | 44 4% | 44 5% | 44 4% | 45 5% | 45 5% | 45 4% | 44 4% | 44 4% | 44 5% | 44 4% | 45 5% | 44 4% | 44 5% | - // | 50 | 46 5% | 47 5% | 47 5% | 46 4% | 47 5% | 47 4% | 47 5% | 46 4% | 46 5% | 46 4% | 47 5% | 47 5% | 47 4% | 46 4% | 46 4% | 46 5% | 46 4% | 47 5% | 46 5% | 46 5% | - // | 52 | 48 5% | 49 5% | 49 5% | 48 4% | 49 5% | 49 4% | 49 5% | 48 4% | 48 5% | 48 4% | 49 5% | 49 5% | 49 4% | 48 4% | 48 5% | 48 5% | 48 4% | 49 5% | 48 5% | 48 5% | - // | 54 | 50 5% | 51 5% | 51 5% | 50 4% | 51 5% | 51 4% | 51 5% | 50 4% | 50 5% | 50 4% | 51 5% | 51 5% | 51 4% | 50 4% | 50 4% | 50 5% | 50 4% | 51 5% | 50 5% | 50 5% | - // | 56 | 52 5% | 53 4% | 53 5% | 52 4% | 53 5% | 53 4% | 53 5% | 52 4% | 52 5% | 52 4% | 53 5% | 53 5% | 53 4% | 52 4% | 52 4% | 52 5% | 52 4% | 53 5% | 52 4% | 52 5% | - // | 58 | 54 5% | 55 4% | 55 5% | 54 4% | 55 5% | 55 4% | 55 5% | 54 4% | 54 5% | 54 4% | 55 5% | 55 5% | 55 4% | 54 4% | 54 5% | 54 4% | 54 4% | 55 5% | 54 5% | 54 5% | - // | 60 | 56 5% | 57 4% | 57 5% | 56 4% | 57 5% | 57 4% | 57 5% | 56 4% | 56 5% | 56 4% | 57 5% | 57 5% | 57 4% | 56 4% | 56 5% | 56 4% | 56 4% | 57 5% | 56 4% | 56 5% | - // | 62 | 58 5% | 59 4% | 59 5% | 58 4% | 59 5% | 59 4% | 59 5% | 58 4% | 58 5% | 58 4% | 59 5% | 59 5% | 59 4% | 58 4% | 58 5% | 58 4% | 58 4% | 59 5% | 58 4% | 58 5% | - // | 64 | 60 5% | 61 4% | 61 5% | 60 4% | 61 5% | 61 4% | 61 5% | 60 4% | 60 5% | 60 4% | 61 5% | 61 5% | 61 4% | 60 4% | 60 5% | 60 4% | 60 4% | 61 5% | 60 4% | 60 5% | - // | 66 | 62 5% | 63 5% | 63 5% | 62 4% | 63 5% | 63 4% | 63 5% | 62 4% | 62 5% | 62 4% | 63 5% | 63 5% | 63 4% | 62 4% | 62 5% | 62 4% | 62 4% | 63 5% | 62 4% | 62 5% | - // | 68 | 64 5% | 65 5% | 65 5% | 64 4% | 65 5% | 65 4% | 65 4% | 64 4% | 64 5% | 64 4% | 65 5% | 65 5% | 65 4% | 64 4% | 64 5% | 64 4% | 64 4% | 65 5% | 64 5% | 64 5% | - // | 70 | 66 5% | 67 5% | 67 4% | 66 4% | 67 5% | 67 4% | 67 4% | 66 4% | 66 5% | 66 4% | 67 5% | 67 5% | 67 4% | 66 4% | 66 5% | 66 4% | 66 4% | 67 5% | 66 4% | 66 5% | - // | 72 | 68 5% | 69 4% | 69 4% | 68 4% | 69 5% | 69 4% | 69 4% | 68 4% | 68 5% | 68 4% | 69 5% | 69 5% | 69 4% | 68 4% | 68 5% | 68 4% | 68 4% | 69 5% | 68 4% | 68 5% | - // | 74 | 70 4% | 71 4% | 71 4% | 70 4% | 71 5% | 71 4% | 71 4% | 70 4% | 70 5% | 70 5% | 71 5% | 71 5% | 71 4% | 70 4% | 70 5% | 70 4% | 70 4% | 71 5% | 70 4% | 70 5% | - // | 76 | 72 4% | 73 5% | 73 4% | 72 4% | 73 5% | 73 4% | 73 4% | 72 4% | 72 5% | 72 5% | 73 5% | 73 5% | 73 4% | 72 4% | 72 5% | 72 4% | 72 4% | 73 5% | 72 4% | 72 5% | - // | 78 | 74 4% | 75 5% | 75 4% | 74 4% | 75 5% | 75 4% | 75 4% | 74 4% | 74 5% | 74 5% | 75 5% | 75 5% | 75 4% | 74 4% | 74 5% | 74 4% | 74 4% | 75 5% | 74 5% | 74 5% | - // | 80 | 76 5% | 77 5% | 77 4% | 76 4% | 77 5% | 77 4% | 77 4% | 76 4% | 76 5% | 76 5% | 77 5% | 77 5% | 77 4% | 76 4% | 76 5% | 76 4% | 76 4% | 77 5% | 76 5% | 76 5% | - // | 82 | 78 5% | 79 5% | 79 4% | 78 4% | 79 5% | 79 4% | 79 4% | 78 4% | 78 5% | 78 5% | 79 5% | 79 5% | 79 4% | 78 4% | 78 5% | 78 4% | 78 4% | 79 5% | 78 5% | 78 5% | - // | 84 | 80 5% | 81 4% | 81 4% | 80 4% | 81 5% | 81 4% | 81 4% | 80 4% | 80 5% | 80 5% | 81 5% | 81 5% | 81 4% | 80 4% | 80 5% | 80 4% | 80 4% | 81 5% | 80 5% | 80 5% | - // | 86 | 82 4% | 83 4% | 83 5% | 82 4% | 83 5% | 83 4% | 83 4% | 82 4% | 82 5% | 82 5% | 83 5% | 83 5% | 83 4% | 82 4% | 82 5% | 82 4% | 82 4% | 83 5% | 82 5% | 82 5% | - // | 88 | 84 4% | 85 4% | 85 5% | 84 4% | 85 5% | 85 4% | 85 4% | 84 4% | 84 5% | 84 5% | 85 5% | 85 5% | 85 4% | 84 4% | 84 5% | 84 4% | 84 4% | 85 5% | 84 4% | 84 5% | - // | 90 | 86 4% | 87 4% | 87 5% | 86 4% | 87 5% | 87 4% | 87 4% | 86 4% | 86 5% | 86 5% | 87 5% | 87 5% | 87 4% | 86 4% | 86 5% | 86 4% | 86 4% | 87 5% | 86 5% | 86 5% | - // | 92 | 88 4% | 89 4% | 89 5% | 88 4% | 89 5% | 89 4% | 89 4% | 88 4% | 88 5% | 88 5% | 89 5% | 89 5% | 89 4% | 88 4% | 88 5% | 88 4% | 88 4% | 89 5% | 88 5% | 88 5% | - // | 94 | 90 4% | 91 4% | 91 4% | 90 4% | 91 5% | 91 4% | 91 4% | 90 4% | 90 5% | 90 5% | 91 5% | 91 5% | 91 4% | 90 4% | 90 5% | 90 4% | 90 4% | 91 5% | 90 5% | 90 4% | - // | 96 | 92 4% | 93 5% | 93 5% | 92 4% | 93 5% | 93 4% | 93 4% | 92 4% | 92 5% | 92 5% | 93 5% | 93 5% | 93 4% | 92 4% | 92 5% | 92 4% | 92 4% | 93 5% | 92 5% | 92 4% | - // | 98 | 94 4% | 95 5% | 95 4% | 94 4% | 95 5% | 95 4% | 95 4% | 94 4% | 94 5% | 94 5% | 95 5% | 95 5% | 95 4% | 94 4% | 94 5% | 94 4% | 94 4% | 95 5% | 94 5% | 94 4% | + // | 2 | 2 28% | 0 0% | 0 0% | 1 19% | 0 0% | 1 5% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 1 18% | 0 0% | 1 12% | 0 0% | 0 0% | 0 0% | 1 14% | 0 0% | + // | 4 | 2 15% | 1 4% | 0 0% | 2 5% | 0 0% | 3 4% | 1 1% | 1 5% | 2 6% | 0 0% | 1 12% | 0 0% | 2 3% | 0 0% | 2 11% | 1 0% | 2 2% | 0 0% | 1 17% | 1 7% | + // | 6 | 3 8% | 3 7% | 3 4% | 3 7% | 3 5% | 4 3% | 3 0% | 3 6% | 3 7% | 0 0% | 3 5% | 3 7% | 3 4% | 0 0% | 3 3% | 3 6% | 3 2% | 2 1% | 3 8% | 3 6% | + // | 8 | 5 6% | 5 7% | 5 6% | 5 6% | 5 5% | 5 3% | 5 2% | 5 5% | 5 6% | 0 0% | 5 5% | 5 6% | 5 6% | 2 1% | 5 5% | 5 4% | 5 3% | 4 3% | 5 6% | 5 5% | + // | 10 | 6 5% | 6 5% | 6 5% | 6 5% | 6 5% | 6 2% | 6 2% | 6 4% | 6 5% | 17 16% | 6 3% | 6 6% | 6 5% | 5 2% | 6 5% | 6 4% | 6 1% | 6 4% | 6 5% | 6 4% | + // | 12 | 8 5% | 8 5% | 8 5% | 8 5% | 8 5% | 8 3% | 8 2% | 8 5% | 8 4% | 17 13% | 8 3% | 8 6% | 9 5% | 8 2% | 8 5% | 8 4% | 8 2% | 8 4% | 8 6% | 8 3% | + // | 14 | 10 4% | 10 4% | 10 5% | 10 5% | 11 6% | 11 4% | 10 3% | 10 5% | 10 4% | 17 9% | 10 3% | 10 5% | 11 5% | 10 2% | 10 5% | 10 4% | 10 2% | 10 5% | 10 5% | 10 3% | + // | 16 | 13 5% | 12 4% | 12 5% | 12 5% | 13 7% | 13 4% | 12 3% | 12 5% | 12 4% | 17 8% | 12 3% | 12 6% | 13 5% | 12 3% | 12 5% | 13 4% | 12 3% | 12 4% | 12 5% | 12 3% | + // | 18 | 15 4% | 14 4% | 14 5% | 14 5% | 15 6% | 15 4% | 15 3% | 14 5% | 14 4% | 17 7% | 14 3% | 14 6% | 15 5% | 14 3% | 14 5% | 15 4% | 14 3% | 14 4% | 15 5% | 14 3% | + // | 20 | 17 5% | 16 4% | 16 5% | 16 5% | 17 5% | 17 4% | 17 4% | 16 5% | 17 4% | 18 6% | 16 3% | 16 6% | 17 5% | 16 3% | 16 5% | 17 4% | 16 3% | 16 4% | 17 5% | 16 3% | + // | 22 | 19 5% | 18 4% | 18 5% | 18 5% | 19 6% | 19 4% | 19 4% | 18 5% | 19 3% | 20 6% | 18 4% | 18 6% | 19 5% | 18 3% | 18 5% | 19 4% | 18 3% | 18 4% | 19 5% | 18 4% | + // | 24 | 21 5% | 20 4% | 20 4% | 20 5% | 21 5% | 21 4% | 21 4% | 20 5% | 21 3% | 22 5% | 20 4% | 20 6% | 21 5% | 20 3% | 20 5% | 21 5% | 20 4% | 20 4% | 21 5% | 20 4% | + // | 26 | 23 5% | 22 4% | 22 4% | 22 5% | 23 5% | 23 4% | 23 4% | 22 5% | 23 3% | 24 5% | 22 4% | 22 6% | 23 5% | 22 4% | 22 5% | 23 5% | 22 4% | 22 4% | 23 5% | 22 4% | + // | 28 | 25 5% | 24 4% | 24 4% | 24 4% | 25 5% | 25 4% | 25 4% | 24 5% | 25 3% | 26 5% | 24 4% | 24 5% | 25 5% | 24 4% | 24 5% | 25 5% | 24 4% | 24 4% | 25 5% | 24 4% | + // | 30 | 27 5% | 26 4% | 26 4% | 26 5% | 27 5% | 27 4% | 27 4% | 26 5% | 27 3% | 28 5% | 26 4% | 26 5% | 27 5% | 26 4% | 26 5% | 27 5% | 26 4% | 26 4% | 27 5% | 26 4% | + // | 32 | 29 5% | 28 4% | 28 4% | 28 5% | 29 5% | 29 4% | 29 4% | 28 5% | 29 3% | 30 5% | 28 4% | 28 5% | 29 5% | 28 4% | 28 5% | 29 5% | 28 4% | 28 4% | 29 5% | 28 4% | + // | 34 | 31 5% | 30 4% | 30 4% | 30 4% | 31 5% | 31 4% | 31 4% | 30 5% | 31 3% | 32 5% | 30 4% | 30 5% | 31 5% | 30 4% | 30 5% | 31 5% | 30 4% | 30 4% | 31 5% | 30 4% | + // | 36 | 33 5% | 32 4% | 32 5% | 32 4% | 33 5% | 33 4% | 33 4% | 32 5% | 33 4% | 34 5% | 32 5% | 32 5% | 33 5% | 32 4% | 32 5% | 33 5% | 32 4% | 32 4% | 33 5% | 32 4% | + // | 38 | 35 5% | 34 4% | 34 5% | 34 4% | 35 5% | 35 4% | 35 4% | 34 5% | 35 4% | 36 5% | 34 5% | 34 5% | 35 5% | 34 4% | 34 5% | 35 5% | 34 4% | 34 4% | 35 5% | 34 4% | + // | 40 | 37 4% | 36 4% | 36 5% | 36 4% | 37 5% | 37 4% | 37 4% | 36 5% | 37 4% | 38 5% | 36 5% | 36 5% | 37 5% | 36 4% | 36 4% | 37 5% | 36 4% | 36 4% | 37 5% | 36 4% | + // | 42 | 39 5% | 38 4% | 38 5% | 38 4% | 39 5% | 39 4% | 39 4% | 38 5% | 39 4% | 40 5% | 38 4% | 38 5% | 39 5% | 38 4% | 38 4% | 39 5% | 38 4% | 38 4% | 39 5% | 38 4% | + // | 44 | 41 4% | 40 4% | 40 5% | 40 4% | 41 5% | 41 4% | 41 4% | 40 5% | 41 4% | 42 5% | 40 4% | 40 6% | 41 5% | 40 4% | 40 4% | 41 5% | 40 4% | 40 4% | 41 5% | 40 4% | + // | 46 | 43 5% | 42 4% | 42 5% | 42 4% | 43 5% | 43 4% | 43 4% | 42 5% | 43 4% | 44 5% | 42 5% | 42 5% | 43 5% | 42 4% | 42 5% | 43 5% | 42 4% | 42 4% | 43 5% | 42 4% | + // | 48 | 45 5% | 44 4% | 44 5% | 44 4% | 45 5% | 45 4% | 45 4% | 44 5% | 45 4% | 46 5% | 44 4% | 44 5% | 45 5% | 44 4% | 44 5% | 45 5% | 44 4% | 44 4% | 45 5% | 44 4% | + // | 50 | 47 5% | 46 4% | 46 5% | 46 4% | 47 5% | 47 4% | 47 4% | 46 5% | 47 4% | 48 5% | 46 4% | 46 5% | 47 5% | 46 4% | 46 4% | 47 5% | 46 4% | 46 4% | 47 5% | 46 4% | + // | 52 | 49 5% | 48 4% | 48 5% | 48 4% | 49 5% | 49 4% | 49 4% | 48 5% | 49 4% | 50 5% | 48 5% | 48 5% | 49 5% | 48 4% | 48 4% | 49 5% | 48 4% | 48 4% | 49 5% | 48 4% | + // | 54 | 51 5% | 50 4% | 50 5% | 50 4% | 51 5% | 51 4% | 51 4% | 50 5% | 51 4% | 52 5% | 50 5% | 50 5% | 51 5% | 50 4% | 50 4% | 51 5% | 50 4% | 50 4% | 51 5% | 50 4% | + // | 56 | 53 5% | 52 4% | 52 5% | 52 4% | 53 5% | 53 4% | 53 4% | 52 5% | 53 4% | 54 5% | 52 5% | 52 5% | 53 5% | 52 4% | 52 4% | 53 5% | 52 4% | 52 4% | 53 5% | 52 4% | + // | 58 | 55 5% | 54 4% | 54 5% | 54 4% | 55 5% | 55 4% | 55 4% | 54 5% | 55 4% | 56 5% | 54 5% | 54 5% | 55 5% | 54 4% | 54 4% | 55 5% | 54 4% | 54 4% | 55 5% | 54 4% | + // | 60 | 57 5% | 56 4% | 56 5% | 56 4% | 57 5% | 57 4% | 57 4% | 56 5% | 57 4% | 58 5% | 56 5% | 56 5% | 57 5% | 56 4% | 56 4% | 57 5% | 56 4% | 56 4% | 57 5% | 56 4% | + // | 62 | 59 5% | 58 4% | 58 5% | 58 4% | 59 5% | 59 4% | 59 4% | 58 5% | 59 4% | 60 5% | 58 5% | 58 5% | 59 5% | 58 4% | 58 4% | 59 5% | 58 4% | 58 4% | 59 5% | 58 4% | + // | 64 | 61 5% | 60 4% | 60 4% | 60 4% | 61 5% | 61 4% | 61 4% | 60 5% | 61 4% | 62 5% | 60 5% | 60 5% | 61 5% | 60 4% | 60 4% | 61 5% | 60 4% | 60 4% | 61 5% | 60 4% | + // | 66 | 63 5% | 62 4% | 62 4% | 62 4% | 63 5% | 63 4% | 63 4% | 62 5% | 63 4% | 64 5% | 62 5% | 62 5% | 63 5% | 62 4% | 62 4% | 63 5% | 62 4% | 62 4% | 63 5% | 62 4% | + // | 68 | 65 5% | 64 4% | 64 4% | 64 4% | 65 5% | 65 5% | 65 4% | 64 5% | 65 4% | 66 5% | 64 5% | 64 5% | 65 5% | 64 4% | 64 4% | 65 5% | 64 4% | 64 4% | 65 5% | 64 4% | + // | 70 | 67 5% | 66 4% | 66 4% | 66 4% | 67 5% | 67 5% | 67 4% | 66 5% | 67 4% | 68 5% | 66 5% | 66 5% | 67 5% | 66 4% | 66 4% | 67 5% | 66 4% | 66 4% | 67 5% | 66 4% | + // | 72 | 69 5% | 68 4% | 68 4% | 68 4% | 69 5% | 69 5% | 69 4% | 68 5% | 69 4% | 70 5% | 68 5% | 68 5% | 69 5% | 68 4% | 68 5% | 69 5% | 68 4% | 68 4% | 69 5% | 68 4% | + // | 74 | 71 5% | 70 4% | 70 4% | 70 4% | 71 5% | 71 5% | 71 4% | 70 5% | 71 4% | 72 5% | 70 5% | 70 5% | 71 5% | 70 4% | 70 5% | 71 5% | 70 4% | 70 4% | 71 5% | 70 4% | + // | 76 | 73 5% | 72 4% | 72 4% | 72 4% | 73 5% | 73 5% | 73 4% | 72 5% | 73 4% | 74 5% | 72 5% | 72 5% | 73 5% | 72 4% | 72 5% | 73 5% | 72 4% | 72 4% | 73 5% | 72 4% | + // | 78 | 75 5% | 74 4% | 74 4% | 74 4% | 75 5% | 75 5% | 75 4% | 74 5% | 75 4% | 76 5% | 74 5% | 74 5% | 75 5% | 74 4% | 74 5% | 75 5% | 74 4% | 74 4% | 75 5% | 74 4% | + // | 80 | 77 5% | 76 4% | 76 4% | 76 4% | 77 5% | 77 5% | 77 4% | 76 5% | 77 4% | 78 5% | 76 5% | 76 5% | 77 5% | 76 4% | 76 5% | 77 5% | 76 4% | 76 4% | 77 5% | 76 4% | + // | 82 | 79 5% | 78 4% | 78 4% | 78 4% | 79 5% | 79 5% | 79 4% | 78 5% | 79 4% | 80 5% | 78 5% | 78 5% | 79 5% | 78 4% | 78 4% | 79 5% | 78 4% | 78 4% | 79 5% | 78 4% | + // | 84 | 81 5% | 80 4% | 80 4% | 80 4% | 81 5% | 81 5% | 81 4% | 80 5% | 81 4% | 82 5% | 80 5% | 80 5% | 81 5% | 80 4% | 80 5% | 81 5% | 80 4% | 80 4% | 81 5% | 80 4% | + // | 86 | 83 5% | 82 4% | 82 4% | 82 4% | 83 5% | 83 5% | 83 4% | 82 5% | 83 4% | 84 5% | 82 5% | 82 5% | 83 5% | 82 4% | 82 5% | 83 5% | 82 4% | 82 4% | 83 5% | 82 4% | + // | 88 | 85 5% | 84 4% | 84 4% | 84 4% | 85 5% | 85 5% | 85 4% | 84 5% | 85 4% | 86 5% | 84 5% | 84 5% | 85 5% | 84 4% | 84 5% | 85 5% | 84 4% | 84 4% | 85 5% | 84 4% | + // | 90 | 87 5% | 86 4% | 86 4% | 86 4% | 87 5% | 87 5% | 87 4% | 86 5% | 87 4% | 88 5% | 86 5% | 86 5% | 87 5% | 86 4% | 86 5% | 87 4% | 86 4% | 86 4% | 87 5% | 86 4% | + // | 92 | 89 5% | 88 4% | 88 4% | 88 4% | 89 5% | 89 5% | 89 4% | 88 5% | 89 4% | 90 5% | 88 5% | 88 5% | 89 5% | 88 4% | 88 5% | 89 4% | 88 4% | 88 4% | 89 5% | 88 4% | + // | 94 | 91 5% | 90 4% | 90 4% | 90 4% | 91 5% | 91 5% | 91 4% | 90 5% | 91 4% | 92 5% | 90 5% | 90 5% | 91 5% | 90 4% | 90 5% | 91 4% | 90 4% | 90 4% | 91 5% | 90 4% | + // | 96 | 93 5% | 92 4% | 92 4% | 92 4% | 93 5% | 93 5% | 93 4% | 92 5% | 93 4% | 94 5% | 92 5% | 92 5% | 93 5% | 92 4% | 92 5% | 93 4% | 92 4% | 92 4% | 93 5% | 92 4% | + // | 98 | 95 5% | 94 4% | 94 4% | 94 5% | 95 5% | 95 5% | 95 4% | 94 5% | 95 4% | 96 5% | 94 5% | 94 5% | 95 5% | 94 4% | 94 5% | 95 4% | 94 4% | 94 4% | 95 5% | 94 4% | // +-----+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+ - // Total bytes=998843773, ranges=1909 + // Total bytes=997754756, ranges=1910 } diff --git a/pkg/storage/rule_solver.go b/pkg/storage/rule_solver.go index 4e915d5ca12e..3104fd55161d 100644 --- a/pkg/storage/rule_solver.go +++ b/pkg/storage/rule_solver.go @@ -38,7 +38,8 @@ func (c candidate) String() string { c.store.StoreID, c.valid, c.constraint, c.balance) } -// less first compares constraint scores, then balance scores. +// less first compares valid, then constraint scores, then balance +// scores. func (c candidate) less(o candidate) bool { if !o.valid { return false @@ -67,6 +68,31 @@ func (cl candidateList) String() string { return buffer.String() } +// byScore implements sort.Interface to sort by scores. +type byScore candidateList + +var _ sort.Interface = byScore(nil) + +func (c byScore) Len() int { return len(c) } +func (c byScore) Less(i, j int) bool { return c[i].less(c[j]) } +func (c byScore) Swap(i, j int) { c[i], c[j] = c[j], c[i] } + +// byScoreAndID implements sort.Interface to sort by scores and ids. +type byScoreAndID candidateList + +var _ sort.Interface = byScoreAndID(nil) + +func (c byScoreAndID) Len() int { return len(c) } +func (c byScoreAndID) Less(i, j int) bool { + if c[i].constraint == c[j].constraint && + c[i].balance == c[j].balance && + c[i].valid == c[j].valid { + return c[i].store.StoreID < c[j].store.StoreID + } + return c[i].less(c[j]) +} +func (c byScoreAndID) Swap(i, j int) { c[i], c[j] = c[j], c[i] } + // 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-- { @@ -163,124 +189,170 @@ func (cl candidateList) selectBad(randGen allocatorRand) candidate { return worst } -// solveState is used to pass solution state information into a rule. -type solveState struct { - constraints config.Constraints - sl StoreList - existing []roachpb.ReplicaDescriptor - existingNodeLocalities map[roachpb.NodeID]roachpb.Locality -} - -// rule is a function that given a solveState will score and possibly -// disqualify a store. The store in solveState can be disqualified by -// returning false. Unless disqualified, the higher the returned score, the -// more likely the store will be picked as a candidate. -type rule func(roachpb.StoreDescriptor, solveState) (bool, float64, float64) - -// ruleSolver is used to test a collection of rules against stores. -type ruleSolver []rule - -// allocateRuleSolver is the set of rules used for adding new replicas. -var allocateRuleSolver = ruleSolver{ - ruleReplicasUniqueNodes, - ruleConstraints, - ruleCapacityMax, - ruleDiversity, - ruleCapacityToMean, - ruleCapacity, -} - -// removeRuleSolver is the set of rules used for removing existing replicas. -var removeRuleSolver = ruleSolver{ - ruleConstraints, - ruleCapacityMax, - ruleDiversity, - ruleCapacityFromMean, - ruleCapacity, -} - -var rebalanceExisting = ruleSolver{ - ruleConstraints, - ruleCapacityMax, - ruleDiversity, - ruleCapacityToMean, - ruleCapacity, -} - -var rebalance = ruleSolver{ - ruleConstraints, - ruleCapacityMax, - ruleDiversity, - ruleCapacityFromMean, - ruleCapacity, -} - -// Solve runs the rules against the stores in the store list and returns all -// candidate stores and their scores ordered from best to worst score. -func (rs ruleSolver) Solve( +// allocateCandidates creates a candidate list of all stores that can used for +// allocating a new replica ordered from the best to the worst. Only stores +// that meet the criteria are included in the list. +func allocateCandidates( sl StoreList, - c config.Constraints, + constraints config.Constraints, existing []roachpb.ReplicaDescriptor, existingNodeLocalities map[roachpb.NodeID]roachpb.Locality, deterministic bool, ) candidateList { - candidates := make(candidateList, len(sl.stores), len(sl.stores)) - state := solveState{ - constraints: c, - sl: sl, - existing: existing, - existingNodeLocalities: existingNodeLocalities, - } + var candidates candidateList + for _, s := range sl.stores { + if !preexistingReplicaCheck(s.Node.NodeID, existing) { + continue + } + constraintsOk, preferredMatched := constraintCheck(s, constraints) + if !constraintsOk { + continue + } + if !maxCapacityCheck(s) { + continue + } - for i, store := range sl.stores { - candidates[i] = rs.computeCandidate(store, state) + constraintScore := diversityScore(s, existingNodeLocalities) + float64(preferredMatched) + balanceScore := capacityScore(s) + candidates = append(candidates, candidate{ + store: s, + valid: true, + constraint: constraintScore, + balance: balanceScore, + }) } if deterministic { sort.Sort(sort.Reverse(byScoreAndID(candidates))) } else { sort.Sort(sort.Reverse(byScore(candidates))) } + return candidates +} +// removeCandidates creates a candidate list of all existing replicas' stores +// ordered from least qualified for removal to most qualified. Stores that are +// marked as not valid, are in violation of a criteria. +func removeCandidates( + sl StoreList, + constraints config.Constraints, + existingNodeLocalities map[roachpb.NodeID]roachpb.Locality, + deterministic bool, +) candidateList { + var candidates candidateList + for _, s := range sl.stores { + constraintsOk, preferredMatched := constraintCheck(s, constraints) + if !constraintsOk { + candidates = append(candidates, candidate{store: s, valid: false}) + continue + } + if !maxCapacityCheck(s) { + candidates = append(candidates, candidate{store: s, valid: false}) + continue + } + constraintScore := diversityRemovalScore(s.Node.NodeID, existingNodeLocalities) + float64(preferredMatched) + if rebalanceToConvergesOnMean(sl, s) { + constraintScore++ + } + balanceScore := capacityScore(s) + candidates = append(candidates, candidate{ + store: s, + valid: true, + constraint: constraintScore, + balance: balanceScore, + }) + } + if deterministic { + sort.Sort(sort.Reverse(byScoreAndID(candidates))) + } else { + sort.Sort(sort.Reverse(byScore(candidates))) + } 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 { - var totalConstraintScore, totalBalanceScore float64 - for _, rule := range rs { - valid, constraintScore, balanceScore := rule(store, state) - if !valid { - return candidate{ - store: store, - valid: false, +// rebalanceCandidates creates two candidate list. The first contains all +// existing replica's stores, order from least qualified for rebalancing to +// most qualified. The second list is of all potential stores that could be +// used as rebalancing receivers, ordered from best to worst. +func rebalanceCandidates( + sl StoreList, + constraints config.Constraints, + existing []roachpb.ReplicaDescriptor, + existingNodeLocalities map[roachpb.NodeID]roachpb.Locality, + deterministic bool, +) (candidateList, candidateList) { + // Load the exiting storesIDs into a map so to eliminate having to loop + // through the existing descriptors more than once. + existingStoreIDs := make(map[roachpb.StoreID]struct{}) + for _, repl := range existing { + existingStoreIDs[repl.StoreID] = struct{}{} + } + + var existingCandidates candidateList + var candidates candidateList + for _, s := range sl.stores { + constraintsOk, preferredMatched := constraintCheck(s, constraints) + maxCapacityOK := maxCapacityCheck(s) + if _, ok := existingStoreIDs[s.StoreID]; ok { + if !constraintsOk { + existingCandidates = append(existingCandidates, candidate{store: s, valid: false}) + continue + } + if !maxCapacityOK { + existingCandidates = append(existingCandidates, candidate{store: s, valid: false}) + continue + } + constraintScore := diversityRemovalScore(s.Node.NodeID, existingNodeLocalities) + float64(preferredMatched) + if rebalanceToConvergesOnMean(sl, s) { + constraintScore++ } + balanceScore := capacityScore(s) + existingCandidates = append(existingCandidates, candidate{ + store: s, + valid: true, + constraint: constraintScore, + balance: balanceScore, + }) + } else { + if !constraintsOk || !maxCapacityOK { + continue + } + constraintScore := diversityScore(s, existingNodeLocalities) + float64(preferredMatched) + if rebalanceToConvergesOnMean(sl, s) { + constraintScore++ + } + balanceScore := capacityScore(s) + candidates = append(candidates, candidate{ + store: s, + valid: true, + constraint: constraintScore, + balance: balanceScore, + }) } - totalConstraintScore += constraintScore - totalBalanceScore += balanceScore } - return candidate{ - store: store, - valid: true, - constraint: totalConstraintScore, - balance: totalBalanceScore, + + if deterministic { + sort.Sort(sort.Reverse(byScoreAndID(existingCandidates))) + sort.Sort(sort.Reverse(byScoreAndID(candidates))) + } else { + sort.Sort(sort.Reverse(byScore(existingCandidates))) + sort.Sort(sort.Reverse(byScore(candidates))) } + + return existingCandidates, candidates } -// ruleReplicasUniqueNodes returns true iff no existing replica is present on -// the candidate's node. All other scores are always 0. -func ruleReplicasUniqueNodes( - store roachpb.StoreDescriptor, state solveState, -) (bool, float64, float64) { - for _, r := range state.existing { - if r.NodeID == store.Node.NodeID { - return false, 0, 0 +// preexistingReplicaCheck returns true if no existing replica is present on +// the candidate's node. +func preexistingReplicaCheck(nodeID roachpb.NodeID, existing []roachpb.ReplicaDescriptor) bool { + for _, r := range existing { + if r.NodeID == nodeID { + return false } } - return true, 0, 0 + return true } -// storeHasConstraint returns whether a store descriptor attributes or locality +// storeHasConstraint returns whether a store's attributes or node's locality // matches the key value pair in the constraint. func storeHasConstraint(store roachpb.StoreDescriptor, c config.Constraint) bool { if c.Key == "" { @@ -301,112 +373,73 @@ func storeHasConstraint(store roachpb.StoreDescriptor, c config.Constraint) bool return false } -// ruleConstraints returns true iff all required and prohibited constraints are +// constraintCheck returns true iff all required and prohibited constraints are // satisfied. Stores with attributes or localities that match the most positive // constraints return higher scores. -func ruleConstraints(store roachpb.StoreDescriptor, state solveState) (bool, float64, float64) { - if len(state.constraints.Constraints) == 0 { - return true, 0, 0 +func constraintCheck(store roachpb.StoreDescriptor, constraints config.Constraints) (bool, int) { + if len(constraints.Constraints) == 0 { + return true, 0 } - matched := 0 - for _, c := range state.constraints.Constraints { - hasConstraint := storeHasConstraint(store, c) + positive := 0 + for _, constraint := range constraints.Constraints { + hasConstraint := storeHasConstraint(store, constraint) switch { - case c.Type == config.Constraint_REQUIRED && !hasConstraint: - return false, 0, 0 - case c.Type == config.Constraint_PROHIBITED && hasConstraint: - return false, 0, 0 - case (c.Type == config.Constraint_POSITIVE && hasConstraint) || - (c.Type == config.Constraint_REQUIRED && hasConstraint) || - (c.Type == config.Constraint_PROHIBITED && !hasConstraint): - matched++ + case constraint.Type == config.Constraint_REQUIRED && !hasConstraint: + return false, 0 + case constraint.Type == config.Constraint_PROHIBITED && hasConstraint: + return false, 0 + case (constraint.Type == config.Constraint_POSITIVE && hasConstraint): + positive++ } } - - return true, float64(matched) / float64(len(state.constraints.Constraints)), 0 + return true, positive } -// ruleDiversity returns higher scores for stores with the fewest locality tiers -// in common with already existing replicas. It always returns true. -func ruleDiversity(store roachpb.StoreDescriptor, state solveState) (bool, float64, float64) { +// diversityScore returns a score between 1 and 0 where higher scores are stores +// with the fewest locality tiers in common with already existing replicas. +func diversityScore( + store roachpb.StoreDescriptor, existingNodeLocalities map[roachpb.NodeID]roachpb.Locality, +) float64 { minScore := 1.0 - for _, locality := range state.existingNodeLocalities { + for _, locality := range existingNodeLocalities { if newScore := store.Node.Locality.DiversityScore(locality); newScore < minScore { minScore = newScore } } - return true, minScore, 0 + return minScore } -/* -// ruleDiversityExisting returns higher scores for stores with the fewest locality tiers -// in common with already existing replicas. It always returns true. -func ruleDiversity(store roachpb.StoreDescriptor, state solveState) (bool, float64, float64) { - minScore := 1.0 - for _, locality := range state.existingNodeLocalities { - if newScore := store.Node.Locality.DiversityScore(locality); newScore < minScore { - minScore = newScore +// diversityRemovalScore is similar to diversityScore but instead of calculating +// the score if a new node is added, it calculates the remaining diversity if a +// node is removed. +func diversityRemovalScore( + nodeID roachpb.NodeID, existingNodeLocalities map[roachpb.NodeID]roachpb.Locality, +) float64 { + var maxScore float64 + for nodeIDx, localityX := range existingNodeLocalities { + if nodeIDx == nodeID { + continue + } + for nodeIDy, localityY := range existingNodeLocalities { + if nodeIDy == nodeID || nodeIDx >= nodeIDy { + continue + } + if newScore := localityX.DiversityScore(localityY); newScore > maxScore { + maxScore = newScore + } } } - return true, minScore, 0 -} -*/ - -// ruleCapacity returns a balance score that is inversely proportional to the -// number of ranges on the candidate store such that the most empty store will -// have the highest scores. Scores are always between 0 and 1. -func ruleCapacity(store roachpb.StoreDescriptor, state solveState) (bool, float64, float64) { - return true, 0, 1 / float64(store.Capacity.RangeCount+1) -} - -// ruleCapacityMax ensures that we don't try to overfill a store. -func ruleCapacityMax(store roachpb.StoreDescriptor, state solveState) (bool, float64, float64) { - if store.Capacity.FractionUsed() > maxFractionUsedThreshold { - return false, 0, 0 - } - return true, 0, 0 -} - -// ruleCapacityFromMean is designed for removals and yields a lower constraint -// score if the removal of this store would push the store closer away from -// the mean number of ranges. -func ruleCapacityFromMean(store roachpb.StoreDescriptor, state solveState) (bool, float64, float64) { - if rebalanceFromConvergesOnMean(state.sl, store) { - return true, 0, 0 - } - return true, 0.1, 0 + return maxScore } -// ruleCapacityToMean is designed for rebalancing and yields a higher constraint -// score if the addition of this store would push the store closer to the mean -// number of ranges. -func ruleCapacityToMean(store roachpb.StoreDescriptor, state solveState) (bool, float64, float64) { - if rebalanceToConvergesOnMean(state.sl, store) { - return true, 0, 0 - } - return true, 0.1, 0 +// capacityScore returns a score between 0 and 1 that is inversely proportional +// to the number of ranges on the store such that the most empty store will have +// the highest scores. +func capacityScore(store roachpb.StoreDescriptor) float64 { + return 1.0 / float64(store.Capacity.RangeCount+1) } -// byScore implements sort.Interface to sort by scores. -type byScore candidateList - -var _ sort.Interface = byScore(nil) - -func (c byScore) Len() int { return len(c) } -func (c byScore) Less(i, j int) bool { return c[i].less(c[j]) } -func (c byScore) Swap(i, j int) { c[i], c[j] = c[j], c[i] } - -type byScoreAndID candidateList - -var _ sort.Interface = byScoreAndID(nil) - -func (c byScoreAndID) Len() int { return len(c) } -func (c byScoreAndID) Less(i, j int) bool { - if c[i].constraint == c[j].constraint && - c[i].balance == c[j].balance && - c[i].valid == c[j].valid { - return c[i].store.StoreID < c[j].store.StoreID - } - return c[i].less(c[j]) +// maxCapacityCheck returns true if the store has room for a new replica. +func maxCapacityCheck(store roachpb.StoreDescriptor) bool { + return store.Capacity.FractionUsed() < maxFractionUsedThreshold } -func (c byScoreAndID) Swap(i, j int) { c[i], c[j] = c[j], c[i] } diff --git a/pkg/storage/rule_solver_test.go b/pkg/storage/rule_solver_test.go index d191a7150eaa..3cd9081bd085 100644 --- a/pkg/storage/rule_solver_test.go +++ b/pkg/storage/rule_solver_test.go @@ -31,268 +31,6 @@ import ( "github.com/cockroachdb/cockroach/pkg/util/leaktest" ) -// TODO(bram): This test suite is not even close to exhaustive. The scores are -// not checked and each rule should have many more test cases. Also add a -// corrupt replica test and remove the 0 range ID used when calling -// getStoreList. -func TestRuleSolver(t *testing.T) { - defer leaktest.AfterTest(t)() - - stopper, _, _, storePool := createTestStorePool( - TestTimeUntilStoreDeadOff, - /* deterministic */ true, - ) - defer stopper.Stop() - - storeUSa15 := roachpb.StoreID(1) // us-a-1-5 - storeUSa1 := roachpb.StoreID(2) // us-a-1 - storeUSb := roachpb.StoreID(3) // us-b - storeDead := roachpb.StoreID(4) - storeEurope := roachpb.StoreID(5) // eur-a-1-5 - - mockStorePool(storePool, []roachpb.StoreID{storeUSa15, storeUSa1, storeUSb, storeEurope}, []roachpb.StoreID{storeDead}, nil) - - // tierSetup returns a tier struct constructed using the passed in values. - // If any value is an empty string, it is not included. - tierSetup := func(datacenter, floor, rack, slot string) []roachpb.Tier { - var tiers []roachpb.Tier - if datacenter != "" { - tiers = append(tiers, roachpb.Tier{Key: "datacenter", Value: datacenter}) - } - if floor != "" { - tiers = append(tiers, roachpb.Tier{Key: "floor", Value: floor}) - } - if rack != "" { - tiers = append(tiers, roachpb.Tier{Key: "rack", Value: rack}) - } - if slot != "" { - tiers = append(tiers, roachpb.Tier{Key: "slot", Value: slot}) - } - return tiers - } - - // capacitySetup returns a store capacity in which the total capacity is - // always 100 and available and range count are passed in. - capacitySetup := func(available int64, rangeCount int32) roachpb.StoreCapacity { - return roachpb.StoreCapacity{ - Capacity: 100, - Available: available, - RangeCount: rangeCount, - } - } - - storePool.mu.Lock() - - storePool.mu.storeDetails[storeUSa15].desc.Attrs.Attrs = []string{"a"} - storePool.mu.storeDetails[storeUSa15].desc.Node.Locality.Tiers = tierSetup("us", "a", "1", "5") - storePool.mu.storeDetails[storeUSa15].desc.Capacity = capacitySetup(1, 99) - storePool.mu.nodeLocalities[roachpb.NodeID(storeUSa15)] = storePool.mu.storeDetails[storeUSa15].desc.Node.Locality - - storePool.mu.storeDetails[storeUSa1].desc.Attrs.Attrs = []string{"a", "b"} - storePool.mu.storeDetails[storeUSa1].desc.Node.Locality.Tiers = tierSetup("us", "a", "1", "") - storePool.mu.storeDetails[storeUSa1].desc.Capacity = capacitySetup(100, 0) - storePool.mu.nodeLocalities[roachpb.NodeID(storeUSa1)] = storePool.mu.storeDetails[storeUSa1].desc.Node.Locality - - storePool.mu.storeDetails[storeUSb].desc.Attrs.Attrs = []string{"a", "b", "c"} - storePool.mu.storeDetails[storeUSb].desc.Node.Locality.Tiers = tierSetup("us", "b", "", "") - storePool.mu.storeDetails[storeUSb].desc.Capacity = capacitySetup(50, 50) - storePool.mu.nodeLocalities[roachpb.NodeID(storeUSb)] = storePool.mu.storeDetails[storeUSb].desc.Node.Locality - - storePool.mu.storeDetails[storeEurope].desc.Node.Locality.Tiers = tierSetup("eur", "a", "1", "5") - storePool.mu.storeDetails[storeEurope].desc.Capacity = capacitySetup(60, 40) - storePool.mu.nodeLocalities[roachpb.NodeID(storeEurope)] = storePool.mu.storeDetails[storeEurope].desc.Node.Locality - - storePool.mu.Unlock() - - testCases := []struct { - name string - rule rule - c config.Constraints - existing []roachpb.ReplicaDescriptor - expectedValid []roachpb.StoreID - expectedInvalid []roachpb.StoreID - }{ - { - name: "no constraints or rules", - expectedValid: []roachpb.StoreID{storeEurope, storeUSb, storeUSa1, storeUSa15}, - }, - { - name: "white list rule", - rule: func(store roachpb.StoreDescriptor, _ solveState) (bool, float64, float64) { - switch store.StoreID { - case storeUSa15: - return true, 0, 0 - case storeUSb: - return true, 1, 0 - default: - return false, 0, 0 - } - }, - expectedValid: []roachpb.StoreID{storeUSb, storeUSa15}, - expectedInvalid: []roachpb.StoreID{storeEurope, storeUSa1}, - }, - { - name: "ruleReplicasUniqueNodes - 2 available nodes", - rule: ruleReplicasUniqueNodes, - existing: []roachpb.ReplicaDescriptor{ - {NodeID: roachpb.NodeID(storeUSa15)}, - {NodeID: roachpb.NodeID(storeUSb)}, - }, - expectedValid: []roachpb.StoreID{storeEurope, storeUSa1}, - expectedInvalid: []roachpb.StoreID{storeUSb, storeUSa15}, - }, - { - name: "ruleReplicasUniqueNodes - 0 available nodes", - rule: ruleReplicasUniqueNodes, - existing: []roachpb.ReplicaDescriptor{ - {NodeID: roachpb.NodeID(storeUSa15)}, - {NodeID: roachpb.NodeID(storeUSa1)}, - {NodeID: roachpb.NodeID(storeUSb)}, - {NodeID: roachpb.NodeID(storeEurope)}, - }, - expectedValid: nil, - expectedInvalid: []roachpb.StoreID{storeEurope, storeUSb, storeUSa1, storeUSa15}, - }, - { - name: "ruleConstraints - required constraints", - rule: ruleConstraints, - c: config.Constraints{ - Constraints: []config.Constraint{ - {Value: "b", Type: config.Constraint_REQUIRED}, - }, - }, - expectedValid: []roachpb.StoreID{storeUSb, storeUSa1}, - expectedInvalid: []roachpb.StoreID{storeEurope, storeUSa15}, - }, - { - name: "ruleConstraints - required locality constraints", - rule: ruleConstraints, - c: config.Constraints{ - Constraints: []config.Constraint{ - {Key: "datacenter", Value: "us", Type: config.Constraint_REQUIRED}, - }, - }, - expectedValid: []roachpb.StoreID{storeUSb, storeUSa1, storeUSa15}, - expectedInvalid: []roachpb.StoreID{storeEurope}, - }, - { - name: "ruleConstraints - prohibited constraints", - rule: ruleConstraints, - c: config.Constraints{ - Constraints: []config.Constraint{ - {Value: "b", Type: config.Constraint_PROHIBITED}, - }, - }, - expectedValid: []roachpb.StoreID{storeEurope, storeUSa15}, - expectedInvalid: []roachpb.StoreID{storeUSb, storeUSa1}, - }, - { - name: "ruleConstraints - prohibited locality constraints", - rule: ruleConstraints, - c: config.Constraints{ - Constraints: []config.Constraint{ - {Key: "datacenter", Value: "us", Type: config.Constraint_PROHIBITED}, - }, - }, - expectedValid: []roachpb.StoreID{storeEurope}, - expectedInvalid: []roachpb.StoreID{storeUSb, storeUSa1, storeUSa15}, - }, - { - name: "ruleConstraints - positive constraints", - rule: ruleConstraints, - c: config.Constraints{ - Constraints: []config.Constraint{ - {Value: "a"}, - {Value: "b"}, - {Value: "c"}, - }, - }, - expectedValid: []roachpb.StoreID{storeUSb, storeUSa1, storeUSa15, storeEurope}, - }, - { - name: "ruleConstraints - positive locality constraints", - rule: ruleConstraints, - c: config.Constraints{ - Constraints: []config.Constraint{ - {Key: "datacenter", Value: "eur"}, - }, - }, - expectedValid: []roachpb.StoreID{storeEurope, storeUSb, storeUSa1, storeUSa15}, - }, - { - name: "ruleDiversity - no existing replicas", - rule: ruleDiversity, - expectedValid: []roachpb.StoreID{storeEurope, storeUSb, storeUSa1, storeUSa15}, - }, - { - name: "ruleDiversity - one existing replicas", - rule: ruleDiversity, - existing: []roachpb.ReplicaDescriptor{ - {NodeID: roachpb.NodeID(storeUSa15)}, - }, - expectedValid: []roachpb.StoreID{storeEurope, storeUSb, storeUSa1, storeUSa15}, - }, - { - name: "ruleDiversity - two existing replicas", - rule: ruleDiversity, - existing: []roachpb.ReplicaDescriptor{ - {NodeID: roachpb.NodeID(storeUSa15)}, - {NodeID: roachpb.NodeID(storeEurope)}, - }, - expectedValid: []roachpb.StoreID{storeUSb, storeUSa1, storeEurope, storeUSa15}, - }, - { - name: "ruleCapacityMax", - rule: ruleCapacityMax, - expectedValid: []roachpb.StoreID{storeEurope, storeUSb, storeUSa1}, - expectedInvalid: []roachpb.StoreID{storeUSa15}, - }, - { - name: "ruleCapacity", - rule: ruleCapacity, - expectedValid: []roachpb.StoreID{storeUSa1, storeEurope, storeUSb, storeUSa15}, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - var solver ruleSolver - if tc.rule != nil { - solver = ruleSolver{tc.rule} - } - sl, _, _ := storePool.getStoreList(roachpb.RangeID(0)) - candidates := solver.Solve( - sl, - tc.c, - tc.existing, - storePool.getNodeLocalities(tc.existing), - storePool.deterministic, - ) - valid := candidates.onlyValid() - invalid := candidates[len(valid):] - - if len(valid) != len(tc.expectedValid) { - t.Fatalf("length of valid %+v should match %+v", valid, tc.expectedValid) - } - for i, expected := range tc.expectedValid { - if actual := valid[i].store.StoreID; actual != expected { - t.Errorf("valid[%d].store.StoreID = %d; not %d; %+v", - i, actual, expected, valid) - } - } - if len(invalid) != len(tc.expectedInvalid) { - t.Fatalf("length of invalids %+v should match %+v", invalid, tc.expectedInvalid) - } - for i, expected := range tc.expectedInvalid { - if actual := invalid[i].store.StoreID; actual != expected { - t.Errorf("invalid[%d].store.StoreID = %d; not %d; %+v", - i, actual, expected, invalid) - } - } - }) - } -} - func TestOnlyValid(t *testing.T) { defer leaktest.AfterTest(t)() @@ -523,3 +261,369 @@ func TestBetterThan(t *testing.T) { } } } + +func TestPreexistingReplicaCheck(t *testing.T) { + defer leaktest.AfterTest(t)() + + var existing []roachpb.ReplicaDescriptor + for i := 2; i < 10; i += 2 { + existing = append(existing, roachpb.ReplicaDescriptor{NodeID: roachpb.NodeID(i)}) + } + for i := 1; i < 10; i++ { + if e, a := i%2 != 0, preexistingReplicaCheck(roachpb.NodeID(i), existing); e != a { + t.Errorf("NodeID %d expected to be %t, got %t", i, e, a) + } + } +} + +// testStoreTierSetup returns a tier struct constructed using the passed in values. +// If any value is an empty string, it is not included. +func testStoreTierSetup(datacenter, floor, rack, slot string) []roachpb.Tier { + var tiers []roachpb.Tier + if datacenter != "" { + tiers = append(tiers, roachpb.Tier{Key: "datacenter", Value: datacenter}) + } + if floor != "" { + tiers = append(tiers, roachpb.Tier{Key: "floor", Value: floor}) + } + if rack != "" { + tiers = append(tiers, roachpb.Tier{Key: "rack", Value: rack}) + } + if slot != "" { + tiers = append(tiers, roachpb.Tier{Key: "slot", Value: slot}) + } + return tiers +} + +// testStoreCapacitySetup returns a store capacity in which the total capacity +// is always 100 and available and range count are passed in. +func testStoreCapacitySetup(available int64, rangeCount int32) roachpb.StoreCapacity { + return roachpb.StoreCapacity{ + Capacity: 100, + Available: available, + RangeCount: rangeCount, + } +} + +// This is a collection of test stores used by a suite of tests. +var ( + testStoreUSa15 = roachpb.StoreID(1) // us-a-1-5 + testStoreUSa1 = roachpb.StoreID(2) // us-a-1 + testStoreUSb = roachpb.StoreID(3) // us-b + testStoreEurope = roachpb.StoreID(4) // eur-a-1-5 + + testStores = []roachpb.StoreDescriptor{ + { + StoreID: testStoreUSa15, + Attrs: roachpb.Attributes{ + Attrs: []string{"a"}, + }, + Node: roachpb.NodeDescriptor{ + NodeID: roachpb.NodeID(testStoreUSa15), + Locality: roachpb.Locality{ + Tiers: testStoreTierSetup("us", "a", "1", "5"), + }, + }, + Capacity: testStoreCapacitySetup(1, 99), + }, + { + StoreID: testStoreUSa1, + Attrs: roachpb.Attributes{ + Attrs: []string{"a", "b"}, + }, + Node: roachpb.NodeDescriptor{ + NodeID: roachpb.NodeID(testStoreUSa1), + Locality: roachpb.Locality{ + Tiers: testStoreTierSetup("us", "a", "1", ""), + }, + }, + Capacity: testStoreCapacitySetup(100, 0), + }, + { + StoreID: testStoreUSb, + Attrs: roachpb.Attributes{ + Attrs: []string{"a", "b", "c"}, + }, + Node: roachpb.NodeDescriptor{ + NodeID: roachpb.NodeID(testStoreUSb), + Locality: roachpb.Locality{ + Tiers: testStoreTierSetup("us", "b", "", ""), + }, + }, + Capacity: testStoreCapacitySetup(50, 50), + }, + { + StoreID: testStoreEurope, + Node: roachpb.NodeDescriptor{ + NodeID: roachpb.NodeID(testStoreEurope), + Locality: roachpb.Locality{ + Tiers: testStoreTierSetup("eur", "a", "1", "5"), + }, + }, + Capacity: testStoreCapacitySetup(60, 40), + }, + } +) + +func TestConstraintCheck(t *testing.T) { + defer leaktest.AfterTest(t)() + + testCases := []struct { + name string + constraints []config.Constraint + expected map[roachpb.StoreID]int + }{ + { + name: "required constraint", + constraints: []config.Constraint{ + {Value: "b", Type: config.Constraint_REQUIRED}, + }, + expected: map[roachpb.StoreID]int{ + testStoreUSa1: 0, + testStoreUSb: 0, + }, + }, + { + name: "required locality constraints", + constraints: []config.Constraint{ + {Key: "datacenter", Value: "us", Type: config.Constraint_REQUIRED}, + }, + expected: map[roachpb.StoreID]int{ + testStoreUSa15: 0, + testStoreUSa1: 0, + testStoreUSb: 0, + }, + }, + { + name: "prohibited constraints", + constraints: []config.Constraint{ + {Value: "b", Type: config.Constraint_PROHIBITED}, + }, + expected: map[roachpb.StoreID]int{ + testStoreUSa15: 0, + testStoreEurope: 0, + }, + }, + { + name: "prohibited locality constraints", + constraints: []config.Constraint{ + {Key: "datacenter", Value: "us", Type: config.Constraint_PROHIBITED}, + }, + expected: map[roachpb.StoreID]int{ + testStoreEurope: 0, + }, + }, + { + name: "positive constraints", + constraints: []config.Constraint{ + {Value: "a"}, + {Value: "b"}, + {Value: "c"}, + }, + expected: map[roachpb.StoreID]int{ + testStoreUSa15: 1, + testStoreUSa1: 2, + testStoreUSb: 3, + testStoreEurope: 0, + }, + }, + { + name: "positive locality constraints", + constraints: []config.Constraint{ + {Key: "datacenter", Value: "eur"}, + }, + expected: map[roachpb.StoreID]int{ + testStoreUSa15: 0, + testStoreUSa1: 0, + testStoreUSb: 0, + testStoreEurope: 1, + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + for _, s := range testStores { + valid, positive := constraintCheck(s, config.Constraints{Constraints: tc.constraints}) + expectedPositive, ok := tc.expected[s.StoreID] + if valid != ok { + t.Errorf("expected store %d to be %t, but got %t", s.StoreID, ok, valid) + continue + } + if positive != expectedPositive { + t.Errorf("expected store %d to have %d positives, but got %d", s.StoreID, expectedPositive, positive) + } + } + }) + } +} + +func TestDiversityScore(t *testing.T) { + defer leaktest.AfterTest(t)() + + testCases := []struct { + name string + existing []roachpb.NodeID + expected map[roachpb.StoreID]float64 + }{ + { + name: "no existing replicas", + expected: map[roachpb.StoreID]float64{ + testStoreUSa15: 1, + testStoreUSa1: 1, + testStoreUSb: 1, + testStoreEurope: 1, + }, + }, + { + name: "one existing replicas", + existing: []roachpb.NodeID{ + roachpb.NodeID(testStoreUSa15), + }, + expected: map[roachpb.StoreID]float64{ + testStoreUSa15: 0, + testStoreUSa1: 1.0 / 4.0, + testStoreUSb: 1.0 / 2.0, + testStoreEurope: 1, + }, + }, + { + name: "two existing replicas", + existing: []roachpb.NodeID{ + roachpb.NodeID(testStoreUSa15), + roachpb.NodeID(testStoreEurope), + }, + expected: map[roachpb.StoreID]float64{ + testStoreUSa15: 0, + testStoreUSa1: 1.0 / 4.0, + testStoreUSb: 1.0 / 2.0, + testStoreEurope: 0, + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + existingNodeLocalities := make(map[roachpb.NodeID]roachpb.Locality) + for _, nodeID := range tc.existing { + for _, s := range testStores { + if s.Node.NodeID == nodeID { + existingNodeLocalities[roachpb.NodeID(s.Node.NodeID)] = s.Node.Locality + } + } + } + for _, s := range testStores { + actualScore := diversityScore(s, existingNodeLocalities) + expectedScore, ok := tc.expected[s.StoreID] + if !ok { + t.Fatalf("no expected score found for storeID %d", s.StoreID) + } + if actualScore != expectedScore { + t.Errorf("store %d expected diversity score: %.2f, actual %.2f", s.StoreID, expectedScore, actualScore) + } + } + }) + } +} + +func TestDiversityRemovalScore(t *testing.T) { + defer leaktest.AfterTest(t)() + + testCases := []struct { + name string + expected map[roachpb.StoreID]float64 + }{ + { + name: "four existing replicas", + expected: map[roachpb.StoreID]float64{ + testStoreUSa15: 1, + testStoreUSa1: 1, + testStoreUSb: 1, + testStoreEurope: 1.0 / 2.0, + }, + }, + { + name: "three existing replicas - testStoreUSa15", + expected: map[roachpb.StoreID]float64{ + testStoreUSa1: 1, + testStoreUSb: 1, + testStoreEurope: 1.0 / 2.0, + }, + }, + { + name: "three existing replicas - testStoreUSa1", + expected: map[roachpb.StoreID]float64{ + testStoreUSa15: 1, + testStoreUSb: 1, + testStoreEurope: 1.0 / 2.0, + }, + }, + { + name: "three existing replicas - testStoreUSb", + expected: map[roachpb.StoreID]float64{ + testStoreUSa15: 1, + testStoreUSa1: 1, + testStoreEurope: 1.0 / 4.0, + }, + }, + { + name: "three existing replicas - testStoreEurope", + expected: map[roachpb.StoreID]float64{ + testStoreUSa15: 1.0 / 2.0, + testStoreUSa1: 1.0 / 2.0, + testStoreUSb: 1.0 / 4.0, + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + existingNodeLocalities := make(map[roachpb.NodeID]roachpb.Locality) + for _, s := range testStores { + if _, ok := tc.expected[s.StoreID]; ok { + existingNodeLocalities[roachpb.NodeID(s.Node.NodeID)] = s.Node.Locality + } + } + for _, s := range testStores { + if _, ok := tc.expected[s.StoreID]; !ok { + continue + } + actualScore := diversityRemovalScore(s.Node.NodeID, existingNodeLocalities) + expectedScore, ok := tc.expected[s.StoreID] + if !ok { + t.Fatalf("no expected score found for storeID %d", s.StoreID) + } + if actualScore != expectedScore { + t.Errorf("store %d expected diversity removal score: %.2f, actual %.2f", s.StoreID, expectedScore, actualScore) + } + } + }) + } +} + +// TestCapacityScore tests both capacityScore and maxCapacityCheck. +func TestCapacityScore(t *testing.T) { + defer leaktest.AfterTest(t)() + + expectedCheck := map[roachpb.StoreID]bool{ + testStoreUSa15: false, + testStoreUSa1: true, + testStoreUSb: true, + testStoreEurope: true, + } + expectedScore := map[roachpb.StoreID]float64{ + testStoreUSa15: 1.0 / 100.0, + testStoreUSa1: 1.0, + testStoreUSb: 1.0 / 51.0, + testStoreEurope: 1.0 / 41.0, + } + + for _, s := range testStores { + if e, a := expectedScore[s.StoreID], capacityScore(s); e != a { + t.Errorf("store %d expected capacity score: %.2f, actual %.2f", s.StoreID, e, a) + } + if e, a := expectedCheck[s.StoreID], maxCapacityCheck(s); e != a { + t.Errorf("store %d expected max capacity check: %t, actual %t", s.StoreID, e, a) + } + } +}