package deliver import ( "hash/fnv" "sort" "strconv" ) type PickNode struct { Key string Weight int } type ConsistentHashPicker struct { nodes []*PickNode replicas int salt string hashKeys []uint32 // sorted hashKeys hashNodes map[uint32]*PickNode length int } func NewConsistentHashPicker(nodes []*PickNode, replicas int, salt string) *ConsistentHashPicker { return &ConsistentHashPicker{ nodes: nodes, replicas: replicas, salt: salt, hashKeys: make([]uint32, 0, len(nodes)*replicas), hashNodes: make(map[uint32]*PickNode, len(nodes)*replicas), } } var ( DefaultReplicas = 10 DefaultSalt = "this_is_salt" ) func (c *ConsistentHashPicker) hashFnv32(data []byte) uint32 { f := fnv.New32() _, err := f.Write(data) if err != nil { panic(err) } return f.Sum32() } // Init 构建hash环 func (c *ConsistentHashPicker) Init() { for _, node := range c.nodes { weight := node.Weight if weight < 1 { weight = 1 } for i := 0; i < weight; i++ { for j := 0; j < c.replicas; j++ { key := c.hashFnv32([]byte(strconv.Itoa(i) + node.Key + strconv.Itoa(j) + c.salt)) if _, ok := c.hashNodes[key]; !ok { c.hashKeys = append(c.hashKeys, key) } c.hashNodes[key] = node } } } sort.Slice(c.hashKeys, func(i, j int) bool { return c.hashKeys[i] < c.hashKeys[j] }) c.length = len(c.hashKeys) } func (c *ConsistentHashPicker) Pick(key string) (node *PickNode, ok bool) { if c.length == 0 { return } hash := c.hashFnv32([]byte(key + c.salt)) idx := sort.Search(c.length, func(i int) bool { // 二分查找最小为true的index return c.hashKeys[i] >= hash }) if idx == c.length { idx = 0 } node, ok = c.hashNodes[c.hashKeys[idx]] return } // PickOffset // key: pick key // offset: 顺时针第n个节点 // return node: pick node // return same: node是否是offset=0时的相同节点 func (c *ConsistentHashPicker) PickOffset(key string, offset int) (node *PickNode, same bool, ok bool) { if c.length == 0 { return } hash := c.hashFnv32([]byte(key + c.salt)) idx := sort.Search(c.length, func(i int) bool { return c.hashKeys[i] >= hash }) if idx == c.length { idx = 0 } hit := (idx + offset) % c.length if hit < 0 { hit = c.length + hit } node, ok = c.hashNodes[c.hashKeys[hit]] same = c.hashNodes[c.hashKeys[idx]] == node return }