81 lines
1.8 KiB
Go
81 lines
1.8 KiB
Go
package loadbalancer
|
|
|
|
import (
|
|
"github.com/kercylan98/minotaur/utils/super"
|
|
"hash/fnv"
|
|
"sort"
|
|
"sync"
|
|
)
|
|
|
|
func NewConsistentHash(replicas int) *ConsistentHash {
|
|
return &ConsistentHash{
|
|
replicas: replicas,
|
|
keys: []int{},
|
|
hashMap: make(map[int]string),
|
|
mutex: sync.RWMutex{},
|
|
}
|
|
}
|
|
|
|
type ConsistentHash struct {
|
|
replicas int // 虚拟节点倍数
|
|
keys []int // 哈希环上的所有节点的哈希值
|
|
hashMap map[int]string // 哈希值到真实节点的映射
|
|
mutex sync.RWMutex // 用于保护数据结构
|
|
}
|
|
|
|
// Add 添加一个节点到哈希环
|
|
func (c *ConsistentHash) Add(node string) {
|
|
c.mutex.Lock()
|
|
defer c.mutex.Unlock()
|
|
|
|
for i := 0; i < c.replicas; i++ {
|
|
hash := c.hash(node + super.IntToString(i))
|
|
c.keys = append(c.keys, hash)
|
|
c.hashMap[hash] = node
|
|
}
|
|
sort.Ints(c.keys)
|
|
}
|
|
|
|
// Remove 从哈希环中移除一个节点
|
|
func (c *ConsistentHash) Remove(node string) {
|
|
c.mutex.Lock()
|
|
defer c.mutex.Unlock()
|
|
|
|
for i := 0; i < c.replicas; i++ {
|
|
hash := c.hash(node + super.IntToString(i))
|
|
delete(c.hashMap, hash)
|
|
// 从 keys 中移除节点的哈希值
|
|
for j, k := range c.keys {
|
|
if k == hash {
|
|
c.keys = append(c.keys[:j], c.keys[j+1:]...)
|
|
break
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Get 返回给定 key 所在的节点
|
|
func (c *ConsistentHash) Get(key string) string {
|
|
c.mutex.RLock()
|
|
defer c.mutex.RUnlock()
|
|
|
|
if len(c.keys) == 0 {
|
|
return ""
|
|
}
|
|
|
|
hash := c.hash(key)
|
|
// 顺时针找到第一个比 key 大的哈希值，即对应的节点
|
|
idx := sort.Search(len(c.keys), func(i int) bool { return c.keys[i] >= hash })
|
|
if idx == len(c.keys) {
|
|
idx = 0 // 如果 key 大于所有哈希值，则返回第一个节点
|
|
}
|
|
return c.hashMap[c.keys[idx]]
|
|
}
|
|
|
|
// hash 计算字符串的哈希值
|
|
func (c *ConsistentHash) hash(key string) int {
|
|
h := fnv.New32a()
|
|
_, _ = h.Write([]byte(key))
|
|
return int(h.Sum32())
|
|
}
|