initial import

core/hash/consistenthash.go

@@ -0,0 +1,180 @@
+package hash
+
+import (
+	"fmt"
+	"sort"
+	"strconv"
+	"sync"
+
+	"zero/core/lang"
+	"zero/core/mapping"
+)
+
+const (
+	TopWeight = 100
+
+	minReplicas = 100
+	prime       = 16777619
+)
+
+type (
+	HashFunc func(data []byte) uint64
+
+	ConsistentHash struct {
+		hashFunc HashFunc
+		replicas int
+		keys     []uint64
+		ring     map[uint64][]interface{}
+		nodes    map[string]lang.PlaceholderType
+		lock     sync.RWMutex
+	}
+)
+
+func NewConsistentHash() *ConsistentHash {
+	return NewCustomConsistentHash(minReplicas, Hash)
+}
+
+func NewCustomConsistentHash(replicas int, fn HashFunc) *ConsistentHash {
+	if replicas < minReplicas {
+		replicas = minReplicas
+	}
+
+	if fn == nil {
+		fn = Hash
+	}
+
+	return &ConsistentHash{
+		hashFunc: fn,
+		replicas: replicas,
+		ring:     make(map[uint64][]interface{}),
+		nodes:    make(map[string]lang.PlaceholderType),
+	}
+}
+
+// Add adds the node with the number of h.replicas,
+// the later call will overwrite the replicas of the former calls.
+func (h *ConsistentHash) Add(node interface{}) {
+	h.AddWithReplicas(node, h.replicas)
+}
+
+// AddWithReplicas adds the node with the number of replicas,
+// replicas will be truncated to h.replicas if it's larger than h.replicas,
+// the later call will overwrite the replicas of the former calls.
+func (h *ConsistentHash) AddWithReplicas(node interface{}, replicas int) {
+	h.Remove(node)
+
+	if replicas > h.replicas {
+		replicas = h.replicas
+	}
+
+	nodeRepr := repr(node)
+	h.lock.Lock()
+	defer h.lock.Unlock()
+	h.addNode(nodeRepr)
+
+	for i := 0; i < replicas; i++ {
+		hash := h.hashFunc([]byte(nodeRepr + strconv.Itoa(i)))
+		h.keys = append(h.keys, hash)
+		h.ring[hash] = append(h.ring[hash], node)
+	}
+
+	sort.Slice(h.keys, func(i int, j int) bool {
+		return h.keys[i] < h.keys[j]
+	})
+}
+
+// AddWithWeight adds the node with weight, the weight can be 1 to 100, indicates the percent,
+// the later call will overwrite the replicas of the former calls.
+func (h *ConsistentHash) AddWithWeight(node interface{}, weight int) {
+	// don't need to make sure weight not larger than TopWeight,
+	// because AddWithReplicas makes sure replicas cannot be larger than h.replicas
+	replicas := h.replicas * weight / TopWeight
+	h.AddWithReplicas(node, replicas)
+}
+
+func (h *ConsistentHash) Get(v interface{}) (interface{}, bool) {
+	h.lock.RLock()
+	defer h.lock.RUnlock()
+
+	if len(h.ring) == 0 {
+		return nil, false
+	}
+
+	hash := h.hashFunc([]byte(repr(v)))
+	index := sort.Search(len(h.keys), func(i int) bool {
+		return h.keys[i] >= hash
+	}) % len(h.keys)
+
+	nodes := h.ring[h.keys[index]]
+	switch len(nodes) {
+	case 0:
+		return nil, false
+	case 1:
+		return nodes[0], true
+	default:
+		innerIndex := h.hashFunc([]byte(innerRepr(v)))
+		pos := int(innerIndex % uint64(len(nodes)))
+		return nodes[pos], true
+	}
+}
+
+func (h *ConsistentHash) Remove(node interface{}) {
+	nodeRepr := repr(node)
+
+	h.lock.Lock()
+	defer h.lock.Unlock()
+
+	if !h.containsNode(nodeRepr) {
+		return
+	}
+
+	for i := 0; i < h.replicas; i++ {
+		hash := h.hashFunc([]byte(nodeRepr + strconv.Itoa(i)))
+		index := sort.Search(len(h.keys), func(i int) bool {
+			return h.keys[i] >= hash
+		})
+		if index < len(h.keys) {
+			h.keys = append(h.keys[:index], h.keys[index+1:]...)
+		}
+		h.removeRingNode(hash, nodeRepr)
+	}
+
+	h.removeNode(nodeRepr)
+}
+
+func (h *ConsistentHash) removeRingNode(hash uint64, nodeRepr string) {
+	if nodes, ok := h.ring[hash]; ok {
+		newNodes := nodes[:0]
+		for _, x := range nodes {
+			if repr(x) != nodeRepr {
+				newNodes = append(newNodes, x)
+			}
+		}
+		if len(newNodes) > 0 {
+			h.ring[hash] = newNodes
+		} else {
+			delete(h.ring, hash)
+		}
+	}
+}
+
+func (h *ConsistentHash) addNode(nodeRepr string) {
+	h.nodes[nodeRepr] = lang.Placeholder
+}
+
+func (h *ConsistentHash) containsNode(nodeRepr string) bool {
+	_, ok := h.nodes[nodeRepr]
+	return ok
+}
+
+func (h *ConsistentHash) removeNode(nodeRepr string) {
+	delete(h.nodes, nodeRepr)
+}
+
+func innerRepr(node interface{}) string {
+	return fmt.Sprintf("%d:%v", prime, node)
+}
+
+func repr(node interface{}) string {
+	return mapping.Repr(node)
+}

core/hash/consistenthash_test.go

@@ -0,0 +1,182 @@
+package hash
+
+import (
+	"fmt"
+	"strconv"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"zero/core/mathx"
+)
+
+const (
+	keySize     = 20
+	requestSize = 1000
+)
+
+func BenchmarkConsistentHashGet(b *testing.B) {
+	ch := NewConsistentHash()
+	for i := 0; i < keySize; i++ {
+		ch.Add("localhost:" + strconv.Itoa(i))
+	}
+
+	for i := 0; i < b.N; i++ {
+		ch.Get(i)
+	}
+}
+
+func TestConsistentHash(t *testing.T) {
+	ch := NewCustomConsistentHash(0, nil)
+	val, ok := ch.Get("any")
+	assert.False(t, ok)
+	assert.Nil(t, val)
+
+	for i := 0; i < keySize; i++ {
+		ch.AddWithReplicas("localhost:"+strconv.Itoa(i), minReplicas<<1)
+	}
+
+	keys := make(map[string]int)
+	for i := 0; i < requestSize; i++ {
+		key, ok := ch.Get(requestSize + i)
+		assert.True(t, ok)
+		keys[key.(string)]++
+	}
+
+	mi := make(map[interface{}]int, len(keys))
+	for k, v := range keys {
+		mi[k] = v
+	}
+	entropy := mathx.CalcEntropy(mi, requestSize)
+	assert.True(t, entropy > .95)
+}
+
+func TestConsistentHashIncrementalTransfer(t *testing.T) {
+	prefix := "anything"
+	create := func() *ConsistentHash {
+		ch := NewConsistentHash()
+		for i := 0; i < keySize; i++ {
+			ch.Add(prefix + strconv.Itoa(i))
+		}
+		return ch
+	}
+
+	originCh := create()
+	keys := make(map[int]string, requestSize)
+	for i := 0; i < requestSize; i++ {
+		key, ok := originCh.Get(requestSize + i)
+		assert.True(t, ok)
+		assert.NotNil(t, key)
+		keys[i] = key.(string)
+	}
+
+	node := fmt.Sprintf("%s%d", prefix, keySize)
+	for i := 0; i < 10; i++ {
+		laterCh := create()
+		laterCh.AddWithWeight(node, 10*(i+1))
+
+		for i := 0; i < requestSize; i++ {
+			key, ok := laterCh.Get(requestSize + i)
+			assert.True(t, ok)
+			assert.NotNil(t, key)
+			value := key.(string)
+			assert.True(t, value == keys[i] || value == node)
+		}
+	}
+}
+
+func TestConsistentHashTransferOnFailure(t *testing.T) {
+	index := 41
+	keys, newKeys := getKeysBeforeAndAfterFailure(t, "localhost:", index)
+	var transferred int
+	for k, v := range newKeys {
+		if v != keys[k] {
+			transferred++
+		}
+	}
+
+	ratio := float32(transferred) / float32(requestSize)
+	assert.True(t, ratio < 2.5/float32(keySize), fmt.Sprintf("%d: %f", index, ratio))
+}
+
+func TestConsistentHashLeastTransferOnFailure(t *testing.T) {
+	prefix := "localhost:"
+	index := 41
+	keys, newKeys := getKeysBeforeAndAfterFailure(t, prefix, index)
+	for k, v := range keys {
+		newV := newKeys[k]
+		if v != prefix+strconv.Itoa(index) {
+			assert.Equal(t, v, newV)
+		}
+	}
+}
+
+func TestConsistentHash_Remove(t *testing.T) {
+	ch := NewConsistentHash()
+	ch.Add("first")
+	ch.Add("second")
+	ch.Remove("first")
+	for i := 0; i < 100; i++ {
+		val, ok := ch.Get(i)
+		assert.True(t, ok)
+		assert.Equal(t, "second", val)
+	}
+}
+
+func TestConsistentHash_RemoveInterface(t *testing.T) {
+	const key = "any"
+	ch := NewConsistentHash()
+	node1 := newMockNode(key, 1)
+	node2 := newMockNode(key, 2)
+	ch.AddWithWeight(node1, 80)
+	ch.AddWithWeight(node2, 50)
+	assert.Equal(t, 1, len(ch.nodes))
+	node, ok := ch.Get(1)
+	assert.True(t, ok)
+	assert.Equal(t, key, node.(*MockNode).Addr)
+	assert.Equal(t, 2, node.(*MockNode).Id)
+}
+
+func getKeysBeforeAndAfterFailure(t *testing.T, prefix string, index int) (map[int]string, map[int]string) {
+	ch := NewConsistentHash()
+	for i := 0; i < keySize; i++ {
+		ch.Add(prefix + strconv.Itoa(i))
+	}
+
+	keys := make(map[int]string, requestSize)
+	for i := 0; i < requestSize; i++ {
+		key, ok := ch.Get(requestSize + i)
+		assert.True(t, ok)
+		assert.NotNil(t, key)
+		keys[i] = key.(string)
+	}
+
+	remove := fmt.Sprintf("%s%d", prefix, index)
+	ch.Remove(remove)
+	newKeys := make(map[int]string, requestSize)
+	for i := 0; i < requestSize; i++ {
+		key, ok := ch.Get(requestSize + i)
+		assert.True(t, ok)
+		assert.NotNil(t, key)
+		assert.NotEqual(t, remove, key)
+		newKeys[i] = key.(string)
+	}
+
+	return keys, newKeys
+}
+
+type MockNode struct {
+	Addr string
+	Id   int
+}
+
+func newMockNode(addr string, id int) *MockNode {
+	return &MockNode{
+		Addr: addr,
+		Id:   id,
+	}
+}
+
+func (n *MockNode) String() string {
+	return n.Addr
+}

core/hash/hash.go

@@ -0,0 +1,22 @@
+package hash
+
+import (
+	"crypto/md5"
+	"fmt"
+
+	"github.com/spaolacci/murmur3"
+)
+
+func Hash(data []byte) uint64 {
+	return murmur3.Sum64(data)
+}
+
+func Md5(data []byte) []byte {
+	digest := md5.New()
+	digest.Write(data)
+	return digest.Sum(nil)
+}
+
+func Md5Hex(data []byte) string {
+	return fmt.Sprintf("%x", Md5(data))
+}

core/hash/hash_test.go

@@ -0,0 +1,42 @@
+package hash
+
+import (
+	"crypto/md5"
+	"fmt"
+	"hash/fnv"
+	"math/big"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	text      = "hello, world!\n"
+	md5Digest = "910c8bc73110b0cd1bc5d2bcae782511"
+)
+
+func TestMd5(t *testing.T) {
+	actual := fmt.Sprintf("%x", Md5([]byte(text)))
+	assert.Equal(t, md5Digest, actual)
+}
+
+func BenchmarkHashFnv(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		h := fnv.New32()
+		new(big.Int).SetBytes(h.Sum([]byte(text))).Int64()
+	}
+}
+
+func BenchmarkHashMd5(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		h := md5.New()
+		bytes := h.Sum([]byte(text))
+		new(big.Int).SetBytes(bytes).Int64()
+	}
+}
+
+func BenchmarkMurmur3(b *testing.B) {
+	for i := 0; i < b.N; i++ {
+		Hash([]byte(text))
+	}
+}