gopkg/consistent_hash/consistent.go

// Package consistent_hash...
//
// Description : consistent_hash...
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 2021-11-21 4:57 下午
package consistent_hash

import (
	"sort"
	"strconv"
	"sync"
)

// New 获取实例
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 5:03 下午 2021/11/21
func New(replicas int, fn HashFunc) *Consistent {
	if nil == fn {
		fn = DefaultHashFunc
	}
	return &Consistent{
		hashFunc: fn,
		replicas: replicas,
		keys:     make([]int, 0),
		hashMap:  make(map[int]string),
		lock:     &sync.RWMutex{},
	}
}

// Consistent ...
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 5:00 下午 2021/11/21
type Consistent struct {
	hashFunc HashFunc       // 哈希函数
	replicas int            // 虚拟的数量
	keys     []int          // 哈希环
	hashMap  map[int]string // 虚拟节点与真实节点的映射表
	lock     *sync.RWMutex  // 锁
}

// Add 添加节点
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 5:10 下午 2021/11/21
func (c *Consistent) Add(keyList ...string) {
	c.lock.Lock()
	defer c.lock.Unlock()
	for _, key := range keyList {
		for i := 0; i < c.replicas; i++ {
			hash := int(c.hashFunc([]byte(strconv.Itoa(i) + key)))
			c.keys = append(c.keys, hash)
			c.hashMap[hash] = key
		}
	}
	// 哈希值排序
	sort.Ints(c.keys)
}

// Get 获取节点
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 5:16 下午 2021/11/21
func (c *Consistent) Get(key string) string {
	c.lock.RLock()
	defer c.lock.Unlock()
	if len(c.keys) == 0 {
		return ""
	}
	hash := int(c.hashFunc([]byte(key)))
	// 查找节点索引
	idx := sort.Search(len(c.keys), func(i int) bool {
		return c.keys[i] >= hash
	})
	return c.hashMap[c.keys[idx%len(c.keys)]]
}

// Delete 删除一个节点,本质逻辑 : 重建哈希环
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 5:35 下午 2021/11/21
func (c *Consistent) Delete(key string) {
	c.lock.Lock()
	defer c.lock.Unlock()
	tmpKeyList := make([]string, 0)
	for _, k := range c.hashMap {
		if k == key {
			continue
		}
		tmpKeyList = append(tmpKeyList, k)
	}
	c.hashMap = make(map[int]string)
	c.keys = make([]int, 0)
	c.Add(tmpKeyList...)
}