diff --git a/cluster/hashring.go b/cluster/hashring.go index 0d5f99c..a469962 100644 --- a/cluster/hashring.go +++ b/cluster/hashring.go @@ -5,10 +5,12 @@ import ( "encoding/binary" "fmt" "sort" + "sync" ) // ConsistentHashRing implements a consistent hash ring for shard distribution type ConsistentHashRing struct { + mu sync.RWMutex ring map[uint32]string // hash -> node ID sortedHashes []uint32 // sorted hash keys nodes map[string]bool // active nodes @@ -35,6 +37,9 @@ func NewConsistentHashRingWithConfig(config HashRingConfig) *ConsistentHashRing // AddNode adds a node to the hash ring func (chr *ConsistentHashRing) AddNode(nodeID string) { + chr.mu.Lock() + defer chr.mu.Unlock() + if chr.nodes[nodeID] { return // Node already exists } @@ -56,6 +61,9 @@ func (chr *ConsistentHashRing) AddNode(nodeID string) { // RemoveNode removes a node from the hash ring func (chr *ConsistentHashRing) RemoveNode(nodeID string) { + chr.mu.Lock() + defer chr.mu.Unlock() + if !chr.nodes[nodeID] { return // Node doesn't exist } @@ -76,6 +84,9 @@ func (chr *ConsistentHashRing) RemoveNode(nodeID string) { // GetNode returns the node responsible for a given key func (chr *ConsistentHashRing) GetNode(key string) string { + chr.mu.RLock() + defer chr.mu.RUnlock() + if len(chr.sortedHashes) == 0 { return "" } @@ -103,6 +114,9 @@ func (chr *ConsistentHashRing) hash(key string) uint32 { // GetNodes returns all active nodes in the ring func (chr *ConsistentHashRing) GetNodes() []string { + chr.mu.RLock() + defer chr.mu.RUnlock() + nodes := make([]string, 0, len(chr.nodes)) for nodeID := range chr.nodes { nodes = append(nodes, nodeID) @@ -112,6 +126,9 @@ func (chr *ConsistentHashRing) GetNodes() []string { // IsEmpty returns true if the ring has no nodes func (chr *ConsistentHashRing) IsEmpty() bool { + chr.mu.RLock() + defer chr.mu.RUnlock() + return len(chr.nodes) == 0 }