一致性哈希

package consistenthash

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

// Hash 定义哈希函数类型
type Hash func(data []byte) uint32

// Map 一致性哈希结构
type Map struct {
	hash     Hash           // 哈希函数
	replicas int            // 虚拟节点倍数
	keys     []int          // 哈希环
	hashMap  map[int]string // 虚拟节点与真实节点的映射
}

// New 创建一致性哈希实例
func New(replicas int, fn Hash) *Map {
	m := &Map{
		replicas: replicas,
		hash:     fn,
		hashMap:  make(map[int]string),
	}
	if m.hash == nil {
		m.hash = crc32.ChecksumIEEE // 默认哈希函数
	}
	return m
}

// Add 添加节点到哈希环
func (m *Map) Add(keys ...string) {
	for _, key := range keys {
		for i := 0; i < m.replicas; i++ {
			// 计算虚拟节点的哈希值
			hash := int(m.hash([]byte(strconv.Itoa(i) + key)))
			m.keys = append(m.keys, hash)
			m.hashMap[hash] = key // 虚拟节点映射到真实节点
		}
	}
	sort.Ints(m.keys) // 排序哈希环
}

// Get 根据key获取对应的节点
func (m *Map) Get(key string) string {
	if len(m.keys) == 0 {
		return ""
	}

	// 计算key的哈希值
	hash := int(m.hash([]byte(key)))

	// 二分查找第一个大于等于hash的节点
	idx := sort.Search(len(m.keys), func(i int) bool {
		return m.keys[i] >= hash
	})

	// 如果没找到，则使用第一个节点(环状结构)
	if idx == len(m.keys) {
		idx = 0
	}

	return m.hashMap[m.keys[idx]]
}

// Remove 从哈希环中移除节点
func (m *Map) Remove(key string) {
	for i := 0; i < m.replicas; i++ {
		hash := int(m.hash([]byte(strconv.Itoa(i) + key)))
		// 从排序列表中删除
		idx := sort.SearchInts(m.keys, hash)
		if idx < len(m.keys) && m.keys[idx] == hash {
			m.keys = append(m.keys[:idx], m.keys[idx+1:]...)
		}
		// 从映射中删除
		delete(m.hashMap, hash)
	}
}

package main

import (
	"fmt"
	"yourmodule/consistenthash"
)

func main() {
	// 创建一致性哈希实例，设置虚拟节点数为3
	ch := consistenthash.New(3, nil)
	
	// 添加节点
	ch.Add("Node1", "Node2", "Node3")
	
	// 测试数据分布
	testKeys := []string{"key1", "key2", "key3", "key4", "key5"}
	for _, key := range testKeys {
		fmt.Printf("Key %s is assigned to %s\n", key, ch.Get(key))
	}
	
	// 添加新节点
	fmt.Println("\nAdding Node4...")
	ch.Add("Node4")
	for _, key := range testKeys {
		fmt.Printf("Key %s is now assigned to %s\n", key, ch.Get(key))
	}
	
	// 移除节点
	fmt.Println("\nRemoving Node2...")
	ch.Remove("Node2")
	for _, key := range testKeys {
		fmt.Printf("Key %s is now assigned to %s\n", key, ch.Get(key))
	}
}