Example consistent hashing golang

consistent-hashing.go

package main

import (
	"crypto/md5"
	"fmt"
	"sort"
	"strconv"
)

// ConsistentHashing represents the structure for consistent hashing
type ConsistentHashing struct {
	replicas    int               // Number of virtual nodes per physical node
	hashRing    map[int]string    // Map hash to node
	sortedHashes []int            // Sorted hash keys
}

// NewConsistentHashing initializes the hash ring
func NewConsistentHashing(replicas int) *ConsistentHashing {
	return &ConsistentHashing{
		replicas:    replicas,
		hashRing:    make(map[int]string),
		sortedHashes: []int{},
	}
}

// hash generates a hash for a given key
func (c *ConsistentHashing) hash(key string) int {
	hash := md5.Sum([]byte(key))
	// Convert hash to int (using first 4 bytes for simplicity)
	return int(hash[0])<<24 | int(hash[1])<<16 | int(hash[2])<<8 | int(hash[3])
}

// AddNode adds a new node to the hash ring
func (c *ConsistentHashing) AddNode(node string) {
	for i := 0; i < c.replicas; i++ {
		virtualKey := node + ":" + strconv.Itoa(i)
		hash := c.hash(virtualKey)
		c.hashRing[hash] = node
		c.sortedHashes = append(c.sortedHashes, hash)
	}
	sort.Ints(c.sortedHashes)
}

// RemoveNode removes a node from the hash ring
func (c *ConsistentHashing) RemoveNode(node string) {
	for i := 0; i < c.replicas; i++ {
		virtualKey := node + ":" + strconv.Itoa(i)
		hash := c.hash(virtualKey)
		delete(c.hashRing, hash)
		for idx, val := range c.sortedHashes {
			if val == hash {
				c.sortedHashes = append(c.sortedHashes[:idx], c.sortedHashes[idx+1:]...)
				break
			}
		}
	}
}

// GetNode finds the appropriate node for a given key
func (c *ConsistentHashing) GetNode(key string) string {
	if len(c.sortedHashes) == 0 {
		return ""
	}
	hash := c.hash(key)
	// Find the first node greater than or equal to the hash
	idx := sort.Search(len(c.sortedHashes), func(i int) bool {
		return c.sortedHashes[i] >= hash
	})
	// If no such node exists, wrap around to the first node
	if idx == len(c.sortedHashes) {
		idx = 0
	}
	return c.hashRing[c.sortedHashes[idx]]
}

// Example usage
func main() {
	ch := NewConsistentHashing(3) // 3 virtual nodes per physical node
	ch.AddNode("Node1")
	ch.AddNode("Node2")
	ch.AddNode("Node3")

	fmt.Println("Key1 is mapped to:", ch.GetNode("Key1"))
	fmt.Println("Key2 is mapped to:", ch.GetNode("Key2"))

	ch.RemoveNode("Node2")
	fmt.Println("Key1 after Node2 removal is mapped to:", ch.GetNode("Key1"))
	fmt.Println("Key2 after Node2 removal is mapped to:", ch.GetNode("Key2"))
}