Consistent Hashing (Go & Python)

"""Consistent hash ring with virtual nodes.
`bisect` is a clean fit here: the ring is nothing more than a sorted list of
integers, and "find the first node clockwise from this hash" is exactly the
"find the insertion point in a sorted sequence" problem bisect exists to
solve in O(log n) -- no hand-rolled binary search needed.
"""
from __future__ import annotations
import bisect
import hashlib
import threading
from typing import Optional
class HashRing:
    """Consistent hash ring with virtual nodes."""
    def __init__(self, replicas: int = 150) -> None:
        """Create an empty ring.
        Args:
            replicas: virtual nodes created per physical node. Higher means
                smoother load distribution at the cost of memory. 100-200
                is a common production value.
        """
        self._replicas = replicas
        self._ring: list[int] = []
        self._ring_map: dict[int, str] = {}
        self._lock = threading.RLock()
    @staticmethod
    def _hash(key: str) -> int:
        """Stable 32-bit hash for ring placement.
        Python's builtin hash() is salted per-process (PYTHONHASHSEED) and
        would produce a different ring layout on every run -- the opposite
        of what a stable, shareable hash ring needs. md5 (truncated to 4
        bytes) is deterministic across processes and machines instead.
        """
        digest = hashlib.md5(key.encode("utf-8")).digest()
        return int.from_bytes(digest[:4], byteorder="big")
    def add_node(self, node_id: str) -> None:
        """Register a physical node, adding `replicas` virtual points."""
        with self._lock:
            for i in range(self._replicas):
                pos = self._hash(f"{node_id}#{i}")
                self._ring_map[pos] = node_id
                bisect.insort(self._ring, pos)
    def remove_node(self, node_id: str) -> None:
        """Remove a physical node and all of its virtual points."""
        with self._lock:
            self._ring = [p for p in self._ring if self._ring_map[p] != node_id]
            self._ring_map = {
                p: n for p, n in self._ring_map.items() if n != node_id
            }
    def get_node(self, key: str) -> Optional[str]:
        """Return the physical node owning `key`, or None if the ring is empty.
        Walks clockwise from key's hash position: bisect_left finds the
        index of the first ring entry >= that hash in O(log n); the index
        wraps to 0 if the hash falls past the last entry -- that wraparound
        is the entire "ring" in "hash ring."
        """
        with self._lock:
            if not self._ring:
                return None
            h = self._hash(key)
            idx = bisect.bisect_left(self._ring, h)
            if idx == len(self._ring):
                idx = 0
            return self._ring_map[self._ring[idx]]
    def nodes(self) -> list[str]:
        """Return the distinct physical nodes currently on the ring."""
        with self._lock:
            return sorted(set(self._ring_map.values()))

import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.TreeMap;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.zip.CRC32;
/**
 * Consistent hash ring with virtual nodes, backed by a TreeMap. A TreeMap
 * keeps ring positions sorted by key automatically, so "find the first node
 * clockwise from this hash" is just ceilingKey() (the built-in equivalent of
 * Go's sort.Search / Python's bisect.bisect_left) with a wraparound to
 * firstKey() when the hash falls past the last entry -- no hand-rolled
 * binary search needed.
 */
public class HashRing {
    private final int replicas; // virtual nodes per physical node
    private final TreeMap<Long, String> ring = new TreeMap<>(); // ring position -> physical node ID
    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
    /**
     * Creates a ring with the given number of virtual nodes (replicas) per
     * physical node. Higher replicas = smoother distribution, more memory.
     * 100-200 is a common production value.
     */
    public HashRing(int replicas) {
        this.replicas = replicas;
    }
    /**
     * Produces an unsigned 32-bit hash for any string key. CRC32 is fast and
     * sufficiently uniform for ring placement; production systems sometimes
     * use a stronger hash (murmur3, fnv) but CRC32 is fine here and in the
     * standard library, no extra dependency.
     */
    private static long hashKey(String key) {
        CRC32 crc = new CRC32();
        crc.update(key.getBytes(StandardCharsets.UTF_8));
        return crc.getValue();
    }
    /** Registers a physical node, creating `replicas` virtual points on the ring for it. */
    public void addNode(String nodeId) {
        lock.writeLock().lock();
        try {
            for (int i = 0; i < replicas; i++) {
                String vNodeKey = nodeId + "#" + i;
                long pos = hashKey(vNodeKey);
                ring.put(pos, nodeId);
            }
        } finally {
            lock.writeLock().unlock();
        }
    }
    /** Removes a physical node and all of its virtual points. */
    public void removeNode(String nodeId) {
        lock.writeLock().lock();
        try {
            ring.values().removeIf(id -> id.equals(nodeId));
        } finally {
            lock.writeLock().unlock();
        }
    }
    /**
     * Returns the physical node responsible for key: the first node found
     * walking clockwise from key's hash position, wrapping around to the
     * ring's first entry if the hash falls past the last one.
     */
    public Optional<String> getNode(String key) {
        lock.readLock().lock();
        try {
            if (ring.isEmpty()) {
                return Optional.empty();
            }
            long h = hashKey(key);
            Map.Entry<Long, String> entry = ring.ceilingEntry(h);
            if (entry == null) {
                entry = ring.firstEntry(); // wrap around to the start of the ring
            }
            return Optional.of(entry.getValue());
        } finally {
            lock.readLock().unlock();
        }
    }
    /** Returns the current distinct set of physical nodes on the ring. */
    public List<String> nodes() {
        lock.readLock().lock();
        try {
            return new ArrayList<>(new HashSet<>(ring.values()));
        } finally {
            lock.readLock().unlock();
        }
    }
}

"""Tests for the consistent hash ring."""
from hash_ring import HashRing
def test_basic_assignment_is_deterministic() -> None:
    ring = HashRing(replicas=100)
    ring.add_node("nodeA")
    ring.add_node("nodeB")
    ring.add_node("nodeC")
    node = ring.get_node("user:1234")
    assert node is not None
    # Same key always maps to the same node while the ring is unchanged.
    assert ring.get_node("user:1234") == node
def test_minimal_reshuffle_on_add() -> None:
    ring = HashRing(replicas=100)
    for node_id in ("nodeA", "nodeB", "nodeC"):
        ring.add_node(node_id)
    keys = [f"key:{i}" for i in range(1000)]
    before = {k: ring.get_node(k) for k in keys}
    ring.add_node("nodeD")  # add a 4th node
    moved = sum(1 for k in keys if ring.get_node(k) != before[k])
    # With consistent hashing, expect roughly 1/4 of keys to move (the new
    # node's fair share), not all 1000.
    assert moved < 400, f"too many keys moved on node add: {moved}/1000"
def test_remove_node_redistributes() -> None:
    ring = HashRing(replicas=100)
    for node_id in ("nodeA", "nodeB", "nodeC"):
        ring.add_node(node_id)
    key = "session:abc"
    before = ring.get_node(key)
    ring.remove_node(before)  # remove whichever node currently owns this key
    after = ring.get_node(key)
    assert after is not None
    assert after != before
def test_empty_ring_returns_none() -> None:
    ring = HashRing(replicas=100)
    assert ring.get_node("anything") is None

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
/** Tests for the consistent hash ring. */
public class HashRingTest {
    public static void main(String[] args) {
        testBasicAssignmentIsDeterministic();
        testMinimalReshuffleOnAdd();
        testRemoveNodeRedistributes();
        testEmptyRing();
        System.out.println("ALL TESTS PASSED");
    }
    static void testBasicAssignmentIsDeterministic() {
        HashRing ring = new HashRing(100);
        ring.addNode("nodeA");
        ring.addNode("nodeB");
        ring.addNode("nodeC");
        Optional<String> node = ring.getNode("user:1234");
        assertTrue(node.isPresent(), "expected a node for key");
        // Same key always maps to the same node while the ring is unchanged.
        assertEquals(node, ring.getNode("user:1234"), "getNode not deterministic");
        System.out.println("PASS testBasicAssignmentIsDeterministic");
    }
    static void testMinimalReshuffleOnAdd() {
        HashRing ring = new HashRing(100);
        ring.addNode("nodeA");
        ring.addNode("nodeB");
        ring.addNode("nodeC");
        String[] keys = new String[1000];
        Map<String, Optional<String>> before = new HashMap<>();
        for (int i = 0; i < keys.length; i++) {
            keys[i] = "key:" + i;
            before.put(keys[i], ring.getNode(keys[i]));
        }
        ring.addNode("nodeD"); // add a 4th node
        int moved = 0;
        for (String k : keys) {
            if (!ring.getNode(k).equals(before.get(k))) {
                moved++;
            }
        }
        // With consistent hashing, expect roughly 1/4 of keys to move (the
        // new node's fair share), not all 1000.
        assertTrue(moved < 400, "too many keys moved on node add: " + moved + "/1000");
        System.out.println("PASS testMinimalReshuffleOnAdd (" + moved + "/1000 keys moved)");
    }
    static void testRemoveNodeRedistributes() {
        HashRing ring = new HashRing(100);
        ring.addNode("nodeA");
        ring.addNode("nodeB");
        ring.addNode("nodeC");
        String key = "session:abc";
        String before = ring.getNode(key).orElseThrow();
        ring.removeNode(before); // remove whichever node currently owns this key
        Optional<String> after = ring.getNode(key);
        assertTrue(after.isPresent(), "expected a node after removal");
        assertTrue(!after.get().equals(before), "key should have moved off the removed node");
        System.out.println("PASS testRemoveNodeRedistributes");
    }
    static void testEmptyRing() {
        HashRing ring = new HashRing(100);
        assertTrue(ring.getNode("anything").isEmpty(), "expected empty Optional on empty ring");
        System.out.println("PASS testEmptyRing");
    }
    static void assertTrue(boolean cond, String msg) {
        if (!cond) {
            throw new AssertionError(msg);
        }
    }
    static void assertEquals(Object a, Object b, String msg) {
        if (!a.equals(b)) {
            throw new AssertionError(msg + ": " + a + " != " + b);
        }
    }
}