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)