sonet/pkg/protocol/deliver/consistent_hash_picker.go

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
}