github.com/lingyao2333/mo-zero@v1.4.1/core/hash/consistenthash_test.go (about)

     1  package hash
     2  
     3  import (
     4  	"fmt"
     5  	"strconv"
     6  	"testing"
     7  
     8  	"github.com/lingyao2333/mo-zero/core/mathx"
     9  	"github.com/stretchr/testify/assert"
    10  )
    11  
    12  const (
    13  	keySize     = 20
    14  	requestSize = 1000
    15  )
    16  
    17  func BenchmarkConsistentHashGet(b *testing.B) {
    18  	ch := NewConsistentHash()
    19  	for i := 0; i < keySize; i++ {
    20  		ch.Add("localhost:" + strconv.Itoa(i))
    21  	}
    22  
    23  	for i := 0; i < b.N; i++ {
    24  		ch.Get(i)
    25  	}
    26  }
    27  
    28  func TestConsistentHash(t *testing.T) {
    29  	ch := NewCustomConsistentHash(0, nil)
    30  	val, ok := ch.Get("any")
    31  	assert.False(t, ok)
    32  	assert.Nil(t, val)
    33  
    34  	for i := 0; i < keySize; i++ {
    35  		ch.AddWithReplicas("localhost:"+strconv.Itoa(i), minReplicas<<1)
    36  	}
    37  
    38  	keys := make(map[string]int)
    39  	for i := 0; i < requestSize; i++ {
    40  		key, ok := ch.Get(requestSize + i)
    41  		assert.True(t, ok)
    42  		keys[key.(string)]++
    43  	}
    44  
    45  	mi := make(map[interface{}]int, len(keys))
    46  	for k, v := range keys {
    47  		mi[k] = v
    48  	}
    49  	entropy := mathx.CalcEntropy(mi)
    50  	assert.True(t, entropy > .95)
    51  }
    52  
    53  func TestConsistentHashIncrementalTransfer(t *testing.T) {
    54  	prefix := "anything"
    55  	create := func() *ConsistentHash {
    56  		ch := NewConsistentHash()
    57  		for i := 0; i < keySize; i++ {
    58  			ch.Add(prefix + strconv.Itoa(i))
    59  		}
    60  		return ch
    61  	}
    62  
    63  	originCh := create()
    64  	keys := make(map[int]string, requestSize)
    65  	for i := 0; i < requestSize; i++ {
    66  		key, ok := originCh.Get(requestSize + i)
    67  		assert.True(t, ok)
    68  		assert.NotNil(t, key)
    69  		keys[i] = key.(string)
    70  	}
    71  
    72  	node := fmt.Sprintf("%s%d", prefix, keySize)
    73  	for i := 0; i < 10; i++ {
    74  		laterCh := create()
    75  		laterCh.AddWithWeight(node, 10*(i+1))
    76  
    77  		for j := 0; j < requestSize; j++ {
    78  			key, ok := laterCh.Get(requestSize + j)
    79  			assert.True(t, ok)
    80  			assert.NotNil(t, key)
    81  			value := key.(string)
    82  			assert.True(t, value == keys[j] || value == node)
    83  		}
    84  	}
    85  }
    86  
    87  func TestConsistentHashTransferOnFailure(t *testing.T) {
    88  	index := 41
    89  	keys, newKeys := getKeysBeforeAndAfterFailure(t, "localhost:", index)
    90  	var transferred int
    91  	for k, v := range newKeys {
    92  		if v != keys[k] {
    93  			transferred++
    94  		}
    95  	}
    96  
    97  	ratio := float32(transferred) / float32(requestSize)
    98  	assert.True(t, ratio < 2.5/float32(keySize), fmt.Sprintf("%d: %f", index, ratio))
    99  }
   100  
   101  func TestConsistentHashLeastTransferOnFailure(t *testing.T) {
   102  	prefix := "localhost:"
   103  	index := 41
   104  	keys, newKeys := getKeysBeforeAndAfterFailure(t, prefix, index)
   105  	for k, v := range keys {
   106  		newV := newKeys[k]
   107  		if v != prefix+strconv.Itoa(index) {
   108  			assert.Equal(t, v, newV)
   109  		}
   110  	}
   111  }
   112  
   113  func TestConsistentHash_Remove(t *testing.T) {
   114  	ch := NewConsistentHash()
   115  	ch.Add("first")
   116  	ch.Add("second")
   117  	ch.Remove("first")
   118  	for i := 0; i < 100; i++ {
   119  		val, ok := ch.Get(i)
   120  		assert.True(t, ok)
   121  		assert.Equal(t, "second", val)
   122  	}
   123  }
   124  
   125  func TestConsistentHash_RemoveInterface(t *testing.T) {
   126  	const key = "any"
   127  	ch := NewConsistentHash()
   128  	node1 := newMockNode(key, 1)
   129  	node2 := newMockNode(key, 2)
   130  	ch.AddWithWeight(node1, 80)
   131  	ch.AddWithWeight(node2, 50)
   132  	assert.Equal(t, 1, len(ch.nodes))
   133  	node, ok := ch.Get(1)
   134  	assert.True(t, ok)
   135  	assert.Equal(t, key, node.(*mockNode).addr)
   136  	assert.Equal(t, 2, node.(*mockNode).id)
   137  }
   138  
   139  func getKeysBeforeAndAfterFailure(t *testing.T, prefix string, index int) (map[int]string, map[int]string) {
   140  	ch := NewConsistentHash()
   141  	for i := 0; i < keySize; i++ {
   142  		ch.Add(prefix + strconv.Itoa(i))
   143  	}
   144  
   145  	keys := make(map[int]string, requestSize)
   146  	for i := 0; i < requestSize; i++ {
   147  		key, ok := ch.Get(requestSize + i)
   148  		assert.True(t, ok)
   149  		assert.NotNil(t, key)
   150  		keys[i] = key.(string)
   151  	}
   152  
   153  	remove := fmt.Sprintf("%s%d", prefix, index)
   154  	ch.Remove(remove)
   155  	newKeys := make(map[int]string, requestSize)
   156  	for i := 0; i < requestSize; i++ {
   157  		key, ok := ch.Get(requestSize + i)
   158  		assert.True(t, ok)
   159  		assert.NotNil(t, key)
   160  		assert.NotEqual(t, remove, key)
   161  		newKeys[i] = key.(string)
   162  	}
   163  
   164  	return keys, newKeys
   165  }
   166  
   167  type mockNode struct {
   168  	addr string
   169  	id   int
   170  }
   171  
   172  func newMockNode(addr string, id int) *mockNode {
   173  	return &mockNode{
   174  		addr: addr,
   175  		id:   id,
   176  	}
   177  }
   178  
   179  func (n *mockNode) String() string {
   180  	return n.addr
   181  }