go-redis/lib/consistenthash/consistenthash.go

// Author: simon (ynwdlxm@163.com)
// Date: 2025/10/17 10:50
// Desc: 一致性哈希算法

package consistenthash

import (
	"hash/crc32"
	"sort"
)

// HashFunc 定义哈希函数类型，接收字节切片并返回32位无符号整数
type HashFunc func(data []byte) uint32

// NodeMap 一致性哈希节点映射结构体
// 用于维护虚拟节点到真实节点的映射关系
type NodeMap struct {
	hashFunc HashFunc       // 哈希函数
	nodeHash []int          // 节点哈希值数组
	nodeMap  map[int]string // 哈希值到节点名称的映射
}

// NewNodeMap 创建新的NodeMap实例
// 如果传入的哈希函数为空，则使用默认的crc32.ChecksumIEEE函数
func NewNodeMap(fn HashFunc) *NodeMap {
	m := &NodeMap{
		hashFunc: fn,
		nodeMap:  make(map[int]string),
	}
	if m.hashFunc == nil {
		m.hashFunc = crc32.ChecksumIEEE
	}
	return m
}

// IsEmpty 判断节点映射是否为空
// 返回true表示没有节点，false表示有节点
func (m *NodeMap) IsEmpty() bool {
	return len(m.nodeHash) == 0
}

// AddNode 添加节点到一致性哈希环中
// 支持同时添加多个节点，会为每个节点计算哈希值并维护映射关系
// 空节点会被忽略
func (m *NodeMap) AddNode(keys ...string) {
	for _, key := range keys {
		if len(key) == 0 {
			// key is empty
			continue
		}
		hash := int(m.hashFunc([]byte(key)))
		m.nodeHash = append(m.nodeHash, hash)
		m.nodeMap[hash] = key
	}
	// 排序节点哈希值，保证一致性哈希环的有序性
	sort.Ints(m.nodeHash)
}

// PickNode 根据键值选择对应的节点
// 通过计算键值的哈希值，在一致性哈希环上找到对应的节点
// 如果哈希环为空则返回空字符串
func (m *NodeMap) PickNode(key string) string {
	if m.IsEmpty() {
		return ""
	}
	hash := int(m.hashFunc([]byte(key)))

	// 使用二分查找找到第一个大于等于hash值的位置
	idx := sort.Search(len(m.nodeHash), func(i int) bool {
		return m.nodeHash[i] >= hash
	})
	// 如果找不到合适的节点，则选择第一个节点（环状结构）
	if idx == len(m.nodeHash) {
		idx = 0
	}
	return m.nodeMap[m.nodeHash[idx]]
}