consistenthash_test.go 3.9 KB

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