一致性哈希 Python 实现

Consistent-Hashing.py

import hashlib
from bintrees.rbtree import RBTree
from types import GeneratorType


def md5(data: str) -> int:
    """md5 util function"""
    return int(hashlib.md5(data.encode('utf-8')).hexdigest(), 16)


class Node:
    """Node Implemention
    >>> n1 = Node('192.168.1.203')
    >>> n1.name
    '192.168.1.203'
    >>> n2 = Node('192.168.1.203', prefix='slave')
    >>> n2.name
    'slave#192.168.1.203'
    >>> n1.nid != n2.nid
    True
    """

    algorithms = {
        'md5': md5
    }

    def __init__(self, node_name: str, prefix: str='',
                 nid: int=None, kind: str='md5'):
        self._node_name = node_name
        if prefix:
            self._node_name = '{}#{}'.format(prefix, self._node_name)
        self._nid = nid or self.algorithms[kind](self._node_name)

    @property
    def name(self) -> str:
        return self._node_name

    @property
    def nid(self) -> int:
        return self._nid


class ConsistentHashing:
    """Consistent Hashing Implemention
    >>> ch = ConsistentHashing()
    >>> ch.size
    255
    >>> ch.join(Node('192.168.1.2'))
    >>> ch.node_count
    1
    >>> ch.join(Node('192.168.1.2', prefix='slave'))
    >>> ch.node_count
    2
    >>> ch.join(Node('192.168.1.65'))
    >>> ch.join(Node('192.168.1.232'))
    >>> [nid for nid, _ in ch.nodes]
    [83, 135, 141, 243]
    >>> ch.get_location(0)
    83
    >>> ch.get_location(39)
    83
    >>> ch.get_location(100)
    135
    >>> ch.get_location(141)
    141
    >>> ch.get_location(244)
    83
    """

    def __init__(self, size: int=255):
        self._rbtree = RBTree()
        self._size = size

    @property
    def size(self)-> int:
        return self._size

    @property
    def nodes(self)-> GeneratorType:
        """a generator"""
        return self._rbtree.items()

    @property
    def node_count(self)-> int:
        return self._rbtree.count

    def join(self, node):
        self._rbtree.insert(node.nid % self._size, node)

    def depart(self, node):
        self._rbtree.remove(node.nid % self._size, node)

    def get_location(self, elem)-> int:
        try:
            location = self._rbtree.ceiling_key(elem)
        except KeyError:
            location = self._rbtree.min_key()
        return location


if __name__ == '__main__':
    import doctest
    doctest.testmod(verbose=False)