diff --git a/releasenotes/notes/hashring-algo-8a279397b8ff8a6a.yaml b/releasenotes/notes/hashring-algo-8a279397b8ff8a6a.yaml new file mode 100644 index 00000000..fbd81c74 --- /dev/null +++ b/releasenotes/notes/hashring-algo-8a279397b8ff8a6a.yaml @@ -0,0 +1,5 @@ +--- +features: + - | + ``HashRing`` now accepts a new argument ``hash_function``, allowing + using a different hash function. diff --git a/tooz/hashring.py b/tooz/hashring.py index 51016f2b..f6dac315 100644 --- a/tooz/hashring.py +++ b/tooz/hashring.py @@ -32,16 +32,28 @@ class HashRing(object): DEFAULT_PARTITION_NUMBER = 2**5 - def __init__(self, nodes, partitions=DEFAULT_PARTITION_NUMBER): + DEFAULT_HASH_FUNCTION = 'md5' + + def __init__(self, nodes, partitions=DEFAULT_PARTITION_NUMBER, + hash_function=DEFAULT_HASH_FUNCTION): """Create a new hashring. :param nodes: List of nodes where objects will be mapped onto. :param partitions: Number of partitions to spread objects onto. + :param hash_function: Hash function to use, one of supported by hashlib + :raises: ValueError if `hash_function` is not supported. """ + if hash_function not in hashlib.algorithms_available: + raise ValueError('Hash function %s is not supported on this ' + 'system, supported are %s' + % (hash_function, + ', '.join(hashlib.algorithms_available))) + self.nodes = {} self._ring = dict() self._partitions = [] self._partition_number = partitions + self._hash_function = hash_function self.add_nodes(set(nodes)) @@ -68,7 +80,7 @@ class HashRing(object): """ for node in nodes: key = utils.to_binary(node, 'utf-8') - key_hash = hashlib.md5(key) + key_hash = hashlib.new(self._hash_function, key) for r in range(self._partition_number * weight): key_hash.update(key) self._ring[self._hash2int(key_hash)] = node @@ -90,7 +102,7 @@ class HashRing(object): raise UnknownNode(node) key = utils.to_binary(node, 'utf-8') - key_hash = hashlib.md5(key) + key_hash = hashlib.new(self._hash_function, key) for r in range(self._partition_number * weight): key_hash.update(key) del self._ring[self._hash2int(key_hash)] @@ -102,7 +114,7 @@ class HashRing(object): return int(key.hexdigest(), 16) def _get_partition(self, data): - hashed_key = self._hash2int(hashlib.md5(data)) + hashed_key = self._hash2int(hashlib.new(self._hash_function, data)) position = bisect.bisect(self._partitions, hashed_key) return position if position < len(self._partitions) else 0 diff --git a/tooz/tests/test_hashring.py b/tooz/tests/test_hashring.py index eeb9edb9..d44c1d16 100644 --- a/tooz/tests/test_hashring.py +++ b/tooz/tests/test_hashring.py @@ -31,7 +31,7 @@ class HashRingTestCase(testcase.TestCase): # fake -> foo, bar, baz # fake-again -> bar, baz, foo - @mock.patch.object(hashlib, 'md5', autospec=True) + @mock.patch.object(hashlib, 'new', autospec=True) def test_hash2int_returns_int(self, mock_md5): r1 = 32 * 'a' r2 = 32 * 'b' @@ -45,12 +45,20 @@ class HashRingTestCase(testcase.TestCase): self.assertIn(int(r1, 16), ring._ring) self.assertIn(int(r2, 16), ring._ring) + mock_md5.assert_called_with('md5', mock.ANY) + def test_create_ring(self): nodes = {'foo', 'bar'} ring = hashring.HashRing(nodes) self.assertEqual(nodes, set(ring.nodes.keys())) self.assertEqual(2 ** 5 * 2, len(ring)) + def test_wrong_hash_function(self): + nodes = {'foo', 'bar'} + self.assertRaisesRegex(ValueError, 'is not supported', + hashring.HashRing, nodes, + hash_function='fold twice and leave to dry') + def test_add_node(self): nodes = {'foo', 'bar'} ring = hashring.HashRing(nodes)