Consistent Hash Ring With Replicas

from __future__ import annotations

import bisect
import hashlib
from collections.abc import Iterable


class ConsistentHashRing:
    def __init__(
        self,
        nodes: Iterable[str] = (),
        *,
        replicas: int = 64,
    ) -> None:
        if (
            not isinstance(replicas, int)
            or isinstance(replicas, bool)
            or replicas < 1
        ):
            raise ValueError("replicas must be a positive integer")
        self.replicas = replicas
        self._ring: list[int] = []
        self._owners: dict[int, str] = {}
        self._nodes: set[str] = set()
        for node in nodes:
            self.add_node(node)

    @staticmethod
    def _hash(value: str) -> int:
        digest = hashlib.sha256(value.encode("utf-8")).digest()
        return int.from_bytes(digest[:8], "big")

    @staticmethod
    def _validate_node(node: str) -> None:
        if not isinstance(node, str) or not node.strip():
            raise ValueError("node must be a non-empty string")

    def add_node(self, node: str) -> bool:
        self._validate_node(node)
        if node in self._nodes:
            return False
        for replica in range(self.replicas):
            point = self._hash(f"{node}#{replica}")
            if point in self._owners:
                continue
            bisect.insort(self._ring, point)
            self._owners[point] = node
        self._nodes.add(node)
        return True

    def remove_node(self, node: str) -> bool:
        if node not in self._nodes:
            return False
        for replica in range(self.replicas):
            point = self._hash(f"{node}#{replica}")
            if self._owners.get(point) != node:
                continue
            del self._owners[point]
            index = bisect.bisect_left(self._ring, point)
            if index < len(self._ring) and self._ring[index] == point:
                self._ring.pop(index)
        self._nodes.remove(node)
        return True

    def get_node(self, key: str) -> str:
        if not isinstance(key, str) or not key:
            raise ValueError("key must be a non-empty string")
        if not self._ring:
            raise RuntimeError("hash ring is empty")
        point = self._hash(key)
        index = bisect.bisect_left(self._ring, point)
        if index == len(self._ring):
            index = 0
        return self._owners[self._ring[index]]

    @property
    def nodes(self) -> tuple[str, ...]:
        return tuple(sorted(self._nodes))


def main() -> None:
    from collections import Counter

    ring = ConsistentHashRing(["cache-a", "cache-b", "cache-c"], replicas=64)
    keys = [f"user:{index}" for index in range(12)]
    before = {key: ring.get_node(key) for key in keys}
    print(Counter(before.values()))

    ring.remove_node("cache-b")
    ring.add_node("cache-d")
    after = {key: ring.get_node(key) for key in keys}
    print(Counter(after.values()))
    print("moved keys:", sum(before[key] != after[key] for key in keys))
    print(ring.nodes)

    try:
        ConsistentHashRing().get_node("user:1")
    except RuntimeError as error:
        print(error)


if __name__ == "__main__":
    main()