class RDBMultiClient(DictNature):
pool_size = 5 # use this many threads per RDBClient
def __init__(self, weights):
self.weights = weights
self.nodes = set(x[0] for x in weights)
self.hasher = ConsistantHasher(weights)
self.clients = dict((node, RDBClient(node))
for node in self.nodes)
self.parallel_transfer = True
if self.parallel_transfer:
# a thread-pool to support concurrent bulk requests that
# span multiple nodes
self.thread_pool = ThreadPool(len(self.nodes) * self.pool_size)
def get(self, key, *a, **kw):
self.clients[self.hasher[key]].get(key, *a, **kw)
def put(self, key, value, *a, **kw):
self.clients[self.hasher[key]].put(key, value, *a, **kw)
def delete(self, key, *a, **kw):
self.clients[self.hasher[key]].delete(key, *a, **kw)
def get_multi(self, keys):
return self.bulk(get = keys)
def put_multi(self, keys):
return self.bulk(put = keys)
def delete_multi(self, keys):
return self.bulk(delete = keys)
def bulk(self, get = [], put = {}, delete = []):
"""Do multiple _bulk requests in parallel"""
if not isinstance(put, dict):
put = dict(put)
by_node = {}
for key in get:
by_node.setdefault(self.hasher[key],
{}).setdefault('get', []).append(key)
for key in delete:
by_node.setdefault(self.hasher[key],
{}).setdefault('delete', []).append(key)
for key, val in put.iteritems():
by_node.setdefault(self.hasher[key],
{}).setdefault('put', {})[key] = val
funcs = []
for node, ops in by_node.iteritems():
def fetch(_node, _ops):
def _fetch():
return self.clients[_node].bulk(**_ops)
return _fetch
funcs.append(fetch(node, ops))
if self.parallel_transfer and len(funcs) > 1:
bulks = self.thread_pool.pmap(funcs)
else:
bulks = [f() for f in funcs]
ret = {}
for bulk in bulks:
ret.update(bulk)
return ret
def _by_node(self, keys):
ret = {}
for key in keys:
ret.setdefault(self.hasher[key], []).append(key)
return ret.items()
def __repr__(self):
return "%s(%r)" % (self.__class__.__name__, self.weights)
