Implementing a Lock-Free Consistent Hashing Ring with Bounded Loads in Go — code snippets

snippet-1.go

package main

import (
	"sync/atomic"
)

// Node represents a physical backend server (e.g., a Redis or gRPC instance).
type Node struct {
	ID   string
	Addr string
	Load int64 // Dynamic load counter incremented/decremented atomically
}

// VirtualNode represents a point on the consistent hashing ring.
type VirtualNode struct {
	Hash uint32
	Node *Node
}

// Ring represents an immutable state of the hash ring at a specific point in time.
type Ring struct {
	vNodes  []VirtualNode
	nodes   []*Node
	epsilon float64 // Over-capacity factor, e.g., 0.25 for 125% capacity limit
}

snippet-2.go

package main

import (
	"fmt"
	"hash/fnv"
	"sort"
)

// hashKey calculates the FNV-1a hash of a string key.
func hashKey(key string) uint32 {
	h := fnv.New32a()
	_, _ = h.Write([]byte(key))
	return h.Sum32()
}

// NewRing constructs an immutable Ring instance with virtual nodes.
func NewRing(nodes []*Node, vNodeCount int, epsilon float64) *Ring {
	if epsilon <= 0 {
		epsilon = 0.25 // Default to 25% load limit headroom
	}

	vNodes := make([]VirtualNode, 0, len(nodes)*vNodeCount)
	for _, node := range nodes {
		for i := 0; i < vNodeCount; i++ {
			// Generate deterministic virtual node identifier
			vNodeKey := fmt.Sprintf("%s#%d", node.ID, i)
			vNodes = append(vNodes, VirtualNode{
				Hash: hashKey(vNodeKey),
				Node: node,
			})
		}
	}

	// Sort the virtual nodes by hash value to form the ring structure
	sort.Slice(vNodes, func(i, j int) bool {
		return vNodes[i].Hash < vNodes[j].Hash
	})

	return &Ring{
		vNodes:  vNodes,
		nodes:   nodes,
		epsilon: epsilon,
	}
}

snippet-3.go

package main

import (
	"sync"
	"sync/atomic"
)

// ConsistentHashRing manages the atomic swapping of immutable Ring states.
type ConsistentHashRing struct {
	active     atomic.Pointer[Ring]
	writeMutex sync.Mutex // Protects the write path from concurrent modifications
	vNodeCount int
	epsilon    float64
}

// NewConsistentHashRing initializes a new consistent hash ring manager.
func NewConsistentHashRing(nodes []*Node, vNodeCount int, epsilon float64) *ConsistentHashRing {
	chr := &ConsistentHashRing{
		vNodeCount: vNodeCount,
		epsilon:    epsilon,
	}
	chr.active.Store(NewRing(nodes, vNodeCount, epsilon))
	return chr
}

// GetRing returns the current active ring state atomically. Lock-free.
func (chr *ConsistentHashRing) GetRing() *Ring {
	return chr.active.Load()
}

// AddNode adds a new physical node, rebuilds the ring, and swaps it atomically.
func (chr *ConsistentHashRing) AddNode(node *Node) {
	chr.writeMutex.Lock()
	defer chr.writeMutex.Unlock()

	current := chr.active.Load()

	// Clone the existing nodes array and append the new node
	newNodes := make([]*Node, len(current.nodes), len(current.nodes)+1)
	copy(newNodes, current.nodes)
	newNodes = append(newNodes, node)

	// Rebuild and swap the active ring pointer
	chr.active.Store(NewRing(newNodes, chr.vNodeCount, chr.epsilon))
}

// RemoveNode removes a physical node by its ID, rebuilds the ring, and swaps.
func (chr *ConsistentHashRing) RemoveNode(id string) {
	chr.writeMutex.Lock()
	defer chr.writeMutex.Unlock()

	current := chr.active.Load()
	newNodes := make([]*Node, 0, len(current.nodes))
	for _, node := range current.nodes {
		if node.ID != id {
			newNodes = append(newNodes, node)
		}
	}

	chr.active.Store(NewRing(newNodes, chr.vNodeCount, chr.epsilon))
}

snippet-4.go

package main

import (
	"math"
	"sort"
	"sync/atomic"
)

// Get routes a key to a physical Node ensuring the node's load does not exceed capacity.
// It returns the selected Node and a release callback to decrement the load when done.
func (r *Ring) Get(key string) (*Node, func()) {
	if len(r.nodes) == 0 {
		return nil, func() {}
	}

	// 1. Calculate active total load across all physical nodes
	var totalLoad int64
	for _, node := range r.nodes {
		totalLoad += atomic.LoadInt64(&node.Load)
	}

	avgLoad := float64(totalLoad) / float64(len(r.nodes))
	// Dynamic capacity threshold: C = ceil(avgLoad * (1 + epsilon))
	capacityLimit := int64(math.Ceil(avgLoad * (1.0 + r.epsilon)))
	if capacityLimit < 1 {
		capacityLimit = 1 // Ensure headroom exists when total load is 0
	}

	keyHash := hashKey(key)

	// 2. Binary search to locate the starting virtual node clockwise
	idx := sort.Search(len(r.vNodes), func(i int) bool {
		return r.vNodes[i].Hash >= keyHash
	})

	// 3. Walk clockwise to find the first candidate node under capacity
	var selectedNode *Node
	for i := 0; i < len(r.vNodes); i++ {
		vIdx := (idx + i) % len(r.vNodes)
		candidate := r.vNodes[vIdx].Node

		if atomic.LoadInt64(&candidate.Load) < capacityLimit {
			selectedNode = candidate
			break
		}
	}

	// 4. Fallback: If all nodes are overloaded, route to the absolute least-loaded node
	if selectedNode == nil {
		selectedNode = r.nodes[0]
		minLoad := atomic.LoadInt64(&selectedNode.Load)
		for _, node := range r.nodes[1:] {
			load := atomic.LoadInt64(&node.Load)
			if load < minLoad {
				selectedNode = node
				minLoad = load
			}
		}
	}

	// 5. Commit the load to the selected node
	atomic.AddInt64(&selectedNode.Load, 1)

	// Return node and release closure to decrement the load
	return selectedNode, func() {
		atomic.AddInt64(&selectedNode.Load, -1)
	}
}

snippet-5.go

package main

import (
	"context"
	"errors"
	"net/http"
	"time"
)

// HTTPClientRouter routes HTTP requests to backends using the lock-free ring.
type HTTPClientRouter struct {
	ringRing   *ConsistentHashRing
	httpClient *http.Client
}

func NewHTTPClientRouter(ringRing *ConsistentHashRing) *HTTPClientRouter {
	return &HTTPClientRouter{
		ringRing: ringRing,
		httpClient: &http.Client{
			Timeout: 3 * time.Second,
		},
	}
}

// SendRoutedRequest routes an HTTP request based on key to a node address.
func (router *HTTPClientRouter) SendRoutedRequest(ctx context.Context, routingKey string, path string) (*http.Response, error) {
	// Retrieve active ring atomically (lock-free)
	ring := router.ringRing.GetRing()
	
	// Get target node and the release callback
	node, release := ring.Get(routingKey)
	if node == nil {
		return nil, errors.New("no available nodes on ring")
	}
	
	// Decrement node load when request completes (success or failure)
	defer release()

	url := "http://" + node.Addr + path
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}

	resp, err := router.httpClient.Do(req)
	if err != nil {
		return nil, err
	}

	return resp, nil
}

snippet-6.go

package main

import (
	"sync"
	"testing"
)

// MutexConsistentHashRing protects read and write paths with a standard Read-Write Mutex.
type MutexConsistentHashRing struct {
	sync.RWMutex
	nodes []string
}

func (m *MutexConsistentHashRing) Get(key string) string {
	m.RLock()
	defer m.RUnlock()
	if len(m.nodes) == 0 {
		return ""
	}
	// Simulated consistent hashing binary search lookup
	return m.nodes[0]
}

// BenchmarkRWMutexReadPath demonstrates RWMutex overhead under high thread contention.
func BenchmarkRWMutexReadPath(b *testing.B) {
	mRing := &MutexConsistentHashRing{
		nodes: []string{"node-1", "node-2", "node-3", "node-4"},
	}

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			_ = mRing.Get("my-routing-key")
		}
	})
}

// BenchmarkLockFreeReadPath demonstrates CPU scalability of our atomic read path.
func BenchmarkLockFreeReadPath(b *testing.B) {
	nodes := []*Node{
		{ID: "node-1", Addr: "10.0.0.1:8080"},
		{ID: "node-2", Addr: "10.0.0.2:8080"},
		{ID: "node-3", Addr: "10.0.0.3:8080"},
		{ID: "node-4", Addr: "10.0.0.4:8080"},
	}
	chr := NewConsistentHashRing(nodes, 150, 0.25)

	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			ring := chr.GetRing()
			node, release := ring.Get("my-routing-key")
			if node != nil {
				release()
			}
		}
	})
}