S074 · Sharding & Partitioning Strategies

#!/usr/bin/env python3
"""consistent_hash.py — a working consistent hash ring in one file.
 
Run: python consistent_hash.py
 
Compares naive `hash % N` to consistent hashing with virtual nodes.
Shows: 'when you add a shard, how many keys move?'
"""
from __future__ import annotations
 
import bisect
import hashlib
from collections import Counter
 
 
def h(x: str) -> int:
    """Stable 128-bit hash truncated to 32 bits for readability."""
    return int(hashlib.md5(x.encode()).hexdigest(), 16) & 0xFFFF_FFFF
 
 
class NaiveHashRouter:
    """The `hash(key) % N` approach. Simple, terrible at resharding."""
    def __init__(self, nodes: list[str]):
        self.nodes = list(nodes)
 
    def route(self, key: str) -> str:
        return self.nodes[h(key) % len(self.nodes)]
 
    def add_node(self, node: str) -> None:
        self.nodes.append(node)
 
 
class ConsistentHashRouter:
    """Consistent hashing with virtual nodes for balance."""
    def __init__(self, nodes: list[str], virtual_per_node: int = 150):
        self.vpn = virtual_per_node
        self.ring: list[tuple[int, str]] = []  # (hash, node)
        for n in nodes:
            self._add(n)
 
    def _add(self, node: str) -> None:
        for v in range(self.vpn):
            self.ring.append((h(f"{node}#{v}"), node))
        self.ring.sort(key=lambda x: x[0])
 
    def route(self, key: str) -> str:
        if not self.ring:
            raise RuntimeError("empty ring")
        hv = h(key)
        # Find the first virtual node whose hash >= key's hash; wrap around.
        idx = bisect.bisect_left(self.ring, (hv, ""))
        if idx == len(self.ring):
            idx = 0
        return self.ring[idx][1]
 
    def add_node(self, node: str) -> None:
        self._add(node)
 
    def remove_node(self, node: str) -> None:
        self.ring = [(hv, n) for hv, n in self.ring if n != node]
 
 
def demo(router_class, name: str, N_INITIAL: int, N_KEYS: int) -> None:
    print(f"\n=== {name} ===")
    nodes = [f"node-{i}" for i in range(N_INITIAL)]
    router = router_class(nodes) if router_class is NaiveHashRouter \
        else router_class(nodes, virtual_per_node=150)
    keys = [f"user-{i}" for i in range(N_KEYS)]
 
    # Distribution before
    before = Counter(router.route(k) for k in keys)
    print(f"Before ({N_INITIAL} nodes):")
    for n, c in sorted(before.items()):
        pct = 100 * c / N_KEYS
        bar = "#" * int(pct)
        print(f"  {n:<10} {c:>6}  {bar}")
 
    # Snapshot mapping
    old_map = {k: router.route(k) for k in keys}
 
    # Add a node
    router.add_node(f"node-{N_INITIAL}")
    new_map = {k: router.route(k) for k in keys}
    moved = sum(1 for k in keys if old_map[k] != new_map[k])
    pct_moved = 100 * moved / N_KEYS
 
    print(f"After adding one node ({N_INITIAL + 1} nodes total):")
    after = Counter(new_map.values())
    for n, c in sorted(after.items()):
        pct = 100 * c / N_KEYS
        bar = "#" * int(pct)
        print(f"  {n:<10} {c:>6}  {bar}")
    print(f"Keys moved: {moved}/{N_KEYS} = {pct_moved:.1f}%")
    print(f"Ideal (1 / (N+1)): {100 / (N_INITIAL + 1):.1f}%")
 
 
def hot_partition_demo() -> None:
    """Simulate the 'celebrity user' effect."""
    print("\n=== Hot-partition simulation ===")
    router = ConsistentHashRouter([f"node-{i}" for i in range(4)])
    # 10 000 normal users, 1 celebrity
    events: list[str] = []
    for i in range(10_000):
        events.append(f"user-{i}")
    # Celebrity generates 100 000 events (all with same key)
    events.extend(["user-42"] * 100_000)
 
    load = Counter(router.route(k) for k in events)
    total = sum(load.values())
    print("Traffic per node (celebrity concentrates load):")
    for n, c in sorted(load.items()):
        pct = 100 * c / total
        bar = "#" * int(pct / 2)
        print(f"  {n:<10} {c:>7}  {pct:5.1f}%  {bar}")
    print("Fix: shard by (user_id, event_id % 16) so one hot user")
    print("     spreads across 16 sub-shards.")
 
 
if __name__ == "__main__":
    demo(NaiveHashRouter, "Naive hash % N", 4, 100_000)
    demo(ConsistentHashRouter, "Consistent hashing (150 vnodes/node)",
         4, 100_000)
    hot_partition_demo()