tarekziade/distribution.py

## distribution.py
""" Consistent load-balancing.

We have a few servers and we want a load-balancer to
distribute incoming requests across them in a deterministic
and consistent way - without keeping any counter to make the
decision.

Removing a backend server should not impact users on other
servers.

Adding a backend will generate the redistribution of
users across other servers.

The goal is to come up with the best algorithm
for 1M users across 5 servers. Speed is a bonus.

Two known techniques:

- RendezVous : https://en.wikipedia.org/wiki/Rendezvous_hashing#Comparison_With_Consistent_Hashing
- Consistent Hashing: https://en.wikipedia.org/wiki/Consistent_hashing

Consistent Hashing implementation inspired by:
http://techspot.zzzeek.org/2012/07/07/the-absolutely-simplest-consistent-hashing-example

Also, good read on hashes:

http://programmers.stackexchange.com/questions/49550/which-hashing-algorithm-is-best-for-uniqueness-and-speed/145633#145633
"""
import hashlib
import bisect
from collections import defaultdict
import binascii
import time
from functools import wraps


class CollisionError(Exception):
    pass


_collisions = {}


def catch_collision(func):
    @wraps(func)
    def _catch(key):
        res = func(key)
        if res in _collisions and key != _collisions[res]:
            raise CollisionError('%s and %s with %s' % (key, _collisions[res],
                func))
        _collisions[res] = key
        return res
    return _catch


@catch_collision
def fnv32a(key):
    hval = 0x811c9dc5
    fnv_32_prime = 0x01000193
    uint32_max = 2 ** 32
    for s in key:
        hval = hval ^ ord(s)
        hval = (hval * fnv_32_prime) % uint32_max
    return hval


@catch_collision
def sha512(key):
    return long(hashlib.sha512(key).hexdigest(), 16)


@catch_collision
def sha256(key):
    return long(hashlib.sha256(key).hexdigest(), 16)


@catch_collision
def md5(key):
    return long(hashlib.md5(key).hexdigest(), 16)


class RendezVous(object):

    def __init__(self, ips=None, hash=md5):
        if ips is None:
            ips = []
        self.ips = ips
        self._hash = hash

    def __str__(self):
        return '<RendezVous with %s hash>' % self._hash

    def add(self, ip):
        self.ips.append(ip)

    def remove(self, ip):
        self.ips.remove(ip)

    def select(self, key):
        high_score = -1
        winner = None
        for ip in self.ips:
            score = self._hash("%s-%s" % (str(ip), str(key)))
            if score > high_score:
                high_score, winner = score, ip

            elif score == high_score:
                high_score, winner = score, max(str(ip), str(winner))
        return winner


def _repl(name, index):
    return '%s:%d' % (name, index)


class ConsistentHashing(object):

    def __init__(self, ips=[], replicas=200, hash=md5):
        self._ips = {}
        self._hashed_ips = []
        self.replicas = replicas
        self._hash = hash

        for ip in ips:
            self.add(ip)

    def __str__(self):
        return '<ConsistentHashing with %s hash>' % self._hash

    def add(self, ip):
        for i in range(self.replicas):
            sip = _repl(ip, i)
            hashed = self._hash(sip)
            self._ips[hashed] = sip
            bisect.insort(self._hashed_ips, hashed)

    def remove(self, ip):
        for i in range(self.replicas):
            sip = _repl(ip, i)
            hashed = self._hash(sip)
            del self._ips[hashed]
            index = bisect.bisect_left(self._hashed_ips, hashed)
            del self._hashed_ips[index]

    def select(self, username):
        hashed = self._hash(username)
        start = bisect.bisect(self._hashed_ips, hashed,
                              hi=len(self._hashed_ips)-1)
        return self._ips[self._hashed_ips[start]].split(':')[0]


NUM_USERS = 1000000


def run_test(servers, users):
    selection = defaultdict(list)
    for user in users:
        user_db = servers.select(user)
        selection[user_db].append(user)

    print '===='
    print('Distribution')
    smallest = NUM_USERS + 1
    biggest = 0
    for db in selection:
        size = len(selection[db])
        if size < smallest:
            smallest = size
        if size > biggest:
            biggest = size

        print('%d users in %s' % (size, db))

    print('span: %d' % (biggest - smallest))

    # removing server 2 and 4
    servers.remove('postgres2')
    print '===='

    selection = defaultdict(list)
    for user in users:
        user_db = servers.select(user)
        selection[user_db].append(user)

    smallest = NUM_USERS + 1
    biggest = 0
    for i, db in enumerate(selection):
        size = len(selection[db])
        if size < smallest:
            smallest = size
        if size > biggest:
            biggest = size

        print('%d users in %s' % (size, db))

    print('span: %d' % (biggest - smallest))


if __name__ == '__main__':
    users = ['%06d' % i for i in range(NUM_USERS)]
    servers = ['postgres5', 'postgres2', 'postgres3', 'postgres4',
               'postgres1']

    for klass in (ConsistentHashing, RendezVous):
        for hash in (md5, sha256, binascii.crc32, fnv32a, sha512):
            try:
                cluster = klass(list(servers), hash=hash)
            except CollisionError:
                print('Collision error with hash %s' % hash)
                continue

            print(cluster)
            start = time.time()
            try:
                run_test(cluster, users)
            except CollisionError:
                print('Collision error..')
            print('Took %d seconds' % (time.time() - start))
            print
            print
	""" Consistent load-balancing.

	We have a few servers and we want a load-balancer to
	distribute incoming requests across them in a deterministic
	and consistent way - without keeping any counter to make the
	decision.

	Removing a backend server should not impact users on other
	servers.

	Adding a backend will generate the redistribution of
	users across other servers.

	The goal is to come up with the best algorithm
	for 1M users across 5 servers. Speed is a bonus.

	Two known techniques:

	- RendezVous : https://en.wikipedia.org/wiki/Rendezvous_hashing#Comparison_With_Consistent_Hashing
	- Consistent Hashing: https://en.wikipedia.org/wiki/Consistent_hashing

	Consistent Hashing implementation inspired by:
	http://techspot.zzzeek.org/2012/07/07/the-absolutely-simplest-consistent-hashing-example

	Also, good read on hashes:

	http://programmers.stackexchange.com/questions/49550/which-hashing-algorithm-is-best-for-uniqueness-and-speed/145633#145633
	"""
	import hashlib
	import bisect
	from collections import defaultdict
	import binascii
	import time
	from functools import wraps


	class CollisionError(Exception):
	pass


	_collisions = {}


	def catch_collision(func):
	@wraps(func)
	def _catch(key):
	res = func(key)
	if res in _collisions and key != _collisions[res]:
	raise CollisionError('%s and %s with %s' % (key, _collisions[res],
	func))
	_collisions[res] = key
	return res
	return _catch


	@catch_collision
	def fnv32a(key):
	hval = 0x811c9dc5
	fnv_32_prime = 0x01000193
	uint32_max = 2 ** 32
	for s in key:
	hval = hval ^ ord(s)
	hval = (hval * fnv_32_prime) % uint32_max
	return hval


	@catch_collision
	def sha512(key):
	return long(hashlib.sha512(key).hexdigest(), 16)


	@catch_collision
	def sha256(key):
	return long(hashlib.sha256(key).hexdigest(), 16)


	@catch_collision
	def md5(key):
	return long(hashlib.md5(key).hexdigest(), 16)



	class RendezVous(object):

	def __init__(self, ips=None, hash=md5):
	if ips is None:
	ips = []
	self.ips = ips
	self._hash = hash

	def __str__(self):
	return '<RendezVous with %s hash>' % self._hash

	def add(self, ip):
	self.ips.append(ip)

	def remove(self, ip):
	self.ips.remove(ip)

	def select(self, key):
	high_score = -1
	winner = None
	for ip in self.ips:
	score = self._hash("%s-%s" % (str(ip), str(key)))
	if score > high_score:
	high_score, winner = score, ip

	elif score == high_score:
	high_score, winner = score, max(str(ip), str(winner))
	return winner


	def _repl(name, index):
	return '%s:%d' % (name, index)


	class ConsistentHashing(object):

	def __init__(self, ips=[], replicas=200, hash=md5):
	self._ips = {}
	self._hashed_ips = []
	self.replicas = replicas
	self._hash = hash

	for ip in ips:
	self.add(ip)

	def __str__(self):
	return '<ConsistentHashing with %s hash>' % self._hash

	def add(self, ip):
	for i in range(self.replicas):
	sip = _repl(ip, i)
	hashed = self._hash(sip)
	self._ips[hashed] = sip
	bisect.insort(self._hashed_ips, hashed)

	def remove(self, ip):
	for i in range(self.replicas):
	sip = _repl(ip, i)
	hashed = self._hash(sip)
	del self._ips[hashed]
	index = bisect.bisect_left(self._hashed_ips, hashed)
	del self._hashed_ips[index]

	def select(self, username):
	hashed = self._hash(username)
	start = bisect.bisect(self._hashed_ips, hashed,
	hi=len(self._hashed_ips)-1)
	return self._ips[self._hashed_ips[start]].split(':')[0]


	NUM_USERS = 1000000


	def run_test(servers, users):
	selection = defaultdict(list)
	for user in users:
	user_db = servers.select(user)
	selection[user_db].append(user)

	print '===='
	print('Distribution')
	smallest = NUM_USERS + 1
	biggest = 0
	for db in selection:
	size = len(selection[db])
	if size < smallest:
	smallest = size
	if size > biggest:
	biggest = size

	print('%d users in %s' % (size, db))

	print('span: %d' % (biggest - smallest))

	# removing server 2 and 4
	servers.remove('postgres2')
	print '===='

	selection = defaultdict(list)
	for user in users:
	user_db = servers.select(user)
	selection[user_db].append(user)

	smallest = NUM_USERS + 1
	biggest = 0
	for i, db in enumerate(selection):
	size = len(selection[db])
	if size < smallest:
	smallest = size
	if size > biggest:
	biggest = size

	print('%d users in %s' % (size, db))

	print('span: %d' % (biggest - smallest))


	if __name__ == '__main__':
	users = ['%06d' % i for i in range(NUM_USERS)]
	servers = ['postgres5', 'postgres2', 'postgres3', 'postgres4',
	'postgres1']

	for klass in (ConsistentHashing, RendezVous):
	for hash in (md5, sha256, binascii.crc32, fnv32a, sha512):
	try:
	cluster = klass(list(servers), hash=hash)
	except CollisionError:
	print('Collision error with hash %s' % hash)
	continue

	print(cluster)
	start = time.time()
	try:
	run_test(cluster, users)
	except CollisionError:
	print('Collision error..')
	print('Took %d seconds' % (time.time() - start))
	print
	print