Merge pull request #20 from vimeo/consistenthash_cleanup

Consistenthash: clarify naming, add lint-happy-making doc-comments & replace int with uint32

consistenthash/consistenthash.go

@@ -23,59 +23,66 @@ import (
 	"strconv"
 )
 
+// Hash maps the data to a uint32 hash-ring
 type Hash func(data []byte) uint32
 
+// Map tracks segments in a hash-ring, mapped to specific keys.
 type Map struct {
-	hash     Hash
-	replicas int
-	keys     []int // Sorted
-	hashMap  map[int]string
+	hash       Hash
+	segsPerKey int
+	keyHashes  []uint32 // Sorted
+	hashMap    map[uint32]string
+	keys       map[string]struct{}
 }
 
-func New(replicas int, fn Hash) *Map {
+// New constructs a new consistenthash hashring, with segsPerKey segments per added key.
+func New(segsPerKey int, fn Hash) *Map {
 	m := &Map{
-		replicas: replicas,
-		hash:     fn,
-		hashMap:  make(map[int]string),
+		segsPerKey: segsPerKey,
+		hash:       fn,
+		hashMap:    make(map[uint32]string),
+		keys:       make(map[string]struct{}),
 	}
 	if m.hash == nil {
 		m.hash = crc32.ChecksumIEEE
 	}
 	return m
 }
 
-// Returns true if there are no items available.
+// IsEmpty returns true if there are no items available.
 func (m *Map) IsEmpty() bool {
-	return len(m.keys) == 0
+	return len(m.keyHashes) == 0
 }
 
-// Adds some keys to the hash.
+// Add adds some keys to the hashring, establishing ownership of segsPerKey
+// segments.
 func (m *Map) Add(keys ...string) {
 	for _, key := range keys {
-		for i := 0; i < m.replicas; i++ {
-			hash := int(m.hash([]byte(strconv.Itoa(i) + key)))
-			m.keys = append(m.keys, hash)
+		m.keys[key] = struct{}{}
+		for i := 0; i < m.segsPerKey; i++ {
+			hash := m.hash([]byte(strconv.Itoa(i) + key))
+			m.keyHashes = append(m.keyHashes, hash)
 			m.hashMap[hash] = key
 		}
 	}
-	sort.Ints(m.keys)
+	sort.Slice(m.keyHashes, func(i, j int) bool { return m.keyHashes[i] < m.keyHashes[j] })
 }
 
-// Gets the closest item in the hash to the provided key.
+// Get gets the closest item in the hash to the provided key.
 func (m *Map) Get(key string) string {
 	if m.IsEmpty() {
 		return ""
 	}
 
-	hash := int(m.hash([]byte(key)))
+	hash := m.hash([]byte(key))
 
 	// Binary search for appropriate replica.
-	idx := sort.Search(len(m.keys), func(i int) bool { return m.keys[i] >= hash })
+	idx := sort.Search(len(m.keyHashes), func(i int) bool { return m.keyHashes[i] >= hash })
 
 	// Means we have cycled back to the first replica.
-	if idx == len(m.keys) {
+	if idx == len(m.keyHashes) {
 		idx = 0
 	}
 
-	return m.hashMap[m.keys[idx]]
+	return m.hashMap[m.keyHashes[idx]]
 }

consistenthash/consistenthash_test.go

@@ -86,25 +86,33 @@ func TestConsistency(t *testing.T) {
 
 }
 
-func BenchmarkGet8(b *testing.B)   { benchmarkGet(b, 8) }
-func BenchmarkGet32(b *testing.B)  { benchmarkGet(b, 32) }
-func BenchmarkGet128(b *testing.B) { benchmarkGet(b, 128) }
-func BenchmarkGet512(b *testing.B) { benchmarkGet(b, 512) }
-
-func benchmarkGet(b *testing.B, shards int) {
-
-	hash := New(50, nil)
-
-	var buckets []string
-	for i := 0; i < shards; i++ {
-		buckets = append(buckets, fmt.Sprintf("shard-%d", i))
-	}
-
-	hash.Add(buckets...)
-
-	b.ResetTimer()
-
-	for i := 0; i < b.N; i++ {
-		hash.Get(buckets[i&(shards-1)])
+func BenchmarkGet(b *testing.B) {
+	for _, itbl := range []struct {
+		segsPerKey int
+		shards     int
+	}{
+		{segsPerKey: 50, shards: 8},
+		{segsPerKey: 50, shards: 32},
+		{segsPerKey: 50, shards: 128},
+		{segsPerKey: 50, shards: 512},
+	} {
+		tbl := itbl
+		b.Run(fmt.Sprintf("segs%d-shards%d", tbl.segsPerKey, tbl.shards), func(b *testing.B) {
+
+			hash := New(tbl.segsPerKey, nil)
+
+			var buckets []string
+			for i := 0; i < tbl.shards; i++ {
+				buckets = append(buckets, fmt.Sprintf("shard-%d", i))
+			}
+
+			hash.Add(buckets...)
+
+			b.ResetTimer()
+
+			for i := 0; i < b.N; i++ {
+				hash.Get(buckets[i&(tbl.shards-1)])
+			}
+		})
 	}
 }