Added prefix match for consistent hash

consistenthash/consistenthash.go

@@ -19,24 +19,78 @@ package consistenthash
 
 import (
  "hash/crc32"
+ "math/bits"
  "sort"
  "strconv"
 )
 
 type Hash func(data []byte) uint32
 
+const defaultHashExpansion = 6
+
 type Map struct {
- hash Hash
+ // Inputs
+
+ // hash is the hash function that will be applied to both added
+ // keys and fetched keys
+ hash Hash
+
+ // replicas is the number of virtual nodes that will be inserted
+ // into the consistent hash ring for each key added
  replicas int
- keys []int // Sorted
- hashMap map[int]string
+
+ // prefixTableExpansion is the multiple of virtual nodes that
+ // will be inserted into the internal hash table for O(1) lookups.
+ prefixTableExpansion int
+
+ // Internal data
+
+ // keys is the hash of the virtual nodes, sorted by hash value
+ keys []int // Sorted
+
+ // hashMap maps the hashed keys back to the input strings.
+ // Note that all virtual nodes will map back to the same input
+ // string
+ hashMap map[int]string
+
+ // prefixShift is the number of bits an input hash should
+ // be right-shifted to act as a lookup in the prefixTable
+ prefixShift uint32
+
+ // prefixTable is a map of the most significant bits of
+ // a hash value to output all hashes with that prefix
+ // map to. If the result is ambiguous (i.e. there is a
+ // hash range split within this prefix) the value will
+ // be blank and we should fall back to a binary search
+ // through keys to find the exact output
+ prefixTable []string
 }
 
+// New returns a blank consistent hash ring that will return
+// the key whose hash comes next after the hash of the input to
+// Map.Get.
+// Increasing the number of replicas will improve the smoothness
+// of the hash ring and reduce the data moved when adding/removing
+// nodes, at the cost of more memory.
 func New(replicas int, fn Hash) *Map {
+ return NewConsistentHash(replicas, defaultHashExpansion, fn)
+}
+
+// NewConsistentHash returns a blank consistent hash ring that will return
+// the key whose hash comes next after the hash of the input to
+// Map.Get.
+// Increasing the number of replicas will improve the smoothness
+// of the hash ring and reduce the data moved when adding/removing
+// nodes.
+// Increasing the tableExpansion will allocate more entries in the
+// internal hash table, reducing the frequency of lg(n) binary
+// searches during calls to the Map.Get method.
+func NewConsistentHash(replicas int, tableExpansion int, fn Hash) *Map {
  m := &Map{
- replicas: replicas,
- hash: fn,
- hashMap: make(map[int]string),
+ replicas: replicas,
+ hash: fn,
+ hashMap: make(map[int]string),
+ prefixTableExpansion: tableExpansion,
  }
  if m.hash == nil {
  m.hash = crc32.ChecksumIEEE
@@ -59,6 +113,37 @@ func (m *Map) Add(keys ...string) {
  }
  }
  sort.Ints(m.keys)
+
+ // Find minimum number of bits to hold |keys| * prefixTableExpansion
+ prefixBits := uint32(bits.Len32(uint32(len(m.keys) * m.prefixTableExpansion)))
+ m.prefixShift = 32 - prefixBits
+
+ prefixTableSize := 1 << prefixBits
+ m.prefixTable = make([]string, prefixTableSize)
+
+ previousKeyPrefix := -1 // Effectively -Inf
+ currentKeyIdx := 0
+ currentKeyPrefix := m.keys[currentKeyIdx] >> m.prefixShift
+
+ for i := range m.prefixTable {
+ if previousKeyPrefix < i && currentKeyPrefix > i {
+ // All keys with this prefix will map to a single value
+ m.prefixTable[i] = m.hashMap[m.keys[currentKeyIdx]]
+ } else {
+ // Several keys might have the same prefix. Walk
+ // over them until it changes
+ previousKeyPrefix = currentKeyPrefix
+ for currentKeyPrefix == previousKeyPrefix {
+ currentKeyIdx++
+ if currentKeyIdx < len(m.keys) {
+ currentKeyPrefix = m.keys[currentKeyIdx] >> m.prefixShift
+ } else {
+ currentKeyIdx = 0
+ currentKeyPrefix = prefixTableSize + 1 // Effectively +Inf
+ }
+ }
+ }
+ }
 }
 
 // Gets the closest item in the hash to the provided key.
@@ -69,6 +154,13 @@ func (m *Map) Get(key string) string {
 
  hash := int(m.hash([]byte(key)))
 
+ // Look for the hash prefix in the prefix table
+ prefixSlot := hash >> m.prefixShift
+ tableResult := m.prefixTable[prefixSlot]
+ if len(tableResult) > 0 {
+ return tableResult
+ }
+
  // Binary search for appropriate replica.
  idx := sort.Search(len(m.keys), func(i int) bool { return m.keys[i] >= hash })
 

consistenthash/consistenthash_test.go

@@ -26,7 +26,7 @@ func TestHashing(t *testing.T) {
 
  // Override the hash function to return easier to reason about values. Assumes
  // the keys can be converted to an integer.
- hash := New(3, func(key []byte) uint32 {
+ hash := NewConsistentHash(3, 6, func(key []byte) uint32 {
  i, err := strconv.Atoi(string(key))
  if err != nil {
  panic(err)
@@ -66,8 +66,8 @@ func TestHashing(t *testing.T) {
 }
 
 func TestConsistency(t *testing.T) {
- hash1 := New(1, nil)
- hash2 := New(1, nil)
+ hash1 := NewConsistentHash(1, 6, nil)
+ hash2 := NewConsistentHash(1, 6, nil)
 
  hash1.Add("Bill", "Bob", "Bonny")
  hash2.Add("Bob", "Bonny", "Bill")
@@ -86,25 +86,31 @@ 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 BenchmarkGet8(b *testing.B) { benchmarkGet(b, 8, 6) }
+func BenchmarkGet32(b *testing.B) { benchmarkGet(b, 32, 6) }
+func BenchmarkGet128(b *testing.B) { benchmarkGet(b, 128, 6) }
+func BenchmarkGet512(b *testing.B) { benchmarkGet(b, 512, 6) }
 
-func benchmarkGet(b *testing.B, shards int) {
+func benchmarkGet(b *testing.B, shards int, expansion int) {
 
- hash := New(50, nil)
+ hash := NewConsistentHash(50, expansion, nil)
 
  var buckets []string
  for i := 0; i < shards; i++ {
  buckets = append(buckets, fmt.Sprintf("shard-%d", i))
  }
 
+ testStringCount := shards
+ var testStrings []string
+ for i := 0; i < testStringCount; i++ {
+ testStrings = append(testStrings, fmt.Sprintf("%d", i))
+ }
+
  hash.Add(buckets...)
 
  b.ResetTimer()
 
  for i := 0; i < b.N; i++ {
- hash.Get(buckets[i&(shards-1)])
+ hash.Get(testStrings[i&(testStringCount-1)])
  }
 }