From e65167476d61d6e39be784fb7a23516e70da0d75 Mon Sep 17 00:00:00 2001 From: Bram Gruneir Date: Fri, 4 Nov 2016 14:30:13 -0400 Subject: [PATCH] storage: Re-apply the rule solver. Instead of applying 1ef40f323b94005dd84378669da8296c9fec8e7c or #10252, this finishes the reapplication of the rule solver. However, this also puts the rule solver under the environment flag COCKROACH_ENABLE_RULE_SOLVER for ease of testing and defaults to not enabled. This commit re-applies the rule solver, specifically the following commits: 1) 4446345bf7a9e1a6eb7c7a3c71ec9c8e3b3d17fc storage: add constraint rule solver for allocation Rules are represented as a single function that returns the candidacy of the store as well as a float value representing the score. These scores are then aggregated from all rules and returns the stores sorted by them. Current rules: - ruleReplicasUniqueNodes ensures that no two replicas are put on the same node. - ruleConstraints enforces that required and prohibited constraints are followed, and that stores with more positive constraints are ranked higher. - ruleDiversity ensures that nodes that have the fewest locality tiers in common are given higher priority. - ruleCapacity prioritizes placing data on empty nodes when the choice is available and prevents data from going onto mostly full nodes. 2) dd3229a718d5e3ec25a784f2f3ff5c9ffd85a3b0 storage: implemented RuleSolver into allocator The follow up to this commit is #10275 and a lot of testing to ensure that the rule solver does indeed perform as expected. Closes #9336 --- pkg/storage/allocator.go | 217 +++++- pkg/storage/allocator_test.go | 1154 ++++++++++++++++++++----------- pkg/storage/balancer.go | 49 -- pkg/storage/replicate_queue.go | 18 +- pkg/storage/rule_solver.go | 277 ++++++++ pkg/storage/rule_solver_test.go | 365 ++++++++++ pkg/storage/simulation/range.go | 7 +- pkg/storage/store.go | 4 + 8 files changed, 1619 insertions(+), 472 deletions(-) create mode 100644 pkg/storage/rule_solver.go create mode 100644 pkg/storage/rule_solver_test.go diff --git a/pkg/storage/allocator.go b/pkg/storage/allocator.go index 43b07ee05bbb..7a90f6620f32 100644 --- a/pkg/storage/allocator.go +++ b/pkg/storage/allocator.go @@ -20,6 +20,7 @@ package storage import ( "fmt" + "math" "math/rand" "golang.org/x/net/context" @@ -27,6 +28,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/config" "github.com/cockroachdb/cockroach/pkg/roachpb" "github.com/cockroachdb/cockroach/pkg/util" + "github.com/cockroachdb/cockroach/pkg/util/envutil" "github.com/cockroachdb/cockroach/pkg/util/log" "github.com/cockroachdb/cockroach/pkg/util/syncutil" "github.com/pkg/errors" @@ -123,28 +125,37 @@ type AllocatorOptions struct { // AllowRebalance allows this store to attempt to rebalance its own // replicas to other stores. AllowRebalance bool + + // UseRuleSolver enables this store to use the updated rules based + // constraint solver instead of the original rebalancer. + UseRuleSolver bool } // Allocator tries to spread replicas as evenly as possible across the stores // in the cluster. type Allocator struct { - storePool *StorePool - randGen allocatorRand - options AllocatorOptions + storePool *StorePool + randGen allocatorRand + options AllocatorOptions + ruleSolver ruleSolver } // MakeAllocator creates a new allocator using the specified StorePool. func MakeAllocator(storePool *StorePool, options AllocatorOptions) Allocator { var randSource rand.Source + // There are number of test cases that make a test store but don't add + // gossip or a store pool. So we can't rely on the existence of the + // store pool in those cases. if storePool != nil && storePool.deterministic { randSource = rand.NewSource(777) } else { randSource = rand.NewSource(rand.Int63()) } return Allocator{ - storePool: storePool, - options: options, - randGen: makeAllocatorRand(randSource), + storePool: storePool, + options: options, + randGen: makeAllocatorRand(randSource), + ruleSolver: makeRuleSolver(defaultSolverRules), } } @@ -203,6 +214,28 @@ func (a *Allocator) AllocateTarget( rangeID roachpb.RangeID, relaxConstraints bool, ) (*roachpb.StoreDescriptor, error) { + 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, err := a.ruleSolver.Solve(sl, constraints, existing) + if err != nil { + return nil, err + } + + if len(candidates) == 0 { + return nil, &allocatorError{ + required: constraints.Constraints, + } + } + // TODO(bram): #10275 Need some randomness here! + return &candidates[0].store, nil + } + existingNodes := make(nodeIDSet, len(existing)) for _, repl := range existing { existingNodes[repl.NodeID] = struct{}{} @@ -244,12 +277,55 @@ func (a *Allocator) AllocateTarget( // make correct decisions in the case of ranges with heterogeneous replica // requirements (i.e. multiple data centers). func (a Allocator) RemoveTarget( - existing []roachpb.ReplicaDescriptor, leaseStoreID roachpb.StoreID, + constraints config.Constraints, + existing []roachpb.ReplicaDescriptor, + leaseStoreID roachpb.StoreID, ) (roachpb.ReplicaDescriptor, error) { if len(existing) == 0 { 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)) + + var worst roachpb.ReplicaDescriptor + worstScore := math.Inf(0) + for _, exist := range existing { + if exist.StoreID == leaseStoreID { + continue + } + desc, ok := a.storePool.getStoreDescriptor(exist.StoreID) + if !ok { + continue + } + + // When a candidate is not valid, the score will be 0 and it should + // be chosen for removal. + if candidate, _ := a.ruleSolver.computeCandidate(solveState{ + constraints: constraints, + store: desc, + existing: nil, + sl: sl, + tierOrder: canonicalTierOrder(sl), + tiers: storeTierMap(sl), + }); candidate.score < worstScore { + worstScore = candidate.score + worst = exist + } + } + + if !math.IsInf(worstScore, 0) { + return worst, 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. var descriptors []roachpb.StoreDescriptor for _, exist := range existing { @@ -269,7 +345,7 @@ func (a Allocator) RemoveTarget( } } } - return roachpb.ReplicaDescriptor{}, errors.Errorf("RemoveTarget() could not select an appropriate replica to be remove") + return roachpb.ReplicaDescriptor{}, errors.New("could not select an appropriate replica to be removed") } // RebalanceTarget returns a suitable store for a rebalance target with @@ -297,9 +373,80 @@ func (a Allocator) RebalanceTarget( existing []roachpb.ReplicaDescriptor, leaseStoreID roachpb.StoreID, rangeID roachpb.RangeID, -) *roachpb.StoreDescriptor { +) (*roachpb.StoreDescriptor, error) { if !a.options.AllowRebalance { - return nil + return nil, nil + } + + 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) + } + + var shouldRebalance bool + for _, repl := range existing { + if leaseStoreID == repl.StoreID { + continue + } + storeDesc, ok := a.storePool.getStoreDescriptor(repl.StoreID) + if ok && a.shouldRebalance(storeDesc, sl) { + shouldRebalance = true + break + } + } + if !shouldRebalance { + 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) + + existingCandidates, err := a.ruleSolver.Solve(existingStoreList, constraints, nil) + if err != nil { + return nil, err + } + candidates, err := a.ruleSolver.Solve(candidateStoreList, constraints, nil) + if err != nil { + return nil, err + } + + // Find all candidates that are better than the worst existing store. + var worstCandidateStore float64 + // If any store from existing is not included in existingCandidates, it + // is because it no longer meets the constraints. If so, its score is + // considered to be 0. + if len(existingCandidates) == len(existing) { + worstCandidateStore = existingCandidates[len(existingCandidates)-1].score + } + + // TODO(bram): #10275 Need some randomness here! + for _, cand := range candidates { + if cand.score > worstCandidateStore { + return &candidates[0].store, nil + } + } + + return nil, nil } sl, _, _ := a.storePool.getStoreList(rangeID) @@ -320,14 +467,14 @@ func (a Allocator) RebalanceTarget( } } if !shouldRebalance { - return nil + return nil, nil } existingNodes := make(nodeIDSet, len(existing)) for _, repl := range existing { existingNodes[repl.NodeID] = struct{}{} } - return a.improve(sl, existingNodes) + return a.improve(sl, existingNodes), nil } // selectGood attempts to select a store from the supplied store list that it @@ -356,11 +503,53 @@ func (a Allocator) improve(sl StoreList, excluded nodeIDSet) *roachpb.StoreDescr return rcb.improve(sl, excluded) } +// rebalanceThreshold is the minimum ratio of a store's range surplus to the +// mean range count that permits rebalances away from that store. +var rebalanceThreshold = envutil.EnvOrDefaultFloat("COCKROACH_REBALANCE_THRESHOLD", 0.05) + // shouldRebalance returns whether the specified store is a candidate for // having a replica removed from it given the candidate store list. func (a Allocator) shouldRebalance(store roachpb.StoreDescriptor, sl StoreList) bool { - rcb := rangeCountBalancer{a.randGen} - return rcb.shouldRebalance(store, sl) + // TODO(peter,bram,cuong): The FractionUsed check seems suspicious. When a + // node becomes fuller than maxFractionUsedThreshold we will always select it + // for rebalancing. + maxCapacityUsed := store.Capacity.FractionUsed() >= maxFractionUsedThreshold + + // Rebalance if we're above the rebalance target, which is + // mean*(1+rebalanceThreshold). + target := int32(math.Ceil(sl.candidateCount.mean * (1 + rebalanceThreshold))) + rangeCountAboveTarget := store.Capacity.RangeCount > target + + // Rebalance if the candidate store has a range count above the mean, and + // there exists another store that is underfull: its range count is smaller + // than mean*(1-rebalanceThreshold). + var rebalanceToUnderfullStore bool + if float64(store.Capacity.RangeCount) > sl.candidateCount.mean { + underfullThreshold := int32(math.Floor(sl.candidateCount.mean * (1 - rebalanceThreshold))) + for _, desc := range sl.stores { + if desc.Capacity.RangeCount < underfullThreshold { + rebalanceToUnderfullStore = true + break + } + } + } + + // Require that moving a replica from the given store makes its range count + // converge on the mean range count. This only affects clusters with a + // small number of ranges. + rebalanceConvergesOnMean := rebalanceFromConvergesOnMean(sl, store) + + result := + (maxCapacityUsed || rangeCountAboveTarget || rebalanceToUnderfullStore) && rebalanceConvergesOnMean + if log.V(2) { + log.Infof(context.TODO(), + "%d: should-rebalance=%t: fraction-used=%.2f range-count=%d "+ + "(mean=%.1f, target=%d, fraction-used=%t, above-target=%t, underfull=%t, converges=%t)", + store.StoreID, result, store.Capacity.FractionUsed(), store.Capacity.RangeCount, + sl.candidateCount.mean, target, maxCapacityUsed, rangeCountAboveTarget, + rebalanceToUnderfullStore, rebalanceConvergesOnMean) + } + return result } // computeQuorum computes the quorum value for the given number of nodes. diff --git a/pkg/storage/allocator_test.go b/pkg/storage/allocator_test.go index 82ba633e26e9..c49da63f20ba 100644 --- a/pkg/storage/allocator_test.go +++ b/pkg/storage/allocator_test.go @@ -173,10 +173,13 @@ var multiDCStores = []*roachpb.StoreDescriptor{ // createTestAllocator creates a stopper, gossip, store pool and allocator for // use in tests. Stopper must be stopped by the caller. func createTestAllocator( - deterministic bool, + deterministic bool, useRuleSolver bool, ) (*stop.Stopper, *gossip.Gossip, *StorePool, Allocator, *hlc.ManualClock) { stopper, g, manualClock, storePool := createTestStorePool(TestTimeUntilStoreDeadOff, deterministic) - a := MakeAllocator(storePool, AllocatorOptions{AllowRebalance: true}) + a := MakeAllocator(storePool, AllocatorOptions{ + AllowRebalance: true, + UseRuleSolver: useRuleSolver, + }) return stopper, g, storePool, a, manualClock } @@ -194,13 +197,20 @@ func mockStorePool( storePool.mu.storeDetails = make(map[roachpb.StoreID]*storeDetail) for _, storeID := range aliveStoreIDs { detail := newStoreDetail(context.TODO()) - detail.desc = &roachpb.StoreDescriptor{StoreID: storeID} + detail.desc = &roachpb.StoreDescriptor{ + StoreID: storeID, + Node: roachpb.NodeDescriptor{NodeID: roachpb.NodeID(storeID)}, + } + storePool.mu.storeDetails[storeID] = detail } for _, storeID := range deadStoreIDs { detail := newStoreDetail(context.TODO()) detail.dead = true - detail.desc = &roachpb.StoreDescriptor{StoreID: storeID} + detail.desc = &roachpb.StoreDescriptor{ + StoreID: storeID, + Node: roachpb.NodeDescriptor{NodeID: roachpb.NodeID(storeID)}, + } storePool.mu.storeDetails[storeID] = detail } for storeID, detail := range storePool.mu.storeDetails { @@ -213,154 +223,193 @@ func mockStorePool( } } +func runToggleRuleSolver(t *testing.T, test func(useRuleSolver bool, t *testing.T)) { + t.Run("without rule solver", func(t *testing.T) { + test(false, t) + }) + t.Run("with rule solver", func(t *testing.T) { + test(true, t) + }) +} + func TestAllocatorSimpleRetrieval(t *testing.T) { defer leaktest.AfterTest(t)() - stopper, g, _, a, _ := createTestAllocator(false /* deterministic */) - defer stopper.Stop() - gossiputil.NewStoreGossiper(g).GossipStores(singleStore, t) - result, err := a.AllocateTarget( - simpleZoneConfig.Constraints, - []roachpb.ReplicaDescriptor{}, - firstRange, - false, - ) - if err != nil { - t.Fatalf("Unable to perform allocation: %v", err) - } - if result.Node.NodeID != 1 || result.StoreID != 1 { - t.Errorf("expected NodeID 1 and StoreID 1: %+v", result) - } + + runToggleRuleSolver(t, func(useRuleSolver bool, t *testing.T) { + stopper, g, _, a, _ := createTestAllocator( + /* deterministic */ false, + /* UseRuleSolver */ useRuleSolver, + ) + defer stopper.Stop() + gossiputil.NewStoreGossiper(g).GossipStores(singleStore, t) + result, err := a.AllocateTarget( + simpleZoneConfig.Constraints, + []roachpb.ReplicaDescriptor{}, + firstRange, + false, + ) + if err != nil { + t.Fatalf("Unable to perform allocation: %v", err) + } + if result.Node.NodeID != 1 || result.StoreID != 1 { + t.Errorf("expected NodeID 1 and StoreID 1: %+v", result) + } + }) } // TestAllocatorCorruptReplica ensures that the allocator never attempts to // allocate a new replica on top of a dead (corrupt) one. func TestAllocatorCorruptReplica(t *testing.T) { defer leaktest.AfterTest(t)() - stopper, g, sp, a, _ := createTestAllocator(false /* deterministic */) - defer stopper.Stop() - gossiputil.NewStoreGossiper(g).GossipStores(sameDCStores, t) - const store1ID = roachpb.StoreID(1) - - // Set store 1 to have a dead replica in the store pool. - sp.mu.Lock() - sp.mu.storeDetails[store1ID].deadReplicas[firstRange] = - []roachpb.ReplicaDescriptor{{ - NodeID: roachpb.NodeID(1), - StoreID: store1ID, - }} - sp.mu.Unlock() - - result, err := a.AllocateTarget( - simpleZoneConfig.Constraints, - []roachpb.ReplicaDescriptor{}, - firstRange, - true, - ) - if err != nil { - t.Fatal(err) - } - if result.Node.NodeID != 2 || result.StoreID != 2 { - t.Errorf("expected NodeID 2 and StoreID 2: %+v", result) - } + + runToggleRuleSolver(t, func(useRuleSolver bool, t *testing.T) { + stopper, g, sp, a, _ := createTestAllocator( + /* deterministic */ false, + /* useRuleSolver */ useRuleSolver, + ) + defer stopper.Stop() + gossiputil.NewStoreGossiper(g).GossipStores(sameDCStores, t) + const store1ID = roachpb.StoreID(1) + + // Set store 1 to have a dead replica in the store pool. + sp.mu.Lock() + sp.mu.storeDetails[store1ID].deadReplicas[firstRange] = + []roachpb.ReplicaDescriptor{{ + NodeID: roachpb.NodeID(1), + StoreID: store1ID, + }} + sp.mu.Unlock() + + result, err := a.AllocateTarget( + simpleZoneConfig.Constraints, + []roachpb.ReplicaDescriptor{}, + firstRange, + true, + ) + if err != nil { + t.Fatal(err) + } + if result.Node.NodeID != 2 || result.StoreID != 2 { + t.Errorf("expected NodeID 2 and StoreID 2: %+v", result) + } + }) } func TestAllocatorNoAvailableDisks(t *testing.T) { defer leaktest.AfterTest(t)() - stopper, _, _, a, _ := createTestAllocator(false /* deterministic */) - defer stopper.Stop() - result, err := a.AllocateTarget( - simpleZoneConfig.Constraints, - []roachpb.ReplicaDescriptor{}, - firstRange, - false, - ) - if result != nil { - t.Errorf("expected nil result: %+v", result) - } - if err == nil { - t.Errorf("allocation succeeded despite there being no available disks: %v", result) - } + + runToggleRuleSolver(t, func(useRuleSolver bool, t *testing.T) { + stopper, _, _, a, _ := createTestAllocator( + /* deterministic */ false, + /* useRuleSolver */ useRuleSolver, + ) + defer stopper.Stop() + result, err := a.AllocateTarget( + simpleZoneConfig.Constraints, + []roachpb.ReplicaDescriptor{}, + firstRange, + false, + ) + if result != nil { + t.Errorf("expected nil result: %+v", result) + } + if err == nil { + t.Errorf("allocation succeeded despite there being no available disks: %v", result) + } + }) } func TestAllocatorTwoDatacenters(t *testing.T) { defer leaktest.AfterTest(t)() - stopper, g, _, a, _ := createTestAllocator(false /* deterministic */) - defer stopper.Stop() - gossiputil.NewStoreGossiper(g).GossipStores(multiDCStores, t) - result1, err := a.AllocateTarget( - multiDCConfig.Constraints, - []roachpb.ReplicaDescriptor{}, - firstRange, - false, - ) - if err != nil { - t.Fatalf("Unable to perform allocation: %v", err) - } - result2, err := a.AllocateTarget( - multiDCConfig.Constraints, - []roachpb.ReplicaDescriptor{{ - NodeID: result1.Node.NodeID, - StoreID: result1.StoreID, - }}, - firstRange, - false, - ) - if err != nil { - t.Fatalf("Unable to perform allocation: %v", err) - } - ids := []int{int(result1.Node.NodeID), int(result2.Node.NodeID)} - sort.Ints(ids) - if expected := []int{1, 2}; !reflect.DeepEqual(ids, expected) { - t.Errorf("Expected nodes %+v: %+v vs %+v", expected, result1.Node, result2.Node) - } - // Verify that no result is forthcoming if we already have a replica. - result3, err := a.AllocateTarget( - multiDCConfig.Constraints, - []roachpb.ReplicaDescriptor{ - { + + runToggleRuleSolver(t, func(useRuleSolver bool, t *testing.T) { + stopper, g, _, a, _ := createTestAllocator( + /* deterministic */ false, + /* useRuleSolver */ useRuleSolver, + ) + defer stopper.Stop() + gossiputil.NewStoreGossiper(g).GossipStores(multiDCStores, t) + result1, err := a.AllocateTarget( + multiDCConfig.Constraints, + []roachpb.ReplicaDescriptor{}, + firstRange, + false, + ) + if err != nil { + t.Fatalf("Unable to perform allocation: %v", err) + } + result2, err := a.AllocateTarget( + multiDCConfig.Constraints, + []roachpb.ReplicaDescriptor{{ NodeID: result1.Node.NodeID, StoreID: result1.StoreID, + }}, + firstRange, + false, + ) + if err != nil { + t.Fatalf("Unable to perform allocation: %v", err) + } + ids := []int{int(result1.Node.NodeID), int(result2.Node.NodeID)} + sort.Ints(ids) + if expected := []int{1, 2}; !reflect.DeepEqual(ids, expected) { + t.Errorf("Expected nodes %+v: %+v vs %+v", expected, result1.Node, result2.Node) + } + // Verify that no result is forthcoming if we already have a replica. + result3, err := a.AllocateTarget( + multiDCConfig.Constraints, + []roachpb.ReplicaDescriptor{ + { + NodeID: result1.Node.NodeID, + StoreID: result1.StoreID, + }, + { + NodeID: result2.Node.NodeID, + StoreID: result2.StoreID, + }, }, - { - NodeID: result2.Node.NodeID, - StoreID: result2.StoreID, - }, - }, - firstRange, - false, - ) - if err == nil { - t.Errorf("expected error on allocation without available stores: %+v", result3) - } + firstRange, + false, + ) + if err == nil { + t.Errorf("expected error on allocation without available stores: %+v", result3) + } + }) } func TestAllocatorExistingReplica(t *testing.T) { defer leaktest.AfterTest(t)() - stopper, g, _, a, _ := createTestAllocator(false /* deterministic */) - defer stopper.Stop() - gossiputil.NewStoreGossiper(g).GossipStores(sameDCStores, t) - result, err := a.AllocateTarget( - config.Constraints{ - Constraints: []config.Constraint{ - {Value: "a"}, - {Value: "hdd"}, + + runToggleRuleSolver(t, func(useRuleSolver bool, t *testing.T) { + stopper, g, _, a, _ := createTestAllocator( + /* deterministic */ false, + /* useRuleSolver */ useRuleSolver, + ) + defer stopper.Stop() + gossiputil.NewStoreGossiper(g).GossipStores(sameDCStores, t) + result, err := a.AllocateTarget( + config.Constraints{ + Constraints: []config.Constraint{ + {Value: "a"}, + {Value: "hdd"}, + }, }, - }, - []roachpb.ReplicaDescriptor{ - { - NodeID: 2, - StoreID: 2, + []roachpb.ReplicaDescriptor{ + { + NodeID: 2, + StoreID: 2, + }, }, - }, - firstRange, - false, - ) - if err != nil { - t.Fatalf("Unable to perform allocation: %v", err) - } - if result.Node.NodeID != 3 || result.StoreID != 4 { - t.Errorf("expected result to have node 3 and store 4: %+v", result) - } + firstRange, + false, + ) + if err != nil { + t.Fatalf("Unable to perform allocation: %v", err) + } + if result.Node.NodeID != 3 || result.StoreID != 4 { + t.Errorf("expected result to have node 3 and store 4: %+v", result) + } + }) } // TestAllocatorRelaxConstraints verifies that attribute constraints @@ -368,62 +417,239 @@ func TestAllocatorExistingReplica(t *testing.T) { // if necessary to find an allocation target. func TestAllocatorRelaxConstraints(t *testing.T) { defer leaktest.AfterTest(t)() - stopper, g, _, a, _ := createTestAllocator(false /* deterministic */) - defer stopper.Stop() - gossiputil.NewStoreGossiper(g).GossipStores(multiDCStores, t) - testCases := []struct { - required []config.Constraint // attribute strings - existing []int // existing store/node ID - relaxConstraints bool // allow constraints to be relaxed? - expID int // expected store/node ID on allocate - expErr bool - }{ - // The two stores in the system have attributes: - // storeID=1 {"a", "ssd"} - // storeID=2 {"b", "ssd"} - {[]config.Constraint{{Value: "a"}, {Value: "ssd"}}, []int{}, true, 1, false}, - {[]config.Constraint{{Value: "a"}, {Value: "ssd"}}, []int{1}, true, 2, false}, - {[]config.Constraint{{Value: "a"}, {Value: "ssd"}}, []int{1}, false, 0, true}, - {[]config.Constraint{{Value: "a"}, {Value: "ssd"}}, []int{1, 2}, true, 0, true}, - {[]config.Constraint{{Value: "b"}, {Value: "ssd"}}, []int{}, true, 2, false}, - {[]config.Constraint{{Value: "b"}, {Value: "ssd"}}, []int{1}, true, 2, false}, - {[]config.Constraint{{Value: "b"}, {Value: "ssd"}}, []int{2}, false, 0, true}, - {[]config.Constraint{{Value: "b"}, {Value: "ssd"}}, []int{2}, true, 1, false}, - {[]config.Constraint{{Value: "b"}, {Value: "ssd"}}, []int{1, 2}, true, 0, true}, - {[]config.Constraint{{Value: "b"}, {Value: "hdd"}}, []int{}, true, 2, false}, - {[]config.Constraint{{Value: "b"}, {Value: "hdd"}}, []int{2}, true, 1, false}, - {[]config.Constraint{{Value: "b"}, {Value: "hdd"}}, []int{2}, false, 0, true}, - {[]config.Constraint{{Value: "b"}, {Value: "hdd"}}, []int{1, 2}, true, 0, true}, - {[]config.Constraint{{Value: "b"}, {Value: "ssd"}, {Value: "gpu"}}, []int{}, true, 2, false}, - {[]config.Constraint{{Value: "b"}, {Value: "hdd"}, {Value: "gpu"}}, []int{}, true, 2, false}, - } - for i, test := range testCases { - var existing []roachpb.ReplicaDescriptor - for _, id := range test.existing { - existing = append(existing, roachpb.ReplicaDescriptor{NodeID: roachpb.NodeID(id), StoreID: roachpb.StoreID(id)}) + t.Run("without rule solver", func(t *testing.T) { + stopper, g, _, a, _ := createTestAllocator( + /* deterministic */ false, + /* useRuleSolver */ false, + ) + defer stopper.Stop() + gossiputil.NewStoreGossiper(g).GossipStores(multiDCStores, t) + + testCases := []struct { + required []config.Constraint // attribute strings + existing []int // existing store/node ID + relaxConstraints bool // allow constraints to be relaxed? + expID int // expected store/node ID on allocate + expErr bool + }{ + // The two stores in the system have attributes: + // storeID=1 {"a", "ssd"} + // storeID=2 {"b", "ssd"} + {[]config.Constraint{{Value: "a"}, {Value: "ssd"}}, []int{}, true, 1, false}, + {[]config.Constraint{{Value: "a"}, {Value: "ssd"}}, []int{1}, true, 2, false}, + {[]config.Constraint{{Value: "a"}, {Value: "ssd"}}, []int{1}, false, 0, true}, + {[]config.Constraint{{Value: "a"}, {Value: "ssd"}}, []int{1, 2}, true, 0, true}, + {[]config.Constraint{{Value: "b"}, {Value: "ssd"}}, []int{}, true, 2, false}, + {[]config.Constraint{{Value: "b"}, {Value: "ssd"}}, []int{1}, true, 2, false}, + {[]config.Constraint{{Value: "b"}, {Value: "ssd"}}, []int{2}, false, 0, true}, + {[]config.Constraint{{Value: "b"}, {Value: "ssd"}}, []int{2}, true, 1, false}, + {[]config.Constraint{{Value: "b"}, {Value: "ssd"}}, []int{1, 2}, true, 0, true}, + {[]config.Constraint{{Value: "b"}, {Value: "hdd"}}, []int{}, true, 2, false}, + {[]config.Constraint{{Value: "b"}, {Value: "hdd"}}, []int{2}, true, 1, false}, + {[]config.Constraint{{Value: "b"}, {Value: "hdd"}}, []int{2}, false, 0, true}, + {[]config.Constraint{{Value: "b"}, {Value: "hdd"}}, []int{1, 2}, true, 0, true}, + {[]config.Constraint{{Value: "b"}, {Value: "ssd"}, {Value: "gpu"}}, []int{}, true, 2, false}, + {[]config.Constraint{{Value: "b"}, {Value: "hdd"}, {Value: "gpu"}}, []int{}, true, 2, false}, } - constraints := config.Constraints{Constraints: test.required} - result, err := a.AllocateTarget( - constraints, - existing, - firstRange, - test.relaxConstraints, + for i, test := range testCases { + var existing []roachpb.ReplicaDescriptor + for _, id := range test.existing { + existing = append(existing, roachpb.ReplicaDescriptor{NodeID: roachpb.NodeID(id), StoreID: roachpb.StoreID(id)}) + } + constraints := config.Constraints{Constraints: test.required} + result, err := a.AllocateTarget( + constraints, + existing, + firstRange, + test.relaxConstraints, + ) + if haveErr := (err != nil); haveErr != test.expErr { + t.Errorf("%d: expected error %t; got %t: %s", i, test.expErr, haveErr, err) + } else if err == nil && roachpb.StoreID(test.expID) != result.StoreID { + t.Errorf("%d: expected result to have store %d; got %+v", i, test.expID, result) + } + } + }) + + t.Run("with rule solver", func(t *testing.T) { + stopper, g, _, a, _ := createTestAllocator( + /* deterministic */ false, + /* useRuleSolver */ true, ) - if haveErr := (err != nil); haveErr != test.expErr { - t.Errorf("%d: expected error %t; got %t: %s", i, test.expErr, haveErr, err) - } else if err == nil && roachpb.StoreID(test.expID) != result.StoreID { - t.Errorf("%d: expected result to have store %d; got %+v", i, test.expID, result) + defer stopper.Stop() + gossiputil.NewStoreGossiper(g).GossipStores(multiDCStores, t) + + testCases := []struct { + name string + constraints []config.Constraint + existing []int // existing store/node ID + expID int // expected store/node ID on allocate + expErr bool + }{ + // The two stores in the system have attributes: + // storeID=1 {"a", "ssd"} + // storeID=2 {"b", "ssd"} + { + name: "positive constraints (matching store 1)", + constraints: []config.Constraint{ + {Value: "a"}, + {Value: "ssd"}, + }, + expID: 1, + }, + { + name: "positive constraints (matching store 2)", + constraints: []config.Constraint{ + {Value: "b"}, + {Value: "ssd"}, + }, + expID: 2, + }, + { + name: "positive constraints (matching store 1) with existing replica (store 1)", + constraints: []config.Constraint{ + {Value: "a"}, + {Value: "ssd"}, + }, + existing: []int{1}, + expID: 2, + }, + { + name: "positive constraints (matching store 1) with two existing replicas", + constraints: []config.Constraint{ + {Value: "a"}, /* remove these?*/ + {Value: "ssd"}, + }, + existing: []int{1, 2}, + expErr: true, + }, + { + name: "positive constraints (matching store 2) with two existing replicas", + constraints: []config.Constraint{ + {Value: "b"}, + {Value: "ssd"}, + }, + existing: []int{1, 2}, + expErr: true, + }, + { + name: "required constraints (matching store 1) with existing replica (store 1)", + constraints: []config.Constraint{ + {Value: "a", Type: config.Constraint_REQUIRED}, + {Value: "ssd", Type: config.Constraint_REQUIRED}, + }, + existing: []int{1}, + expErr: true, + }, + { + name: "required constraints (matching store 2) with exiting replica (store 2)", + constraints: []config.Constraint{ + {Value: "b", Type: config.Constraint_REQUIRED}, + {Value: "ssd", Type: config.Constraint_REQUIRED}, + }, + existing: []int{2}, + expErr: true, + }, + { + name: "positive constraints (matching store 2) with existing replica (store 1)", + constraints: []config.Constraint{ + {Value: "b"}, + {Value: "ssd"}, + }, + existing: []int{1}, + expID: 2, + }, + { + name: "positive constraints (matching store 2) with existing replica (store 2)", + constraints: []config.Constraint{ + {Value: "b"}, + {Value: "ssd"}, + }, + existing: []int{2}, + expID: 1, + }, + { + name: "positive constraints (half matching store 2)", + constraints: []config.Constraint{ + {Value: "b"}, + {Value: "hdd"}, + }, + expID: 2, + }, + { + name: "positive constraints (half matching store 2) with existing replica (store 2)", + constraints: []config.Constraint{ + {Value: "b"}, + {Value: "hdd"}, + }, + existing: []int{2}, + expID: 1, + }, + { + name: "required constraints (half matching store 2) with existing replica (store 2)", + constraints: []config.Constraint{ + {Value: "b", Type: config.Constraint_REQUIRED}, + {Value: "hdd", Type: config.Constraint_REQUIRED}, + }, + existing: []int{2}, + expErr: true, + }, + { + name: "positive constraints (half matching store 2) with two existing replica", + constraints: []config.Constraint{ + {Value: "b"}, + {Value: "hdd"}, + }, + existing: []int{1, 2}, + expErr: true, + }, + { + name: "positive constraints (2/3 matching store 2)", + constraints: []config.Constraint{ + {Value: "b"}, + {Value: "ssd"}, + {Value: "gpu"}, + }, + expID: 2, + }, + { + name: "positive constraints (1/3 matching store 2)", + constraints: []config.Constraint{ + {Value: "b"}, + {Value: "hdd"}, + {Value: "gpu"}, + }, + expID: 2, + }, } - } + for _, test := range testCases { + t.Run(test.name, func(t *testing.T) { + var existing []roachpb.ReplicaDescriptor + for _, id := range test.existing { + existing = append(existing, roachpb.ReplicaDescriptor{NodeID: roachpb.NodeID(id), StoreID: roachpb.StoreID(id)}) + } + result, err := a.AllocateTarget( + config.Constraints{Constraints: test.constraints}, + existing, + firstRange, + false, + ) + if haveErr := (err != nil); haveErr != test.expErr { + t.Errorf("expected error %t; got %t: %s", test.expErr, haveErr, err) + } else if err == nil && roachpb.StoreID(test.expID) != result.StoreID { + t.Errorf("expected result to have store %d; got %+v", test.expID, result) + } + }) + } + }) + } // TestAllocatorRebalance verifies that rebalance targets are chosen // randomly from amongst stores over the minAvailCapacityThreshold. func TestAllocatorRebalance(t *testing.T) { defer leaktest.AfterTest(t)() - stopper, g, _, a, _ := createTestAllocator(false /* deterministic */) - defer stopper.Stop() stores := []*roachpb.StoreDescriptor{ { @@ -467,36 +693,48 @@ func TestAllocatorRebalance(t *testing.T) { }, }, } - gossiputil.NewStoreGossiper(g).GossipStores(stores, t) - - // Every rebalance target must be either stores 1 or 2. - for i := 0; i < 10; i++ { - result := a.RebalanceTarget( - config.Constraints{}, - []roachpb.ReplicaDescriptor{{StoreID: 3}}, - noStore, - firstRange, + + runToggleRuleSolver(t, func(useRuleSolver bool, t *testing.T) { + stopper, g, _, a, _ := createTestAllocator( + /* deterministic */ false, + /* useRuleSolver */ useRuleSolver, ) - if result == nil { - t.Fatal("nil result") - } - if result.StoreID != 1 && result.StoreID != 2 { - t.Errorf("%d: expected store 1 or 2; got %d", i, result.StoreID) - } - } + defer stopper.Stop() - // Verify shouldRebalance results. - for i, store := range stores { - desc, ok := a.storePool.getStoreDescriptor(store.StoreID) - if !ok { - t.Fatalf("%d: unable to get store %d descriptor", i, store.StoreID) + gossiputil.NewStoreGossiper(g).GossipStores(stores, t) + + // Every rebalance target must be either stores 1 or 2. + for i := 0; i < 10; i++ { + result, err := a.RebalanceTarget( + config.Constraints{}, + []roachpb.ReplicaDescriptor{{StoreID: 3}}, + noStore, + firstRange, + ) + if err != nil { + t.Fatal(err) + } + if result == nil { + t.Fatal("nil result") + } + if result.StoreID != 1 && result.StoreID != 2 { + t.Errorf("%d: expected store 1 or 2; got %d", i, result.StoreID) + } } - sl, _, _ := a.storePool.getStoreList(firstRange) - result := a.shouldRebalance(desc, sl) - if expResult := (i >= 2); expResult != result { - t.Errorf("%d: expected rebalance %t; got %t", i, expResult, result) + + // Verify shouldRebalance results. + for i, store := range stores { + desc, ok := a.storePool.getStoreDescriptor(store.StoreID) + if !ok { + t.Fatalf("%d: unable to get store %d descriptor", i, store.StoreID) + } + sl, _, _ := a.storePool.getStoreList(firstRange) + result := a.shouldRebalance(desc, sl) + if expResult := (i >= 2); expResult != result { + t.Errorf("%d: expected rebalance %t; got %t", i, expResult, result) + } } - } + }) } // TestAllocatorRebalanceThrashing tests that the rebalancer does not thrash @@ -546,80 +784,85 @@ func TestAllocatorRebalanceThrashing(t *testing.T) { return stores } - // Each test case defines the range counts for the test stores and whether we - // should rebalance from the store. - testCases := [][]testStore{ - // An evenly balanced cluster should not rebalance. - {{5, false}, {5, false}, {5, false}, {5, false}}, - // A very nearly balanced cluster should not rebalance. - {{5, false}, {5, false}, {5, false}, {6, false}}, - // Adding an empty node to a 3-node cluster triggers rebalancing from - // existing nodes. - {{100, true}, {100, true}, {100, true}, {0, false}}, - // A cluster where all range counts are within RebalanceThreshold should - // not rebalance. This assumes RebalanceThreshold > 2%. - {{98, false}, {99, false}, {101, false}, {102, false}}, - - // 5-nodes, each with a single store above the rebalancer target range - // count. - oneStoreAboveRebalanceTarget(100, 5), - oneStoreAboveRebalanceTarget(1000, 5), - oneStoreAboveRebalanceTarget(10000, 5), - - oneUnderusedStore(1000, 5), - oneUnderusedStore(1000, 10), - } - for i, tc := range testCases { - t.Logf("test case %d: %v", i, tc) - - // It doesn't make sense to test sets of stores containing fewer than 4 - // stores, because 4 stores is the minimum number of stores needed to - // trigger rebalancing with the default replication factor of 3. Also, the - // above local functions need a minimum number of stores to properly create - // the desired distribution of range counts. - const minStores = 4 - if numStores := len(tc); numStores < minStores { - t.Fatalf("%d: numStores %d < min %d", i, numStores, minStores) + runToggleRuleSolver(t, func(useRuleSolver bool, t *testing.T) { + // Each test case defines the range counts for the test stores and whether we + // should rebalance from the store. + testCases := [][]testStore{ + // An evenly balanced cluster should not rebalance. + {{5, false}, {5, false}, {5, false}, {5, false}}, + // A very nearly balanced cluster should not rebalance. + {{5, false}, {5, false}, {5, false}, {6, false}}, + // Adding an empty node to a 3-node cluster triggers rebalancing from + // existing nodes. + {{100, true}, {100, true}, {100, true}, {0, false}}, + // A cluster where all range counts are within RebalanceThreshold should + // not rebalance. This assumes RebalanceThreshold > 2%. + {{98, false}, {99, false}, {101, false}, {102, false}}, + + // 5-nodes, each with a single store above the rebalancer target range + // count. + oneStoreAboveRebalanceTarget(100, 5), + oneStoreAboveRebalanceTarget(1000, 5), + oneStoreAboveRebalanceTarget(10000, 5), + + oneUnderusedStore(1000, 5), + oneUnderusedStore(1000, 10), } - // Deterministic is required when stressing as test case 8 may rebalance - // to different configurations. - stopper, g, _, a, _ := createTestAllocator(true /* deterministic */) - defer stopper.Stop() - - // Create stores with the range counts from the test case and gossip them. - var stores []*roachpb.StoreDescriptor - for j, store := range tc { - stores = append(stores, &roachpb.StoreDescriptor{ - StoreID: roachpb.StoreID(j + 1), - Node: roachpb.NodeDescriptor{NodeID: roachpb.NodeID(j + 1)}, - Capacity: roachpb.StoreCapacity{Capacity: 1, Available: 1, RangeCount: store.rangeCount}, + for i, tc := range testCases { + t.Logf("test case %d: %v", i, tc) + + // It doesn't make sense to test sets of stores containing fewer than 4 + // stores, because 4 stores is the minimum number of stores needed to + // trigger rebalancing with the default replication factor of 3. Also, the + // above local functions need a minimum number of stores to properly create + // the desired distribution of range counts. + const minStores = 4 + if numStores := len(tc); numStores < minStores { + t.Fatalf("%d: numStores %d < min %d", i, numStores, minStores) + } + // Deterministic is required when stressing as test case 8 may rebalance + // to different configurations. + stopper, g, _, a, _ := createTestAllocator( + /* deterministic */ true, + /* UseRuleSolver */ useRuleSolver, + ) + defer stopper.Stop() + + // Create stores with the range counts from the test case and gossip them. + var stores []*roachpb.StoreDescriptor + for j, store := range tc { + stores = append(stores, &roachpb.StoreDescriptor{ + StoreID: roachpb.StoreID(j + 1), + Node: roachpb.NodeDescriptor{NodeID: roachpb.NodeID(j + 1)}, + Capacity: roachpb.StoreCapacity{Capacity: 1, Available: 1, RangeCount: store.rangeCount}, + }) + } + gossiputil.NewStoreGossiper(g).GossipStores(stores, t) + + // Ensure gossiped store descriptor changes have propagated. + util.SucceedsSoon(t, func() error { + sl, _, _ := a.storePool.getStoreList(firstRange) + for j, s := range sl.stores { + if a, e := s.Capacity.RangeCount, tc[j].rangeCount; a != e { + return errors.Errorf("tc %d: range count for %d = %d != expected %d", i, j, a, e) + } + } + return nil }) - } - gossiputil.NewStoreGossiper(g).GossipStores(stores, t) - - // Ensure gossiped store descriptor changes have propagated. - util.SucceedsSoon(t, func() error { sl, _, _ := a.storePool.getStoreList(firstRange) - for j, s := range sl.stores { - if a, e := s.Capacity.RangeCount, tc[j].rangeCount; a != e { - return errors.Errorf("tc %d: range count for %d = %d != expected %d", i, j, a, e) - } - } - return nil - }) - sl, _, _ := a.storePool.getStoreList(firstRange) - // Verify shouldRebalance returns the expected value. - for j, store := range stores { - desc, ok := a.storePool.getStoreDescriptor(store.StoreID) - if !ok { - t.Fatalf("[tc %d,store %d]: unable to get store %d descriptor", i, j, store.StoreID) - } - if a, e := a.shouldRebalance(desc, sl), tc[j].shouldRebalanceFrom; a != e { - t.Errorf("[tc %d,store %d]: shouldRebalance %t != expected %t", i, store.StoreID, a, e) + // Verify shouldRebalance returns the expected value. + for j, store := range stores { + desc, ok := a.storePool.getStoreDescriptor(store.StoreID) + if !ok { + t.Fatalf("[tc %d,store %d]: unable to get store %d descriptor", i, j, store.StoreID) + } + if a, e := a.shouldRebalance(desc, sl), tc[j].shouldRebalanceFrom; a != e { + t.Errorf("[tc %d,store %d]: shouldRebalance %t != expected %t", i, store.StoreID, a, e) + } } } - } + }) } // TestAllocatorRebalanceByCount verifies that rebalance targets are @@ -627,8 +870,6 @@ func TestAllocatorRebalanceThrashing(t *testing.T) { // exceed the maxAvailCapacityThreshold. func TestAllocatorRebalanceByCount(t *testing.T) { defer leaktest.AfterTest(t)() - stopper, g, _, a, _ := createTestAllocator(false /* deterministic */) - defer stopper.Stop() // Setup the stores so that only one is below the standard deviation threshold. stores := []*roachpb.StoreDescriptor{ @@ -653,41 +894,51 @@ func TestAllocatorRebalanceByCount(t *testing.T) { Capacity: roachpb.StoreCapacity{Capacity: 100, Available: 98, RangeCount: 2}, }, } - gossiputil.NewStoreGossiper(g).GossipStores(stores, t) - - // Every rebalance target must be store 4 (or nil for case of missing the only option). - for i := 0; i < 10; i++ { - result := a.RebalanceTarget( - config.Constraints{}, - []roachpb.ReplicaDescriptor{{StoreID: stores[0].StoreID}}, - stores[0].StoreID, - firstRange, + + runToggleRuleSolver(t, func(useRuleSolver bool, t *testing.T) { + stopper, g, _, a, _ := createTestAllocator( + /* deterministic */ false, + /* useRuleSolver */ useRuleSolver, ) - if result != nil && result.StoreID != 4 { - t.Errorf("expected store 4; got %d", result.StoreID) - } - } + defer stopper.Stop() - // Verify shouldRebalance results. - for i, store := range stores { - desc, ok := a.storePool.getStoreDescriptor(store.StoreID) - if !ok { - t.Fatalf("%d: unable to get store %d descriptor", i, store.StoreID) + gossiputil.NewStoreGossiper(g).GossipStores(stores, t) + + // Every rebalance target must be store 4 (or nil for case of missing the only option). + for i := 0; i < 10; i++ { + result, err := a.RebalanceTarget( + config.Constraints{}, + []roachpb.ReplicaDescriptor{{StoreID: stores[0].StoreID}}, + stores[0].StoreID, + firstRange, + ) + if err != nil { + t.Fatal(err) + } + if result != nil && result.StoreID != 4 { + t.Errorf("expected store 4; got %d", result.StoreID) + } } - sl, _, _ := a.storePool.getStoreList(firstRange) - result := a.shouldRebalance(desc, sl) - if expResult := (i < 3); expResult != result { - t.Errorf("%d: expected rebalance %t; got %t", i, expResult, result) + + // Verify shouldRebalance results. + for i, store := range stores { + desc, ok := a.storePool.getStoreDescriptor(store.StoreID) + if !ok { + t.Fatalf("%d: unable to get store %d descriptor", i, store.StoreID) + } + sl, _, _ := a.storePool.getStoreList(firstRange) + result := a.shouldRebalance(desc, sl) + if expResult := (i < 3); expResult != result { + t.Errorf("%d: expected rebalance %t; got %t", i, expResult, result) + } } - } + }) } // TestAllocatorRemoveTarget verifies that the replica chosen by RemoveTarget is // the one with the lowest capacity. func TestAllocatorRemoveTarget(t *testing.T) { defer leaktest.AfterTest(t)() - stopper, g, _, a, _ := createTestAllocator(false /* deterministic */) - defer stopper.Stop() // List of replicas that will be passed to RemoveTarget replicas := []roachpb.ReplicaDescriptor{ @@ -737,60 +988,54 @@ func TestAllocatorRemoveTarget(t *testing.T) { Capacity: roachpb.StoreCapacity{Capacity: 100, Available: 65, RangeCount: 10}, }, } - sg := gossiputil.NewStoreGossiper(g) - sg.GossipStores(stores, t) - // Exclude store 2 as a removal target so that only store 3 is a candidate. - targetRepl, err := a.RemoveTarget(replicas, stores[1].StoreID) - if err != nil { - t.Fatal(err) - } - if a, e := targetRepl, replicas[2]; a != e { - t.Fatalf("RemoveTarget did not select expected replica; expected %v, got %v", e, a) - } + runToggleRuleSolver(t, func(useRuleSolver bool, t *testing.T) { + stopper, g, _, a, _ := createTestAllocator( + /* deterministic */ false, + /* useRuleSolver */ useRuleSolver, + ) + defer stopper.Stop() + sg := gossiputil.NewStoreGossiper(g) + sg.GossipStores(stores, t) - // Now exclude store 3 so that only store 2 is a candidate. - targetRepl, err = a.RemoveTarget(replicas, stores[2].StoreID) - if err != nil { - t.Fatal(err) - } - if a, e := targetRepl, replicas[1]; a != e { - t.Fatalf("RemoveTarget did not select expected replica; expected %v, got %v", e, a) - } + // Exclude store 2 as a removal target so that only store 3 is a candidate. + targetRepl, err := a.RemoveTarget(config.Constraints{}, replicas, stores[1].StoreID) + if err != nil { + t.Fatal(err) + } + if a, e := targetRepl, replicas[2]; a != e { + t.Fatalf("RemoveTarget did not select expected replica; expected %v, got %v", e, a) + } - var counts [4]int - for i := 0; i < 100; i++ { - // Exclude store 1 as a removal target. We should see stores 2 and 3 - // randomly selected as the removal target. - targetRepl, err := a.RemoveTarget(replicas, stores[0].StoreID) + // Now exclude store 3 so that only store 2 is a candidate. + targetRepl, err = a.RemoveTarget(config.Constraints{}, replicas, stores[2].StoreID) if err != nil { t.Fatal(err) } - counts[targetRepl.StoreID-1]++ - } - if counts[0] != 0 || counts[3] != 0 || counts[1] == 0 || counts[2] == 0 { - t.Fatalf("unexpected removal target counts: %d", counts) - } + if a, e := targetRepl, replicas[1]; a != e { + t.Fatalf("RemoveTarget did not select expected replica; expected %v, got %v", e, a) + } + + if !useRuleSolver { + var counts [4]int + for i := 0; i < 100; i++ { + // Exclude store 1 as a removal target. We should see stores 2 and 3 + // randomly selected as the removal target. + targetRepl, err := a.RemoveTarget(config.Constraints{}, replicas, stores[0].StoreID) + if err != nil { + t.Fatal(err) + } + counts[targetRepl.StoreID-1]++ + } + if counts[0] != 0 || counts[1] == 0 || counts[2] == 0 || counts[3] != 0 { + t.Fatalf("unexpected removal target counts: %d", counts) + } + } + }) } func TestAllocatorComputeAction(t *testing.T) { defer leaktest.AfterTest(t)() - stopper, _, sp, a, _ := createTestAllocator(false /* deterministic */) - defer stopper.Stop() - - // Set up eight stores. Stores six and seven are marked as dead. Replica eight - // is dead. - mockStorePool(sp, - []roachpb.StoreID{1, 2, 3, 4, 5, 8}, - []roachpb.StoreID{6, 7}, - []roachpb.ReplicaIdent{{ - RangeID: 0, - Replica: roachpb.ReplicaDescriptor{ - NodeID: 8, - StoreID: 8, - ReplicaID: 8, - }, - }}) // Each test case should describe a repair situation which has a lower // priority than the previous test case. @@ -1116,32 +1361,57 @@ func TestAllocatorComputeAction(t *testing.T) { }, } - lastPriority := float64(999999999) - for i, tcase := range testCases { - action, priority := a.ComputeAction(tcase.zone, &tcase.desc) - if tcase.expectedAction != action { - t.Errorf("Test case %d expected action %d, got action %d", i, tcase.expectedAction, action) - continue - } - if tcase.expectedAction != AllocatorNoop && priority > lastPriority { - t.Errorf("Test cases should have descending priority. Case %d had priority %f, previous case had priority %f", i, priority, lastPriority) + runToggleRuleSolver(t, func(useRuleSolver bool, t *testing.T) { + stopper, _, sp, a, _ := createTestAllocator( + /* deterministic */ false, + /* useRuleSolver */ useRuleSolver, + ) + defer stopper.Stop() + + // Set up eight stores. Stores six and seven are marked as dead. Replica eight + // is dead. + mockStorePool(sp, + []roachpb.StoreID{1, 2, 3, 4, 5, 8}, + []roachpb.StoreID{6, 7}, + []roachpb.ReplicaIdent{{ + RangeID: 0, + Replica: roachpb.ReplicaDescriptor{ + NodeID: 8, + StoreID: 8, + ReplicaID: 8, + }, + }}) + + lastPriority := float64(999999999) + for i, tcase := range testCases { + action, priority := a.ComputeAction(tcase.zone, &tcase.desc) + if tcase.expectedAction != action { + t.Errorf("Test case %d expected action %d, got action %d", i, tcase.expectedAction, action) + continue + } + if tcase.expectedAction != AllocatorNoop && priority > lastPriority { + t.Errorf("Test cases should have descending priority. Case %d had priority %f, previous case had priority %f", i, priority, lastPriority) + } + lastPriority = priority } - lastPriority = priority - } + }) } // TestAllocatorComputeActionNoStorePool verifies that // ComputeAction returns AllocatorNoop when storePool is nil. func TestAllocatorComputeActionNoStorePool(t *testing.T) { defer leaktest.AfterTest(t)() - a := MakeAllocator(nil /* storePool */, AllocatorOptions{}) - action, priority := a.ComputeAction(config.ZoneConfig{}, nil) - if action != AllocatorNoop { - t.Errorf("expected AllocatorNoop, but got %v", action) - } - if priority != 0 { - t.Errorf("expected priority 0, but got %f", priority) - } + + runToggleRuleSolver(t, func(useRuleSolver bool, t *testing.T) { + a := MakeAllocator(nil /* storePool */, AllocatorOptions{UseRuleSolver: useRuleSolver}) + action, priority := a.ComputeAction(config.ZoneConfig{}, nil) + if action != AllocatorNoop { + t.Errorf("expected AllocatorNoop, but got %v", action) + } + if priority != 0 { + t.Errorf("expected priority 0, but got %f", priority) + } + }) } // TestAllocatorError ensures that the correctly formatted error message is @@ -1187,53 +1457,59 @@ func TestAllocatorError(t *testing.T) { // will not be sent to purgatory. func TestAllocatorThrottled(t *testing.T) { defer leaktest.AfterTest(t)() - stopper, g, _, a, _ := createTestAllocator(false /* deterministic */) - defer stopper.Stop() - // First test to make sure we would send the replica to purgatory. - _, err := a.AllocateTarget( - simpleZoneConfig.Constraints, - []roachpb.ReplicaDescriptor{}, - firstRange, - false, - ) - if _, ok := err.(purgatoryError); !ok { - t.Fatalf("expected a purgatory error, got: %v", err) - } + runToggleRuleSolver(t, func(useRuleSolver bool, t *testing.T) { + stopper, g, _, a, _ := createTestAllocator( + /* deterministic */ false, + /* useRuleSolver */ false, + ) + defer stopper.Stop() - // Second, test the normal case in which we can allocate to the store. - gossiputil.NewStoreGossiper(g).GossipStores(singleStore, t) - result, err := a.AllocateTarget( - simpleZoneConfig.Constraints, - []roachpb.ReplicaDescriptor{}, - firstRange, - false, - ) - if err != nil { - t.Fatalf("unable to perform allocation: %v", err) - } - if result.Node.NodeID != 1 || result.StoreID != 1 { - t.Errorf("expected NodeID 1 and StoreID 1: %+v", result) - } + // First test to make sure we would send the replica to purgatory. + _, err := a.AllocateTarget( + simpleZoneConfig.Constraints, + []roachpb.ReplicaDescriptor{}, + firstRange, + false, + ) + if _, ok := err.(purgatoryError); !ok { + t.Fatalf("expected a purgatory error, got: %v", err) + } - // Finally, set that store to be throttled and ensure we don't send the - // replica to purgatory. - a.storePool.mu.Lock() - storeDetail, ok := a.storePool.mu.storeDetails[singleStore[0].StoreID] - if !ok { - t.Fatalf("store:%d was not found in the store pool", singleStore[0].StoreID) - } - storeDetail.throttledUntil = timeutil.Now().Add(24 * time.Hour) - a.storePool.mu.Unlock() - _, err = a.AllocateTarget( - simpleZoneConfig.Constraints, - []roachpb.ReplicaDescriptor{}, - firstRange, - false, - ) - if _, ok := err.(purgatoryError); ok { - t.Fatalf("expected a non purgatory error, got: %v", err) - } + // Second, test the normal case in which we can allocate to the store. + gossiputil.NewStoreGossiper(g).GossipStores(singleStore, t) + result, err := a.AllocateTarget( + simpleZoneConfig.Constraints, + []roachpb.ReplicaDescriptor{}, + firstRange, + false, + ) + if err != nil { + t.Fatalf("unable to perform allocation: %v", err) + } + if result.Node.NodeID != 1 || result.StoreID != 1 { + t.Errorf("expected NodeID 1 and StoreID 1: %+v", result) + } + + // Finally, set that store to be throttled and ensure we don't send the + // replica to purgatory. + a.storePool.mu.Lock() + storeDetail, ok := a.storePool.mu.storeDetails[singleStore[0].StoreID] + if !ok { + t.Fatalf("store:%d was not found in the store pool", singleStore[0].StoreID) + } + storeDetail.throttledUntil = timeutil.Now().Add(24 * time.Hour) + a.storePool.mu.Unlock() + _, err = a.AllocateTarget( + simpleZoneConfig.Constraints, + []roachpb.ReplicaDescriptor{}, + firstRange, + false, + ) + if _, ok := err.(purgatoryError); ok { + t.Fatalf("expected a non purgatory error, got: %v", err) + } + }) } type testStore struct { @@ -1255,7 +1531,7 @@ func (ts *testStore) rebalance(ots *testStore, bytes int64) { ots.Capacity.Available -= bytes } -func Example_rebalancing() { +func exampleRebalancingCore(useRuleSolver bool) { stopper := stop.NewStopper() defer stopper.Stop() @@ -1282,7 +1558,10 @@ func Example_rebalancing() { stopper, /* deterministic */ true, ) - alloc := MakeAllocator(sp, AllocatorOptions{AllowRebalance: true}) + alloc := MakeAllocator(sp, AllocatorOptions{ + AllowRebalance: true, + UseRuleSolver: useRuleSolver, + }) var wg sync.WaitGroup g.RegisterCallback(gossip.MakePrefixPattern(gossip.KeyStorePrefix), func(_ string, _ roachpb.Value) { wg.Done() }) @@ -1329,12 +1608,15 @@ func Example_rebalancing() { // Next loop through test stores and maybe rebalance. for j := 0; j < len(testStores); j++ { ts := &testStores[j] - target := alloc.RebalanceTarget( + target, err := alloc.RebalanceTarget( config.Constraints{}, []roachpb.ReplicaDescriptor{{NodeID: ts.Node.NodeID, StoreID: ts.StoreID}}, noStore, firstRange, ) + if err != nil { + panic(err) + } if target != nil { testStores[j].rebalance(&testStores[int(target.StoreID)], alloc.randGen.Int63n(1<<20)) } @@ -1364,7 +1646,10 @@ func Example_rebalancing() { } table.Render() fmt.Printf("Total bytes=%d, ranges=%d\n", totBytes, totRanges) +} +func Example_rebalancing() { + exampleRebalancingCore(false /* useRuleSolver */) // Output: // +-----+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+ // | 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 | @@ -1422,3 +1707,64 @@ func Example_rebalancing() { // +-----+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+ // Total bytes=915403982, ranges=1748 } + +func Example_rebalancingWithRuleSolver() { + exampleRebalancingCore(true /* useRuleSolver */) + + // Output: + // +-----+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+ + // | 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 88% | 1 11% | 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% | 0 0% | 0 0% | + // | 2 | 1 32% | 2 24% | 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 21% | 2 20% | + // | 4 | 2 9% | 2 22% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 0 0% | 3 5% | 1 0% | 2 9% | 2 19% | 2 21% | 3 9% | + // | 6 | 3 5% | 2 5% | 0 0% | 0 0% | 0 0% | 0 0% | 6 15% | 1 3% | 0 0% | 0 0% | 0 0% | 0 0% | 2 4% | 7 27% | 3 1% | 2 3% | 2 6% | 2 10% | 2 9% | 3 6% | + // | 8 | 3 4% | 3 5% | 0 0% | 0 0% | 0 0% | 0 0% | 6 9% | 3 2% | 7 8% | 7 6% | 3 7% | 0 0% | 3 4% | 7 17% | 4 3% | 3 2% | 3 6% | 3 9% | 3 6% | 3 5% | + // | 10 | 5 3% | 5 4% | 0 0% | 0 0% | 5 4% | 5 6% | 6 6% | 5 4% | 7 5% | 7 5% | 5 8% | 0 0% | 5 5% | 7 11% | 5 3% | 5 4% | 5 5% | 5 9% | 5 5% | 5 5% | + // | 12 | 6 3% | 6 3% | 5 5% | 15 14% | 6 3% | 6 6% | 6 4% | 6 3% | 7 2% | 7 3% | 6 7% | 3 1% | 6 4% | 7 7% | 6 3% | 6 2% | 6 3% | 6 7% | 6 4% | 6 5% | + // | 14 | 8 4% | 8 4% | 8 6% | 15 10% | 8 2% | 8 4% | 8 5% | 8 4% | 8 3% | 8 2% | 8 6% | 9 5% | 8 3% | 8 5% | 8 4% | 8 3% | 8 4% | 8 7% | 8 4% | 8 6% | + // | 16 | 11 5% | 10 4% | 10 5% | 15 8% | 10 3% | 10 5% | 10 4% | 10 4% | 10 3% | 10 2% | 10 5% | 11 4% | 10 4% | 10 5% | 10 4% | 10 4% | 10 4% | 10 6% | 10 4% | 11 6% | + // | 18 | 13 4% | 12 4% | 12 6% | 15 7% | 12 3% | 12 5% | 12 4% | 12 4% | 12 4% | 12 2% | 12 6% | 13 5% | 12 4% | 12 5% | 12 5% | 12 4% | 13 5% | 13 6% | 12 4% | 13 6% | + // | 20 | 15 4% | 14 4% | 14 6% | 16 6% | 14 3% | 14 5% | 14 5% | 14 4% | 14 4% | 14 3% | 14 6% | 15 4% | 14 4% | 14 4% | 14 5% | 15 4% | 15 5% | 15 6% | 14 4% | 15 5% | + // | 22 | 17 4% | 16 4% | 16 6% | 18 6% | 16 3% | 16 5% | 16 5% | 16 4% | 16 4% | 16 3% | 16 6% | 17 4% | 16 4% | 16 4% | 16 5% | 17 4% | 17 5% | 17 6% | 16 5% | 17 5% | + // | 24 | 19 4% | 18 4% | 18 6% | 20 6% | 18 3% | 18 5% | 18 5% | 18 4% | 18 4% | 18 3% | 18 5% | 19 4% | 18 4% | 18 4% | 18 5% | 19 4% | 19 5% | 19 6% | 18 5% | 19 5% | + // | 26 | 21 4% | 20 4% | 20 5% | 22 5% | 20 3% | 20 5% | 20 4% | 20 4% | 20 4% | 20 4% | 20 5% | 21 5% | 20 4% | 20 4% | 20 5% | 21 3% | 21 5% | 21 6% | 20 5% | 21 5% | + // | 28 | 23 4% | 22 4% | 22 5% | 24 5% | 22 3% | 22 5% | 22 5% | 22 4% | 22 4% | 22 4% | 22 5% | 23 5% | 22 4% | 22 4% | 22 5% | 23 3% | 23 5% | 23 5% | 22 5% | 23 5% | + // | 30 | 25 4% | 24 4% | 24 5% | 26 5% | 24 3% | 24 5% | 24 4% | 24 4% | 24 4% | 24 4% | 24 5% | 25 4% | 24 4% | 24 5% | 24 5% | 25 4% | 25 5% | 25 6% | 24 5% | 25 5% | + // | 32 | 27 4% | 26 5% | 26 5% | 28 5% | 26 3% | 26 5% | 26 4% | 26 4% | 26 4% | 26 4% | 26 5% | 27 4% | 26 4% | 26 4% | 26 5% | 27 4% | 27 5% | 27 5% | 26 5% | 27 4% | + // | 34 | 29 4% | 28 4% | 28 5% | 30 5% | 28 3% | 28 5% | 28 4% | 28 4% | 28 4% | 28 4% | 28 5% | 29 5% | 28 4% | 28 4% | 28 5% | 29 4% | 29 5% | 29 5% | 28 5% | 29 5% | + // | 36 | 31 4% | 30 4% | 30 5% | 32 5% | 30 3% | 30 5% | 30 4% | 30 4% | 30 4% | 30 5% | 30 5% | 31 4% | 30 4% | 30 4% | 30 5% | 31 4% | 31 5% | 31 5% | 30 5% | 31 5% | + // | 38 | 33 4% | 32 4% | 32 5% | 34 5% | 32 4% | 32 4% | 32 4% | 32 4% | 32 4% | 32 5% | 32 5% | 33 4% | 32 5% | 32 4% | 32 5% | 33 4% | 33 6% | 33 5% | 32 5% | 33 5% | + // | 40 | 35 4% | 34 5% | 34 5% | 36 5% | 34 4% | 34 5% | 34 4% | 34 4% | 34 4% | 34 5% | 34 5% | 35 4% | 34 5% | 34 4% | 34 5% | 35 4% | 35 6% | 35 5% | 34 5% | 35 5% | + // | 42 | 37 4% | 36 4% | 36 5% | 38 5% | 36 4% | 36 5% | 36 4% | 36 4% | 36 4% | 36 4% | 36 5% | 37 5% | 36 5% | 36 4% | 36 5% | 37 4% | 37 5% | 37 5% | 36 5% | 37 5% | + // | 44 | 39 4% | 38 4% | 38 5% | 40 5% | 38 4% | 38 5% | 38 4% | 38 4% | 38 4% | 38 4% | 38 5% | 39 4% | 38 5% | 38 4% | 38 5% | 39 4% | 39 5% | 39 5% | 38 5% | 39 5% | + // | 46 | 41 4% | 40 4% | 40 5% | 42 5% | 40 4% | 40 5% | 40 4% | 40 4% | 40 4% | 40 4% | 40 5% | 41 4% | 40 5% | 40 4% | 40 5% | 41 4% | 41 5% | 41 5% | 40 5% | 41 5% | + // | 48 | 43 4% | 42 4% | 42 5% | 44 5% | 42 4% | 42 5% | 42 5% | 42 4% | 42 4% | 42 4% | 42 5% | 43 4% | 42 5% | 42 4% | 42 5% | 43 4% | 43 5% | 43 5% | 42 5% | 43 5% | + // | 50 | 45 4% | 44 4% | 44 5% | 46 5% | 44 4% | 44 5% | 44 5% | 44 4% | 44 4% | 44 4% | 44 5% | 45 4% | 44 5% | 44 4% | 44 4% | 45 4% | 45 5% | 45 5% | 44 5% | 45 5% | + // | 52 | 47 4% | 46 4% | 46 5% | 48 5% | 46 4% | 46 5% | 46 5% | 46 4% | 46 4% | 46 4% | 46 5% | 47 5% | 46 5% | 46 4% | 46 4% | 47 4% | 47 5% | 47 5% | 46 5% | 47 5% | + // | 54 | 49 4% | 48 4% | 48 5% | 50 5% | 48 4% | 48 5% | 48 5% | 48 4% | 48 4% | 48 4% | 48 5% | 49 5% | 48 5% | 48 4% | 48 4% | 49 4% | 49 5% | 49 5% | 48 5% | 49 5% | + // | 56 | 51 4% | 50 4% | 50 5% | 52 5% | 50 4% | 50 5% | 50 5% | 50 5% | 50 4% | 50 4% | 50 5% | 51 5% | 50 5% | 50 4% | 50 4% | 51 4% | 51 5% | 51 5% | 50 5% | 51 5% | + // | 58 | 53 4% | 52 4% | 52 5% | 54 5% | 52 4% | 52 5% | 52 5% | 52 5% | 52 4% | 52 4% | 52 5% | 53 5% | 52 5% | 52 4% | 52 4% | 53 4% | 53 5% | 53 5% | 52 5% | 53 5% | + // | 60 | 55 4% | 54 4% | 54 5% | 56 5% | 54 4% | 54 5% | 54 5% | 54 5% | 54 4% | 54 4% | 54 5% | 55 5% | 54 5% | 54 4% | 54 4% | 55 4% | 55 5% | 55 5% | 54 5% | 55 5% | + // | 62 | 57 4% | 56 4% | 56 5% | 58 5% | 56 4% | 56 5% | 56 5% | 56 5% | 56 4% | 56 4% | 56 5% | 57 5% | 56 5% | 56 4% | 56 5% | 57 4% | 57 5% | 57 5% | 56 5% | 57 5% | + // | 64 | 59 4% | 58 4% | 58 5% | 60 5% | 58 4% | 58 5% | 58 4% | 58 5% | 58 4% | 58 4% | 58 5% | 59 4% | 58 5% | 58 4% | 58 5% | 59 4% | 59 5% | 59 5% | 58 5% | 59 5% | + // | 66 | 61 4% | 60 4% | 60 5% | 62 5% | 60 4% | 60 5% | 60 4% | 60 5% | 60 4% | 60 4% | 60 5% | 61 4% | 60 5% | 60 4% | 60 5% | 61 4% | 61 5% | 61 5% | 60 5% | 61 5% | + // | 68 | 63 4% | 62 4% | 62 5% | 64 5% | 62 4% | 62 5% | 62 4% | 62 5% | 62 4% | 62 4% | 62 5% | 63 4% | 62 5% | 62 4% | 62 5% | 63 4% | 63 5% | 63 5% | 62 5% | 63 5% | + // | 70 | 65 4% | 64 4% | 64 5% | 66 5% | 64 4% | 64 5% | 64 4% | 64 5% | 64 5% | 64 4% | 64 5% | 65 4% | 64 5% | 64 4% | 64 5% | 65 4% | 65 5% | 65 5% | 64 5% | 65 5% | + // | 72 | 67 4% | 66 4% | 66 5% | 68 5% | 66 4% | 66 5% | 66 4% | 66 5% | 66 5% | 66 4% | 66 5% | 67 4% | 66 5% | 66 4% | 66 5% | 67 4% | 67 5% | 67 5% | 66 5% | 67 5% | + // | 74 | 69 4% | 68 4% | 68 5% | 70 5% | 68 4% | 68 5% | 68 4% | 68 5% | 68 5% | 68 4% | 68 5% | 69 4% | 68 5% | 68 4% | 68 5% | 69 4% | 69 5% | 69 5% | 68 5% | 69 5% | + // | 76 | 71 4% | 70 4% | 70 5% | 72 5% | 70 4% | 70 5% | 70 4% | 70 4% | 70 5% | 70 4% | 70 5% | 71 4% | 70 5% | 70 4% | 70 5% | 71 4% | 71 5% | 71 5% | 70 5% | 71 5% | + // | 78 | 73 4% | 72 4% | 72 5% | 74 5% | 72 4% | 72 5% | 72 4% | 72 4% | 72 5% | 72 4% | 72 5% | 73 4% | 72 5% | 72 4% | 72 5% | 73 4% | 73 5% | 73 5% | 72 5% | 73 5% | + // | 80 | 75 4% | 74 4% | 74 5% | 76 5% | 74 4% | 74 5% | 74 4% | 74 5% | 74 5% | 74 4% | 74 5% | 75 4% | 74 5% | 74 4% | 74 5% | 75 4% | 75 5% | 75 5% | 74 5% | 75 5% | + // | 82 | 77 4% | 76 4% | 76 5% | 78 5% | 76 4% | 76 5% | 76 4% | 76 5% | 76 5% | 76 4% | 76 5% | 77 4% | 76 5% | 76 4% | 76 5% | 77 4% | 77 5% | 77 5% | 76 5% | 77 5% | + // | 84 | 79 4% | 78 4% | 78 5% | 80 5% | 78 4% | 78 5% | 78 4% | 78 4% | 78 5% | 78 4% | 78 5% | 79 4% | 78 5% | 78 4% | 78 5% | 79 4% | 79 5% | 79 5% | 78 5% | 79 5% | + // | 86 | 81 4% | 80 4% | 80 5% | 82 5% | 80 4% | 80 5% | 80 4% | 80 4% | 80 4% | 80 4% | 80 5% | 81 4% | 80 5% | 80 4% | 80 5% | 81 4% | 81 5% | 81 5% | 80 5% | 81 5% | + // | 88 | 83 4% | 82 4% | 82 5% | 84 5% | 82 4% | 82 5% | 82 4% | 82 5% | 82 4% | 82 4% | 82 5% | 83 4% | 82 5% | 82 4% | 82 5% | 83 4% | 83 5% | 83 5% | 82 5% | 83 5% | + // | 90 | 85 4% | 84 4% | 84 5% | 86 5% | 84 4% | 84 5% | 84 4% | 84 4% | 84 4% | 84 4% | 84 5% | 85 4% | 84 5% | 84 4% | 84 4% | 85 4% | 85 5% | 85 5% | 84 5% | 85 5% | + // | 92 | 87 4% | 86 4% | 86 5% | 88 5% | 86 4% | 86 5% | 86 4% | 86 4% | 86 4% | 86 4% | 86 5% | 87 4% | 86 5% | 86 4% | 86 4% | 87 4% | 87 5% | 87 5% | 86 5% | 87 5% | + // | 94 | 89 4% | 88 4% | 88 5% | 90 5% | 88 4% | 88 5% | 88 4% | 88 5% | 88 4% | 88 4% | 88 5% | 89 4% | 88 5% | 88 4% | 88 4% | 89 4% | 89 5% | 89 5% | 88 5% | 89 5% | + // | 96 | 91 4% | 90 4% | 90 5% | 92 5% | 90 4% | 90 5% | 90 4% | 90 5% | 90 4% | 90 4% | 90 5% | 91 4% | 90 5% | 90 4% | 90 4% | 91 4% | 91 5% | 91 5% | 90 5% | 91 5% | + // | 98 | 93 4% | 92 4% | 92 5% | 94 5% | 92 4% | 92 5% | 92 4% | 92 4% | 92 4% | 92 4% | 92 5% | 93 4% | 92 5% | 92 4% | 92 4% | 93 4% | 93 5% | 93 5% | 92 5% | 93 5% | + // +-----+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+----------+ + // Total bytes=969243192, ranges=1868 +} diff --git a/pkg/storage/balancer.go b/pkg/storage/balancer.go index dad577d53129..82632a30d6d7 100644 --- a/pkg/storage/balancer.go +++ b/pkg/storage/balancer.go @@ -19,12 +19,10 @@ package storage import ( "bytes" "fmt" - "math" "golang.org/x/net/context" "github.com/cockroachdb/cockroach/pkg/roachpb" - "github.com/cockroachdb/cockroach/pkg/util/envutil" "github.com/cockroachdb/cockroach/pkg/util/log" ) @@ -152,53 +150,6 @@ func (rcb rangeCountBalancer) improve(sl StoreList, excluded nodeIDSet) *roachpb return candidate } -// rebalanceThreshold is the minimum ratio of a store's range surplus to the -// mean range count that permits rebalances away from that store. -var rebalanceThreshold = envutil.EnvOrDefaultFloat("COCKROACH_REBALANCE_THRESHOLD", 0.05) - -func (rangeCountBalancer) shouldRebalance(store roachpb.StoreDescriptor, sl StoreList) bool { - // TODO(peter,bram,cuong): The FractionUsed check seems suspicious. When a - // node becomes fuller than maxFractionUsedThreshold we will always select it - // for rebalancing. This is currently utilized by tests. - maxCapacityUsed := store.Capacity.FractionUsed() >= maxFractionUsedThreshold - - // Rebalance if we're above the rebalance target, which is - // mean*(1+RebalanceThreshold). - target := int32(math.Ceil(sl.candidateCount.mean * (1 + rebalanceThreshold))) - rangeCountAboveTarget := store.Capacity.RangeCount > target - - // Rebalance if the candidate store has a range count above the mean, and - // there exists another store that is underfull: its range count is smaller - // than mean*(1-RebalanceThreshold). - var rebalanceToUnderfullStore bool - if float64(store.Capacity.RangeCount) > sl.candidateCount.mean { - underfullThreshold := int32(math.Floor(sl.candidateCount.mean * (1 - rebalanceThreshold))) - for _, desc := range sl.stores { - if desc.Capacity.RangeCount < underfullThreshold { - rebalanceToUnderfullStore = true - break - } - } - } - - // Require that moving a replica from the given store makes its range count - // converge on the mean range count. This only affects clusters with a - // small number of ranges. - rebalanceConvergesOnMean := rebalanceFromConvergesOnMean(sl, store) - - shouldRebalance := - (maxCapacityUsed || rangeCountAboveTarget || rebalanceToUnderfullStore) && rebalanceConvergesOnMean - if log.V(2) { - log.Infof(context.TODO(), - "%d: should-rebalance=%t: fraction-used=%.2f range-count=%d "+ - "(mean=%.1f, target=%d, fraction-used=%t, above-target=%t, underfull=%t, converges=%t)", - store.StoreID, shouldRebalance, store.Capacity.FractionUsed(), - store.Capacity.RangeCount, sl.candidateCount.mean, target, - maxCapacityUsed, rangeCountAboveTarget, rebalanceToUnderfullStore, rebalanceConvergesOnMean) - } - return shouldRebalance -} - // selectRandom chooses up to count random store descriptors from the given // store list, excluding any stores that are too full to accept more replicas. func selectRandom( diff --git a/pkg/storage/replicate_queue.go b/pkg/storage/replicate_queue.go index 0d1792c806b8..63420bc95de0 100644 --- a/pkg/storage/replicate_queue.go +++ b/pkg/storage/replicate_queue.go @@ -119,12 +119,16 @@ func (rq *replicateQueue) shouldQueue( if lease, _ := repl.getLease(); lease != nil { leaseStoreID = lease.Replica.StoreID } - target := rq.allocator.RebalanceTarget( + target, err := rq.allocator.RebalanceTarget( zone.Constraints, desc.Replicas, leaseStoreID, desc.RangeID, ) + if err != nil { + log.ErrEventf(ctx, "rebalance target failed: %s", err) + return false, 0 + } if log.V(2) { if target != nil { log.Infof(ctx, "%s rebalance target found, enqueuing", repl) @@ -180,7 +184,11 @@ func (rq *replicateQueue) process( log.Event(ctx, "removing a replica") // We require the lease in order to process replicas, so // repl.store.StoreID() corresponds to the lease-holder's store ID. - removeReplica, err := rq.allocator.RemoveTarget(desc.Replicas, repl.store.StoreID()) + removeReplica, err := rq.allocator.RemoveTarget( + zone.Constraints, + desc.Replicas, + repl.store.StoreID(), + ) if err != nil { return err } @@ -212,12 +220,16 @@ func (rq *replicateQueue) process( // // We require the lease in order to process replicas, so // repl.store.StoreID() corresponds to the lease-holder's store ID. - rebalanceStore := rq.allocator.RebalanceTarget( + rebalanceStore, err := rq.allocator.RebalanceTarget( zone.Constraints, desc.Replicas, repl.store.StoreID(), desc.RangeID, ) + if err != nil { + log.ErrEventf(ctx, "rebalance target failed %s", err) + return nil + } if rebalanceStore == nil { log.VEventf(ctx, 1, "no suitable rebalance target") // No action was necessary and no rebalance target was found. Return diff --git a/pkg/storage/rule_solver.go b/pkg/storage/rule_solver.go new file mode 100644 index 000000000000..db0439b2beff --- /dev/null +++ b/pkg/storage/rule_solver.go @@ -0,0 +1,277 @@ +// Copyright 2016 The Cockroach Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or +// implied. See the License for the specific language governing +// permissions and limitations under the License. +// +// Author: Tristan Rice (rice@fn.lc) + +package storage + +import ( + "sort" + + "github.com/cockroachdb/cockroach/pkg/config" + "github.com/cockroachdb/cockroach/pkg/roachpb" +) + +// candidate store for allocation. +type candidate struct { + store roachpb.StoreDescriptor + score float64 +} + +// solveState is used to pass solution state information into a rule. +type solveState struct { + constraints config.Constraints + store roachpb.StoreDescriptor + existing []roachpb.ReplicaDescriptor + sl StoreList + tiers map[roachpb.StoreID]map[string]roachpb.Tier + tierOrder []roachpb.Tier +} + +// rule is a generic rule that can be used to solve a constraint problem. +type rule struct { + weight float64 + // run 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. When calculating the + // best candidate, the rule's weight is used as a multiplier for the + // score. + run func(state solveState) (float64, bool) +} + +// defaultSolverRules is the default rule set to use. +var defaultSolverRules = []rule{ + { + weight: 1.0, + run: ruleReplicasUniqueNodes, + }, + { + weight: 1.0, + run: ruleConstraints, + }, + { + weight: 0.01, + run: ruleCapacity, + }, + { + weight: 0.1, + run: ruleDiversity, + }, +} + +// makeRuleSolver makes a new ruleSolver. The order of the rules is the order in +// which they are run. For optimization purposes, less computationally expensive +// rules should run first to eliminate candidates. +func makeRuleSolver(rules []rule) ruleSolver { + return ruleSolver{ + rules: rules, + } +} + +// ruleSolver solves a set of rules for a store. +type ruleSolver struct { + rules []rule +} + +// 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. +func (rs ruleSolver) Solve( + sl StoreList, c config.Constraints, existing []roachpb.ReplicaDescriptor, +) ([]candidate, error) { + candidates := make([]candidate, 0, len(sl.stores)) + state := solveState{ + constraints: c, + existing: existing, + sl: sl, + tierOrder: canonicalTierOrder(sl), + tiers: storeTierMap(sl), + } + + for _, store := range sl.stores { + state.store = store + if cand, ok := rs.computeCandidate(state); ok { + candidates = append(candidates, cand) + } + } + sort.Sort(byScoreDescending(candidates)) + return candidates, nil +} + +// 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(state solveState) (candidate, bool) { + var totalScore float64 + for _, rule := range rs.rules { + score, valid := rule.run(state) + if !valid { + return candidate{}, false + } + totalScore += score * rule.weight + } + return candidate{store: state.store, score: totalScore}, true +} + +// ruleReplicasUniqueNodes returns true iff no existing replica is present on +// the candidate's node. +func ruleReplicasUniqueNodes(state solveState) (float64, bool) { + for _, r := range state.existing { + if r.NodeID == state.store.Node.NodeID { + return 0, false + } + } + return 0, true +} + +// storeHasConstraint returns whether a store descriptor attributes or locality +// matches the key value pair in the constraint. +func storeHasConstraint(store roachpb.StoreDescriptor, c config.Constraint) bool { + if c.Key == "" { + for _, attrs := range []roachpb.Attributes{store.Attrs, store.Node.Attrs} { + for _, attr := range attrs.Attrs { + if attr == c.Value { + return true + } + } + } + } else { + for _, tier := range store.Node.Locality.Tiers { + if c.Key == tier.Key && c.Value == tier.Value { + return true + } + } + } + return false +} + +// ruleConstraints returns true iff all required and prohibited constraints are +// satisfied. Stores with more positive constraints return higher scores. +func ruleConstraints(state solveState) (float64, bool) { + if len(state.constraints.Constraints) == 0 { + return 0, true + } + matched := 0 + for _, c := range state.constraints.Constraints { + hasConstraint := storeHasConstraint(state.store, c) + switch { + case c.Type == config.Constraint_POSITIVE && hasConstraint: + matched++ + case c.Type == config.Constraint_REQUIRED && !hasConstraint: + return 0, false + case c.Type == config.Constraint_PROHIBITED && hasConstraint: + return 0, false + } + } + + return float64(matched) / float64(len(state.constraints.Constraints)), true +} + +// ruleDiversity returns higher scores for stores with the fewest locality tiers +// in common with already existing replicas. It always returns true. +func ruleDiversity(state solveState) (float64, bool) { + storeTiers := state.tiers[state.store.StoreID] + var maxScore, score float64 + for i, tier := range state.tierOrder { + storeTier, ok := storeTiers[tier.Key] + if !ok { + continue + } + tierScore := 1 / (float64(i) + 1) + for _, existing := range state.existing { + existingTier, ok := state.tiers[existing.StoreID][tier.Key] + if ok && existingTier.Value != storeTier.Value { + score += tierScore + } + maxScore += tierScore + } + } + if maxScore == 0 { + return 0, true + } + return score / maxScore, true +} + +// ruleCapacity returns true iff a new replica won't overfill the store. The +// score returned is inversely proportional to the number of ranges on the +// candidate store, with the most empty nodes having the highest scores. +// TODO(bram): consider splitting this into two rules. +func ruleCapacity(state solveState) (float64, bool) { + // Don't overfill stores. + if state.store.Capacity.FractionUsed() > maxFractionUsedThreshold { + return 0, false + } + + return 1 / float64(state.store.Capacity.RangeCount+1), true +} + +// canonicalTierOrder returns the most common key at each tier level in +// order from the most important tier down to the least important, cutting off +// if the majority of stores have no tier at a level. +func canonicalTierOrder(sl StoreList) []roachpb.Tier { + maxTierCount := 0 + for _, store := range sl.stores { + if count := len(store.Node.Locality.Tiers); maxTierCount < count { + maxTierCount = count + } + } + + // Might have up to maxTierCount of tiers. + tiers := make([]roachpb.Tier, 0, maxTierCount) + for i := 0; i < maxTierCount; i++ { + // At each tier, count the number of occurrences of each key. + counts := map[string]int{} + maxKey := "" + for _, store := range sl.stores { + key := "" + if i < len(store.Node.Locality.Tiers) { + key = store.Node.Locality.Tiers[i].Key + } + counts[key]++ + if counts[key] > counts[maxKey] { + maxKey = key + } + } + // Don't add the tier if most nodes don't have that many tiers. + if maxKey != "" { + tiers = append(tiers, roachpb.Tier{Key: maxKey}) + } else { + break + } + } + return tiers +} + +// storeTierMap indexes a store list so it can be used to look up the locality +// tier value from store ID and tier key. +func storeTierMap(sl StoreList) map[roachpb.StoreID]map[string]roachpb.Tier { + m := map[roachpb.StoreID]map[string]roachpb.Tier{} + for _, store := range sl.stores { + sm := map[string]roachpb.Tier{} + m[store.StoreID] = sm + for _, tier := range store.Node.Locality.Tiers { + sm[tier.Key] = tier + } + } + return m +} + +// byScore implements sort.Interface to sort by scores descending from highest +// to lowest. +type byScoreDescending []candidate + +var _ sort.Interface = byScoreDescending(nil) + +func (c byScoreDescending) Len() int { return len(c) } +func (c byScoreDescending) Less(i, j int) bool { return c[i].score > c[j].score } +func (c byScoreDescending) 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 new file mode 100644 index 000000000000..a1bcbcaf937d --- /dev/null +++ b/pkg/storage/rule_solver_test.go @@ -0,0 +1,365 @@ +// Copyright 2016 The Cockroach Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or +// implied. See the License for the specific language governing +// permissions and limitations under the License. +// +// Author: Tristan Rice (rice@fn.lc) + +package storage + +import ( + "reflect" + "sort" + "testing" + + "github.com/cockroachdb/cockroach/pkg/config" + "github.com/cockroachdb/cockroach/pkg/roachpb" + "github.com/cockroachdb/cockroach/pkg/util/leaktest" +) + +type byScoreAndID []candidate + +func (c byScoreAndID) Len() int { return len(c) } +func (c byScoreAndID) Less(i, j int) bool { + if c[i].score == c[j].score { + return c[i].store.StoreID < c[j].store.StoreID + } + return c[i].score > c[j].score +} +func (c byScoreAndID) Swap(i, j int) { c[i], c[j] = c[j], c[i] } + +// 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 */ false, + ) + defer stopper.Stop() + + storeRSa := roachpb.StoreID(1) // datacenter (us), rack, slot + a + storeRab := roachpb.StoreID(2) // datacenter (us), rack + a,b + storeFRabc := roachpb.StoreID(3) // datacenter (us), floor, rack + a,b,c + storeDead := roachpb.StoreID(4) + storeEurope := roachpb.StoreID(5) // datacenter (eur), rack + + mockStorePool(storePool, []roachpb.StoreID{storeRSa, storeRab, storeFRabc, 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[storeRSa].desc.Attrs.Attrs = []string{"a"} + storePool.mu.storeDetails[storeRSa].desc.Node.Locality.Tiers = tierSetup("us", "", "1", "5") + storePool.mu.storeDetails[storeRSa].desc.Capacity = capacitySetup(1, 99) + + storePool.mu.storeDetails[storeRab].desc.Attrs.Attrs = []string{"a", "b"} + storePool.mu.storeDetails[storeRab].desc.Node.Locality.Tiers = tierSetup("us", "", "1", "") + storePool.mu.storeDetails[storeRab].desc.Capacity = capacitySetup(100, 0) + + storePool.mu.storeDetails[storeFRabc].desc.Attrs.Attrs = []string{"a", "b", "c"} + storePool.mu.storeDetails[storeFRabc].desc.Node.Locality.Tiers = tierSetup("us", "1", "2", "") + storePool.mu.storeDetails[storeFRabc].desc.Capacity = capacitySetup(50, 50) + + storePool.mu.storeDetails[storeEurope].desc.Node.Locality.Tiers = tierSetup("eur", "", "1", "") + storePool.mu.storeDetails[storeEurope].desc.Capacity = capacitySetup(60, 40) + + storePool.mu.Unlock() + + testCases := []struct { + name string + rule rule + c config.Constraints + existing []roachpb.ReplicaDescriptor + expected []roachpb.StoreID + }{ + { + name: "no constraints or rules", + expected: []roachpb.StoreID{storeRSa, 2, 3, 5}, + }, + { + name: "white list rule", + rule: rule{ + weight: 1, + run: func(state solveState) (float64, bool) { + switch state.store.StoreID { + case storeRSa: + return 0, true + case storeFRabc: + return 1, true + default: + return 0, false + } + }, + }, + expected: []roachpb.StoreID{3, storeRSa}, + }, + { + name: "ruleReplicasUniqueNodes - 2 available nodes", + rule: rule{weight: 1, run: ruleReplicasUniqueNodes}, + existing: []roachpb.ReplicaDescriptor{ + {NodeID: roachpb.NodeID(storeRSa)}, + {NodeID: roachpb.NodeID(storeFRabc)}, + }, + expected: []roachpb.StoreID{storeRab, 5}, + }, + { + name: "ruleReplicasUniqueNodes - 0 available nodes", + rule: rule{weight: 1, run: ruleReplicasUniqueNodes}, + existing: []roachpb.ReplicaDescriptor{ + {NodeID: roachpb.NodeID(storeRSa)}, + {NodeID: roachpb.NodeID(storeRab)}, + {NodeID: roachpb.NodeID(storeFRabc)}, + {NodeID: roachpb.NodeID(storeEurope)}, + }, + expected: nil, + }, + { + name: "ruleConstraints - required constraints", + rule: rule{weight: 1, run: ruleConstraints}, + c: config.Constraints{ + Constraints: []config.Constraint{ + {Value: "b", Type: config.Constraint_REQUIRED}, + }, + }, + expected: []roachpb.StoreID{storeRab, storeFRabc}, + }, + { + name: "ruleConstraints - required locality constraints", + rule: rule{weight: 1, run: ruleConstraints}, + c: config.Constraints{ + Constraints: []config.Constraint{ + {Key: "datacenter", Value: "us", Type: config.Constraint_REQUIRED}, + }, + }, + expected: []roachpb.StoreID{storeRSa, storeRab, storeFRabc}, + }, + { + name: "ruleConstraints - prohibited constraints", + rule: rule{weight: 1, run: ruleConstraints}, + c: config.Constraints{ + Constraints: []config.Constraint{ + {Value: "b", Type: config.Constraint_PROHIBITED}, + }, + }, + expected: []roachpb.StoreID{storeRSa, storeEurope}, + }, + { + name: "ruleConstraints - prohibited locality constraints", + rule: rule{weight: 1, run: ruleConstraints}, + c: config.Constraints{ + Constraints: []config.Constraint{ + {Key: "datacenter", Value: "us", Type: config.Constraint_PROHIBITED}, + }, + }, + expected: []roachpb.StoreID{storeEurope}, + }, + { + name: "ruleConstraints - positive constraints", + rule: rule{weight: 1, run: ruleConstraints}, + c: config.Constraints{ + Constraints: []config.Constraint{ + {Value: "a"}, + {Value: "b"}, + {Value: "c"}, + }, + }, + expected: []roachpb.StoreID{storeFRabc, storeRab, storeRSa, storeEurope}, + }, + { + name: "ruleConstraints - positive locality constraints", + rule: rule{weight: 1, run: ruleConstraints}, + c: config.Constraints{ + Constraints: []config.Constraint{ + {Key: "datacenter", Value: "eur"}, + }, + }, + expected: []roachpb.StoreID{5, storeRSa, storeRab, storeFRabc}, + }, + { + name: "ruleDiversity - no existing replicas", + rule: rule{weight: 1, run: ruleDiversity}, + existing: nil, + expected: []roachpb.StoreID{storeRSa, storeRab, storeFRabc, 5}, + }, + { + name: "ruleDiversity - one existing replicas", + rule: rule{weight: 1, run: ruleDiversity}, + existing: []roachpb.ReplicaDescriptor{ + {StoreID: storeRSa}, + }, + expected: []roachpb.StoreID{storeEurope, storeFRabc, storeRSa, storeRab}, + }, + { + name: "ruleCapacity", + rule: rule{weight: 1, run: ruleCapacity}, + expected: []roachpb.StoreID{storeRab, storeEurope, storeFRabc}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + var rules []rule + if tc.rule.run != nil { + rules = []rule{tc.rule} + } + solver := makeRuleSolver(rules) + sl, _, _ := storePool.getStoreList(roachpb.RangeID(0)) + candidates, err := solver.Solve(sl, tc.c, tc.existing) + if err != nil { + t.Fatal(err) + } + sort.Sort(byScoreAndID(candidates)) + if len(candidates) != len(tc.expected) { + t.Fatalf("length of %+v should match %+v", candidates, tc.expected) + } + for i, expected := range tc.expected { + if actual := candidates[i].store.StoreID; actual != expected { + t.Errorf("candidates[%d].store.StoreID = %d; not %d; %+v", + i, actual, expected, candidates) + } + } + }) + } +} + +func TestCanonicalTierOrder(t *testing.T) { + defer leaktest.AfterTest(t)() + + testCases := []struct { + name string + tiersPerStore [][]roachpb.Tier + expected []roachpb.Tier + }{ + { + "no tiers at all", + nil, + []roachpb.Tier{}, + }, + { + "one store with two empty tiers", + [][]roachpb.Tier{nil, nil}, + []roachpb.Tier{}, + }, + { + "one store with three tiers", + [][]roachpb.Tier{ + { + {Key: "a"}, + {Key: "b"}, + {Key: "c"}, + }, + }, + []roachpb.Tier{ + {Key: "a"}, + {Key: "b"}, + {Key: "c"}, + }, + }, + { + "3 stores with the same tiers, one with an extra one", + [][]roachpb.Tier{ + { + {Key: "a"}, + {Key: "b"}, + {Key: "c"}, + }, + { + {Key: "a"}, + {Key: "b"}, + {Key: "c"}, + }, + { + {Key: "b"}, + {Key: "c"}, + {Key: "a"}, + {Key: "d"}, + }, + }, + []roachpb.Tier{ + {Key: "a"}, + {Key: "b"}, + {Key: "c"}, + }, + }, + { + "two stores with completely different tiers", + [][]roachpb.Tier{ + { + {Key: "a"}, + {Key: "b"}, + {Key: "c"}, + }, + { + {Key: "e"}, + {Key: "f"}, + {Key: "g"}, + }, + }, + []roachpb.Tier{ + {Key: "a"}, + {Key: "b"}, + {Key: "c"}, + }, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + var descriptors []roachpb.StoreDescriptor + for _, tiers := range tc.tiersPerStore { + descriptors = append(descriptors, roachpb.StoreDescriptor{ + Node: roachpb.NodeDescriptor{ + Locality: roachpb.Locality{Tiers: tiers}, + }, + }) + } + + sl := makeStoreList(descriptors) + if actual := canonicalTierOrder(sl); !reflect.DeepEqual(actual, tc.expected) { + t.Errorf("canonicalTierOrder(%+v) = %+v; not %+v", + tc.tiersPerStore, actual, tc.expected) + } + }) + } +} diff --git a/pkg/storage/simulation/range.go b/pkg/storage/simulation/range.go index 94be037c5948..0effcf37163c 100644 --- a/pkg/storage/simulation/range.go +++ b/pkg/storage/simulation/range.go @@ -131,7 +131,7 @@ func (r *Range) getAllocateTarget() (roachpb.StoreID, error) { func (r *Range) getRemoveTarget() (roachpb.StoreID, error) { // Pass in an invalid store ID since we don't consider range leases as part // of the simulator. - removeStore, err := r.allocator.RemoveTarget(r.desc.Replicas, roachpb.StoreID(-1)) + removeStore, err := r.allocator.RemoveTarget(r.zone.Constraints, r.desc.Replicas, roachpb.StoreID(-1)) if err != nil { return 0, err } @@ -142,12 +142,15 @@ func (r *Range) getRemoveTarget() (roachpb.StoreID, error) { // candidate to add a replica for rebalancing. Returns true only if a target is // found. func (r *Range) getRebalanceTarget(storeID roachpb.StoreID) (roachpb.StoreID, bool) { - rebalanceTarget := r.allocator.RebalanceTarget( + rebalanceTarget, err := r.allocator.RebalanceTarget( r.zone.Constraints, r.desc.Replicas, storeID, r.desc.RangeID, ) + if err != nil { + panic(err) + } if rebalanceTarget == nil { return 0, false } diff --git a/pkg/storage/store.go b/pkg/storage/store.go index 506a9c7fc3c3..367a3f5e277d 100644 --- a/pkg/storage/store.go +++ b/pkg/storage/store.go @@ -102,6 +102,8 @@ var storeSchedulerConcurrency = envutil.EnvOrDefaultInt( var enablePreVote = envutil.EnvOrDefaultBool( "COCKROACH_ENABLE_PREVOTE", false) +var enableRuleSolver = envutil.EnvOrDefaultBool("COCKROACH_ENABLE_RULE_SOLVER", false) + // RaftElectionTimeout returns the raft election timeout, as computed // from the specified tick interval and number of election timeout // ticks. If raftElectionTimeoutTicks is 0, uses the value of @@ -742,6 +744,8 @@ func (sc *StoreConfig) SetDefaults() { if sc.RangeLeaseRenewalDuration == 0 { sc.RangeLeaseRenewalDuration = rangeLeaseRenewalDuration } + + sc.AllocatorOptions.UseRuleSolver = enableRuleSolver } // NewStore returns a new instance of a store.