adm-ads/common.in/utils/consistency_hash.go

package utils

import (
	"fmt"
	"hash/crc32"
	"sort"
)

type ConsistentHash struct {
	VirtualNodes    int
	virtualNodesMap map[uint32]string
	actNodes        map[string]string
	sortRing        []uint32
}

func NewConsistentHash(virtualNodeNumber int) *ConsistentHash {
	return &ConsistentHash{
		VirtualNodes:    virtualNodeNumber,
		virtualNodesMap: make(map[uint32]string),
		actNodes:        make(map[string]string),
		sortRing:        []uint32{},
	}
}

func (c *ConsistentHash) Add(node string) bool {

	if _, ok := c.actNodes[node]; ok {
		return false
	}

	for i := 0; i < c.VirtualNodes; i++ {
		c.virtualNodesMap[c.hashStr(fmt.Sprintf("%s#%d", node, i))] = node
	}
	c.actNodes[node] = node
	c.sortHashRing()
	return true
}

func (c *ConsistentHash) sortHashRing() {
	c.sortRing = []uint32{}
	for k := range c.virtualNodesMap {
		c.sortRing = append(c.sortRing, k)
	}
	sort.Slice(c.sortRing, func(i, j int) bool {
		return c.sortRing[i] < c.sortRing[j]
	})
}

func (c *ConsistentHash) Get(key string) string {
	hash := c.hashStr(key)

	if len(c.virtualNodesMap) == 0 {
		return ""
	}
	i := c.search(hash)
	return c.virtualNodesMap[c.sortRing[i]]
}

func (c *ConsistentHash) hashStr(key string) uint32 {
	return crc32.ChecksumIEEE([]byte(key))
}

func (c *ConsistentHash) search(hash uint32) int {
	i := sort.Search(len(c.sortRing), func(i int) bool { return c.sortRing[i] >= hash })
	if i < len(c.sortRing) {
		if i == len(c.sortRing)-1 {
			return 0
		} else {
			return i
		}
	} else {
		return len(c.sortRing) - 1
	}
}