Skip to content

Commit

Permalink
util/shuffle, kv, storage: Adds a shuffle package for randomizing lists.
Browse files Browse the repository at this point in the history
Also, add a shuffle to the store pool's getStoreList.

Relying on the range over a map to ensure randomness is a little too subtle.  This is much more explisit.

Part of cockroachdb#10275
  • Loading branch information
BramGruneir committed Nov 11, 2016
1 parent c67ee45 commit df54661
Show file tree
Hide file tree
Showing 7 changed files with 137 additions and 46 deletions.
5 changes: 3 additions & 2 deletions pkg/kv/dist_sender.go
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/util/hlc"
"github.com/cockroachdb/cockroach/pkg/util/log"
"github.com/cockroachdb/cockroach/pkg/util/retry"
"github.com/cockroachdb/cockroach/pkg/util/shuffle"
"github.com/cockroachdb/cockroach/pkg/util/timeutil"
"github.com/cockroachdb/cockroach/pkg/util/tracing"
)
Expand Down Expand Up @@ -246,7 +247,7 @@ func (ds *DistSender) RangeLookup(
Reverse: useReverseScan,
})
replicas := newReplicaSlice(ds.gossip, desc)
replicas.Shuffle()
shuffle.Shuffle(replicas)
br, err := ds.sendRPC(ctx, desc.RangeID, replicas, ba)
if err != nil {
return nil, nil, roachpb.NewError(err)
Expand Down Expand Up @@ -288,7 +289,7 @@ func (ds *DistSender) optimizeReplicaOrder(replicas ReplicaSlice) {
nodeDesc := ds.getNodeDescriptor()
// If we don't know which node we're on, don't optimize anything.
if nodeDesc == nil {
replicas.Shuffle()
shuffle.Shuffle(replicas)
return
}
// Sort replicas by attribute affinity (if any), which we treat as a stand-in
Expand Down
29 changes: 9 additions & 20 deletions pkg/kv/replica_slice.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,12 @@
package kv

import (
"math/rand"

"golang.org/x/net/context"

"github.com/cockroachdb/cockroach/pkg/gossip"
"github.com/cockroachdb/cockroach/pkg/roachpb"
"github.com/cockroachdb/cockroach/pkg/util/log"
"github.com/cockroachdb/cockroach/pkg/util/shuffle"
)

// ReplicaInfo extends the Replica structure with the associated node
Expand Down Expand Up @@ -65,10 +64,13 @@ func newReplicaSlice(gossip *gossip.Gossip, desc *roachpb.RangeDescriptor) Repli
return replicas
}

// Swap interchanges the replicas stored at the given indices.
func (rs ReplicaSlice) Swap(i, j int) {
rs[i], rs[j] = rs[j], rs[i]
}
// ReplicaSlice implements shuffle.Interface.

// Len returns the total number of replicas in the slice.
func (rs ReplicaSlice) Len() int { return len(rs) }

// Swap swaps the replicas with indexes i and j.
func (rs ReplicaSlice) Swap(i, j int) { rs[i], rs[j] = rs[j], rs[i] }

// FindReplica returns the index of the replica which matches the specified store
// ID. If no replica matches, -1 is returned.
Expand Down Expand Up @@ -116,7 +118,7 @@ func (rs ReplicaSlice) SortByCommonAttributePrefix(attrs []string) int {
}
}
if topIndex < len(rs)-1 {
rs.randPerm(firstNotOrdered, topIndex, rand.Intn)
shuffle.Shuffle(rs[firstNotOrdered : topIndex+1])
}
if firstNotOrdered == 0 {
return bucket
Expand All @@ -138,16 +140,3 @@ func (rs ReplicaSlice) MoveToFront(i int) {
copy(rs[1:], rs[:i])
rs[0] = front
}

// Shuffle randomizes the order of the replicas.
func (rs ReplicaSlice) Shuffle() {
rs.randPerm(0, len(rs)-1, rand.Intn)
}

func (rs ReplicaSlice) randPerm(startIndex int, topIndex int, intnFn func(int) int) {
length := topIndex - startIndex + 1
for i := 1; i < length; i++ {
j := intnFn(i + 1)
rs.Swap(startIndex+i, startIndex+j)
}
}
19 changes: 0 additions & 19 deletions pkg/kv/replica_slice_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
package kv

import (
"math/rand"
"reflect"
"testing"

Expand Down Expand Up @@ -115,21 +114,3 @@ func TestReplicaSetMoveToFront(t *testing.T) {
t.Errorf("expected order %s, got %s", exp, stores)
}
}

func verifyRandPermOrdering(startIndex int, topIndex int, exp []roachpb.StoreID, t *testing.T) {
r := rand.New(rand.NewSource(0))
rs := createReplicaSlice()
rs.randPerm(startIndex, topIndex, r.Intn)
if stores := getStores(rs); !reflect.DeepEqual(stores, exp) {
t.Errorf("expected order %s, got %s", exp, stores)
}
}

func TestReplicaSetRandPerm(t *testing.T) {
defer leaktest.AfterTest(t)()
verifyRandPermOrdering(2, 2, []roachpb.StoreID{1, 2, 3, 4, 5}, t)
verifyRandPermOrdering(3, 4, []roachpb.StoreID{1, 2, 3, 5, 4}, t)
verifyRandPermOrdering(0, 2, []roachpb.StoreID{3, 1, 2, 4, 5}, t)
verifyRandPermOrdering(1, 3, []roachpb.StoreID{1, 4, 2, 3, 5}, t)
verifyRandPermOrdering(0, 4, []roachpb.StoreID{3, 5, 2, 1, 4}, t)
}
11 changes: 6 additions & 5 deletions pkg/storage/store.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ import (
"fmt"
"io"
"math"
"math/rand"
"runtime"
"sync"
"sync/atomic"
Expand Down Expand Up @@ -52,6 +51,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/util/log"
"github.com/cockroachdb/cockroach/pkg/util/metric"
"github.com/cockroachdb/cockroach/pkg/util/retry"
"github.com/cockroachdb/cockroach/pkg/util/shuffle"
"github.com/cockroachdb/cockroach/pkg/util/stop"
"github.com/cockroachdb/cockroach/pkg/util/syncutil"
"github.com/cockroachdb/cockroach/pkg/util/timeutil"
Expand Down Expand Up @@ -260,6 +260,10 @@ type storeReplicaVisitor struct {
visited int // Number of visited ranges, -1 before first call to Visit()
}

// storeReplicaVisitor implements the shuffle.Interface.
func (rs storeReplicaVisitor) Len() int { return len(rs.repls) }
func (rs storeReplicaVisitor) Swap(i, j int) { rs.repls[i], rs.repls[j] = rs.repls[j], rs.repls[i] }

func newStoreReplicaVisitor(store *Store) *storeReplicaVisitor {
return &storeReplicaVisitor{
store: store,
Expand Down Expand Up @@ -287,10 +291,7 @@ func (rs *storeReplicaVisitor) Visit(visitor func(*Replica) bool) {
//
// TODO(peter): Re-evaluate whether this is necessary after we allow
// rebalancing away from the leaseholder. See TestRebalance_3To5Small.
for i := 1; i < len(rs.repls); i++ {
j := rand.Intn(i + 1)
rs.repls[i], rs.repls[j] = rs.repls[j], rs.repls[i]
}
shuffle.Shuffle(rs)

rs.visited = 0
for _, repl := range rs.repls {
Expand Down
3 changes: 3 additions & 0 deletions pkg/storage/store_pool.go
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/util/envutil"
"github.com/cockroachdb/cockroach/pkg/util/hlc"
"github.com/cockroachdb/cockroach/pkg/util/log"
"github.com/cockroachdb/cockroach/pkg/util/shuffle"
"github.com/cockroachdb/cockroach/pkg/util/stop"
"github.com/cockroachdb/cockroach/pkg/util/syncutil"
"github.com/cockroachdb/cockroach/pkg/util/timeutil"
Expand Down Expand Up @@ -524,6 +525,8 @@ func (sp *StorePool) getStoreList(rangeID roachpb.RangeID) (StoreList, int, int)

if sp.deterministic {
sort.Sort(storeIDs)
} else {
shuffle.Shuffle(storeIDs)
}

var aliveStoreCount int
Expand Down
36 changes: 36 additions & 0 deletions pkg/util/shuffle/shuffle.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
// 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.

package shuffle

import "math/rand"

// Interface for shuffle. When it is satisfied, a collection can be shuffled by
// the routines in this package. The methods require that the elements of the
// collection be enumerable by an integer index. This interface is similar to
// sort.Interface.
type Interface interface {
// Len is the number of elements in the collection.
Len() int
// Swap swaps the elements with indexes i and j.
Swap(i, j int)
}

// Shuffle randomizes the order of the array.
func Shuffle(data Interface) {
n := data.Len()
for i := 1; i < n; i++ {
data.Swap(i, rand.Intn(i+1))
}
}
80 changes: 80 additions & 0 deletions pkg/util/shuffle/shuffle_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
// 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.

package shuffle

import (
"math/rand"
"reflect"
"testing"

"github.com/cockroachdb/cockroach/pkg/util/leaktest"
)

type testSlice []int

// testSlice implements shuffle.Interface.
func (ts testSlice) Len() int { return len(ts) }
func (ts testSlice) Swap(i, j int) { ts[i], ts[j] = ts[j], ts[i] }

func TestShuffle(t *testing.T) {
defer leaktest.AfterTest(t)()
rand.Seed(0)

verify := func(original, expected testSlice) {
Shuffle(original)
if !reflect.DeepEqual(original, expected) {
t.Errorf("expected %v, got %v", expected, original)
}
}

ts := testSlice{}
verify(ts, testSlice{})
verify(ts, testSlice{})

ts = testSlice{1}
verify(ts, testSlice{1})
verify(ts, testSlice{1})

ts = testSlice{1, 2}
verify(ts, testSlice{2, 1})
verify(ts, testSlice{1, 2})

ts = testSlice{1, 2, 3}
verify(ts, testSlice{1, 3, 2})
verify(ts, testSlice{1, 2, 3})
verify(ts, testSlice{1, 2, 3})
verify(ts, testSlice{3, 1, 2})

ts = testSlice{1, 2, 3, 4, 5}
verify(ts, testSlice{2, 1, 3, 5, 4})
verify(ts, testSlice{4, 2, 1, 5, 3})
verify(ts, testSlice{1, 4, 2, 3, 5})
verify(ts, testSlice{2, 5, 4, 1, 3})
verify(ts, testSlice{4, 2, 3, 1, 5})

verify(ts[2:2], testSlice{})
verify(ts[0:0], testSlice{})
verify(ts[5:5], testSlice{})
verify(ts[3:5], testSlice{1, 5})
verify(ts[3:5], testSlice{5, 1})
verify(ts[0:2], testSlice{4, 2})
verify(ts[0:2], testSlice{2, 4})
verify(ts[1:4], testSlice{3, 5, 4})
verify(ts[1:4], testSlice{5, 4, 3})
verify(ts[0:4], testSlice{4, 5, 2, 3})
verify(ts[0:4], testSlice{2, 4, 3, 5})

verify(ts, testSlice{1, 3, 4, 2, 5})
}

0 comments on commit df54661

Please sign in to comment.