"""
Based on https://pypi.python.org/pypi/hash_ring (copyright: 2008 by Amir Salihefendic; license: BSD)
I (William Stein) rewrote all the documentation, fixed some trivial bugs,
more compatible with node-hashring (e.g., vnode support), etc. BSD license.
Example:
import hashring
r = hashring.HashRing({'10.1.1.4':{'vnodes':256}, '10.1.2.4':{'vnodes':128}, '10.1.3.4':{'vnodes':128}})
r.range('foo')
['10.1.1.4', '10.1.3.4', '10.1.2.4']
"""
import math
import sys
from bisect import bisect
import hashlib
md5_constructor = hashlib.md5
class HashRing(object):
def __init__(self, nodes=None, weights=1, vnodes=40, replicas=4):
"""
`nodes` is a list of objects that have a proper __str__ representation.
`weights` is dictionary that sets weights to the nodes. The default
weight is that all nodes are equal.
"""
self.ring = dict()
self._sorted_keys = []
if isinstance(nodes, dict):
weights = dict([(n, nodes[n].get('weight',weights)) for n in nodes])
vnodes = dict([(n, nodes[n].get('vnodes',vnodes)) for n in nodes])
nodes = nodes.keys()
self.nodes = nodes
if not isinstance(weights, dict):
weights = dict([(n,weights) for n in nodes])
else:
for n in nodes:
if n not in weights:
weights[n] = 1
self.weights = weights
self.replicas = replicas
if not isinstance(vnodes, dict):
vnodes = dict([(n,vnodes) for n in nodes])
else:
for n in nodes:
if n not in vnodes:
vnodes[n] = 40
self.vnodes = vnodes
self._generate_circle()
def _generate_circle(self):
"""
Generates the circle.
"""
total_weight = 0
for node in self.nodes:
total_weight += self.weights.get(node, 1)
for node in self.nodes:
weight = self.weights[node]
factor = (self.vnodes[node]*len(self.nodes)*weight) // total_weight
for j in xrange(0, int(factor)):
b_key = self._hash_digest('%s-%s'%(node, j))
for i in xrange(0, self.replicas):
key = self._hash_val(b_key, lambda x: x+i*4)
self.ring[key] = node
self._sorted_keys.append(key)
self._sorted_keys.sort()
def get_node(self, string_key):
"""
Given a string key, return some corresponding node in the hash ring.
If the hash ring is empty, return `None`.
"""
pos = self.get_node_pos(string_key)
if pos is None:
return None
return self.ring[ self._sorted_keys[pos] ]
def range(self, key, size=None, distinct=True):
if size is None:
return list(self.iterate_nodes(key, distinct=distinct))
v = []
for k in self.iterate_nodes(key, distinct=distinct):
v.append(k)
if len(v) >= size:
return v
return v
def __getitem__(self, string_key):
"""
Given a string key, return node that hold the key.
"""
return self.range(string_key)
def get_node_pos(self, string_key):
"""
Given a string key, return a corresponding node in the hash ring
along with it's position in the ring.
If the hash ring is empty, `None` is returned.
"""
if not self.ring:
return None
key = self.gen_key(string_key)
nodes = self._sorted_keys
pos = bisect(nodes, key)
if pos == len(nodes):
return 0
else:
return pos
def iterate_nodes(self, string_key, distinct=True):
"""
Given a string key, return iterator over all nodes that hold the key.
The generator iterates one time through the ring
starting at the correct position.
if `distinct` is set, then the nodes returned will be unique,
i.e. no virtual copies will be returned.
"""
if not self.ring:
yield None, None
returned_values = set()
def distinct_filter(value):
if str(value) not in returned_values:
returned_values.add(str(value))
return value
pos = self.get_node_pos(string_key)
for key in self._sorted_keys[pos:]:
val = distinct_filter(self.ring[key])
if val:
yield val
for i, key in enumerate(self._sorted_keys):
if i < pos:
val = distinct_filter(self.ring[key])
if val:
yield val
def gen_key(self, key):
"""
Given a string key return a long value, which
represents a place on the hash ring.
md5 is currently used because it mixes well.
"""
b_key = self._hash_digest(key)
return self._hash_val(b_key, lambda x: x)
def _hash_val(self, b_key, entry_fn):
return (( b_key[entry_fn(3)] << 24)
|(b_key[entry_fn(2)] << 16)
|(b_key[entry_fn(1)] << 8)
| b_key[entry_fn(0)] )
def _hash_digest(self, key):
m = md5_constructor()
m.update(key)
return map(ord, m.digest())
