hash_ring.py 1.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. import bisect
  2. import hashlib
  3. from typing import Dict, Iterable, Iterator, List, Optional, Tuple
  4. class HashRing:
  5. nodes: List[str] = []
  6. def __init__(self, nodes: Iterable[str] = (), replicas: int = 128) -> None:
  7. self.replicas: int = replicas
  8. self.ring: Dict[str, str] = {}
  9. self.sorted_keys: List[str] = []
  10. for node in nodes:
  11. self.add_node(node)
  12. def add_node(self, node: str) -> None:
  13. self.nodes.append(node)
  14. for x in range(self.replicas):
  15. _key = f"{node}:{x}"
  16. _hash = hashlib.sha256(_key.encode()).hexdigest()
  17. self.ring[_hash] = node
  18. self.sorted_keys.append(_hash)
  19. self.sorted_keys.sort()
  20. def remove_node(self, node: str) -> None:
  21. self.nodes.remove(node)
  22. for x in range(self.replicas):
  23. _hash = hashlib.sha256(f"{node}:{x}".encode()).hexdigest()
  24. del self.ring[_hash]
  25. self.sorted_keys.remove(_hash)
  26. def get_node(self, key: str) -> Optional[str]:
  27. n, i = self.get_node_pos(key)
  28. return n
  29. def get_node_pos(self, key: str) -> Tuple[Optional[str], Optional[int]]:
  30. if len(self.ring) == 0:
  31. return None, None
  32. _hash = hashlib.sha256(key.encode()).hexdigest()
  33. idx = bisect.bisect(self.sorted_keys, _hash)
  34. idx = min(idx - 1, (self.replicas * len(self.nodes)) - 1)
  35. return self.ring[self.sorted_keys[idx]], idx
  36. def iter_nodes(self, key: str) -> Iterator[Tuple[Optional[str], Optional[str]]]:
  37. if len(self.ring) == 0:
  38. yield None, None
  39. node, pos = self.get_node_pos(key)
  40. for k in self.sorted_keys[pos:]:
  41. yield k, self.ring[k]
  42. def __call__(self, key: str) -> Optional[str]:
  43. return self.get_node(key)