consistenthash: replace linear search with binary search

The binary search quickly out-paces the linear search, even for a small
number of shards and replicas.

benchmark          old ns/op    new ns/op    delta
BenchmarkGet8            122          122   +0.00%
BenchmarkGet32           471          137  -70.91%
BenchmarkGet128         5619          254  -95.48%
BenchmarkGet512        90302          406  -99.55%

2 files changed, 31 insertions, 7 deletions

diff --git a/consistenthash/consistenthash.go b/consistenthash/consistenthash.go
index 455ffd5..a9c56f0 100644
--- a/consistenthash/consistenthash.go
+++ b/consistenthash/consistenthash.go
@@ -69,13 +69,13 @@ func (m *Map) Get(key string) string {
 
 	hash := int(m.hash([]byte(key)))
 
-	// Linear search for appropriate replica.
-	for _, v := range m.keys {
-		if v >= hash {
-			return m.hashMap[v]
-		}
-	}
+	// Binary search for appropriate replica.
+	idx := sort.Search(len(m.keys), func(i int) bool { return m.keys[i] >= hash })
 
 	// Means we have cycled back to the first replica.
-	return m.hashMap[m.keys[0]]
+	if idx == len(m.keys) {
+		idx = 0
+	}
+
+	return m.hashMap[m.keys[idx]]
 }
diff --git a/consistenthash/consistenthash_test.go b/consistenthash/consistenthash_test.go
index a1b96db..f2b82d5 100644
--- a/consistenthash/consistenthash_test.go
+++ b/consistenthash/consistenthash_test.go
@@ -17,6 +17,7 @@ limitations under the License.
 package consistenthash
 
 import (
+	"fmt"
 	"strconv"
 	"testing"
 )
@@ -84,3 +85,26 @@ 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(shards, 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)])
+	}
+}