consistenthash.go 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187
  1. package hash
  2. import (
  3. "fmt"
  4. "sort"
  5. "strconv"
  6. "sync"
  7. "github.com/tal-tech/go-zero/core/lang"
  8. "github.com/tal-tech/go-zero/core/mapping"
  9. )
  10. const (
  11. // TopWeight is the top weight that one entry might set.
  12. TopWeight = 100
  13. minReplicas = 100
  14. prime = 16777619
  15. )
  16. type (
  17. // Func defines the hash method.
  18. Func func(data []byte) uint64
  19. // A ConsistentHash is a ring hash implementation.
  20. ConsistentHash struct {
  21. hashFunc Func
  22. replicas int
  23. keys []uint64
  24. ring map[uint64][]interface{}
  25. nodes map[string]lang.PlaceholderType
  26. lock sync.RWMutex
  27. }
  28. )
  29. // NewConsistentHash returns a ConsistentHash.
  30. func NewConsistentHash() *ConsistentHash {
  31. return NewCustomConsistentHash(minReplicas, Hash)
  32. }
  33. // NewCustomConsistentHash returns a ConsistentHash with given replicas and hash func.
  34. func NewCustomConsistentHash(replicas int, fn Func) *ConsistentHash {
  35. if replicas < minReplicas {
  36. replicas = minReplicas
  37. }
  38. if fn == nil {
  39. fn = Hash
  40. }
  41. return &ConsistentHash{
  42. hashFunc: fn,
  43. replicas: replicas,
  44. ring: make(map[uint64][]interface{}),
  45. nodes: make(map[string]lang.PlaceholderType),
  46. }
  47. }
  48. // Add adds the node with the number of h.replicas,
  49. // the later call will overwrite the replicas of the former calls.
  50. func (h *ConsistentHash) Add(node interface{}) {
  51. h.AddWithReplicas(node, h.replicas)
  52. }
  53. // AddWithReplicas adds the node with the number of replicas,
  54. // replicas will be truncated to h.replicas if it's larger than h.replicas,
  55. // the later call will overwrite the replicas of the former calls.
  56. func (h *ConsistentHash) AddWithReplicas(node interface{}, replicas int) {
  57. h.Remove(node)
  58. if replicas > h.replicas {
  59. replicas = h.replicas
  60. }
  61. nodeRepr := repr(node)
  62. h.lock.Lock()
  63. defer h.lock.Unlock()
  64. h.addNode(nodeRepr)
  65. for i := 0; i < replicas; i++ {
  66. hash := h.hashFunc([]byte(nodeRepr + strconv.Itoa(i)))
  67. h.keys = append(h.keys, hash)
  68. h.ring[hash] = append(h.ring[hash], node)
  69. }
  70. sort.Slice(h.keys, func(i int, j int) bool {
  71. return h.keys[i] < h.keys[j]
  72. })
  73. }
  74. // AddWithWeight adds the node with weight, the weight can be 1 to 100, indicates the percent,
  75. // the later call will overwrite the replicas of the former calls.
  76. func (h *ConsistentHash) AddWithWeight(node interface{}, weight int) {
  77. // don't need to make sure weight not larger than TopWeight,
  78. // because AddWithReplicas makes sure replicas cannot be larger than h.replicas
  79. replicas := h.replicas * weight / TopWeight
  80. h.AddWithReplicas(node, replicas)
  81. }
  82. // Get returns the corresponding node from h base on the given v.
  83. func (h *ConsistentHash) Get(v interface{}) (interface{}, bool) {
  84. h.lock.RLock()
  85. defer h.lock.RUnlock()
  86. if len(h.ring) == 0 {
  87. return nil, false
  88. }
  89. hash := h.hashFunc([]byte(repr(v)))
  90. index := sort.Search(len(h.keys), func(i int) bool {
  91. return h.keys[i] >= hash
  92. }) % len(h.keys)
  93. nodes := h.ring[h.keys[index]]
  94. switch len(nodes) {
  95. case 0:
  96. return nil, false
  97. case 1:
  98. return nodes[0], true
  99. default:
  100. innerIndex := h.hashFunc([]byte(innerRepr(v)))
  101. pos := int(innerIndex % uint64(len(nodes)))
  102. return nodes[pos], true
  103. }
  104. }
  105. // Remove removes the given node from h.
  106. func (h *ConsistentHash) Remove(node interface{}) {
  107. nodeRepr := repr(node)
  108. h.lock.Lock()
  109. defer h.lock.Unlock()
  110. if !h.containsNode(nodeRepr) {
  111. return
  112. }
  113. for i := 0; i < h.replicas; i++ {
  114. hash := h.hashFunc([]byte(nodeRepr + strconv.Itoa(i)))
  115. index := sort.Search(len(h.keys), func(i int) bool {
  116. return h.keys[i] >= hash
  117. })
  118. if index < len(h.keys) {
  119. h.keys = append(h.keys[:index], h.keys[index+1:]...)
  120. }
  121. h.removeRingNode(hash, nodeRepr)
  122. }
  123. h.removeNode(nodeRepr)
  124. }
  125. func (h *ConsistentHash) removeRingNode(hash uint64, nodeRepr string) {
  126. if nodes, ok := h.ring[hash]; ok {
  127. newNodes := nodes[:0]
  128. for _, x := range nodes {
  129. if repr(x) != nodeRepr {
  130. newNodes = append(newNodes, x)
  131. }
  132. }
  133. if len(newNodes) > 0 {
  134. h.ring[hash] = newNodes
  135. } else {
  136. delete(h.ring, hash)
  137. }
  138. }
  139. }
  140. func (h *ConsistentHash) addNode(nodeRepr string) {
  141. h.nodes[nodeRepr] = lang.Placeholder
  142. }
  143. func (h *ConsistentHash) containsNode(nodeRepr string) bool {
  144. _, ok := h.nodes[nodeRepr]
  145. return ok
  146. }
  147. func (h *ConsistentHash) removeNode(nodeRepr string) {
  148. delete(h.nodes, nodeRepr)
  149. }
  150. func innerRepr(node interface{}) string {
  151. return fmt.Sprintf("%d:%v", prime, node)
  152. }
  153. func repr(node interface{}) string {
  154. return mapping.Repr(node)
  155. }