Merge "hashring: allow choosing hash function"

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.

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
 

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)