diff --git a/riak/client.py b/riak/client.py index 7bb5e8dc..e66c754c 100644 --- a/riak/client.py +++ b/riak/client.py @@ -36,7 +36,7 @@ class RiakClient(object): Riak. The Riak API uses HTTP, so there is no persistent connection, and the ``RiakClient`` object is extremely lightweight. """ - def __init__(self, host='127.0.0.1', port=8098, prefix='riak', + def __init__(self, host=('127.0.0.1', 8098), pool_size=10, prefix='riak', mapred_prefix='mapred', transport_class=None, client_id=None, solr_transport_class=None, transport_options=None): @@ -60,11 +60,10 @@ def __init__(self, host='127.0.0.1', port=8098, prefix='riak', """ if transport_class is None: transport_class = RiakHttpTransport - + api = getattr(transport_class, 'api', 1) if api >= 2: - hostports = [ (host, port), ] - self._cm = transport_class.default_cm(hostports) + self._cm = transport_class.default_cm(host, pool_size) # If no transport options are provided, then default to the # empty dict, otherwise just pass through what we are provided. @@ -79,8 +78,8 @@ def __init__(self, host='127.0.0.1', port=8098, prefix='riak', else: deprecated('please upgrade the transport to the new API') self._cm = None - self._transport = transport_class(host, port, client_id=client_id) - + self._transport = transport_class(host[0], host[1], client_id=client_id) + self._r = "default" self._w = "default" self._dw = "default" @@ -93,8 +92,7 @@ def __init__(self, host='127.0.0.1', port=8098, prefix='riak', 'text/json': json.loads} self._solr = None self._host = host - self._port = port - + def get_transport(self): """ Get the transport instance the client is using for it's connection. @@ -380,6 +378,6 @@ def get_index(self, bucket, index, startkey, endkey=None): def solr(self): if self._solr is None: - self._solr = RiakSearch(self, host=self._host, port=self._port) + self._solr = RiakSearch(self, host=self._host[0], port=self._host[1]) return self._solr diff --git a/riak/transports/connection.py b/riak/transports/connection.py index 12172df4..6ce947d2 100644 --- a/riak/transports/connection.py +++ b/riak/transports/connection.py @@ -17,9 +17,12 @@ """ import httplib -import socket import contextlib import functools +import socket +from gevent import monkey + +from gevent.queue import Queue class ConnectionManager(object): @@ -28,76 +31,50 @@ class ConnectionManager(object): # Must have a close() method connection_class = None - def __init__(self, hostports=[]): + def __init__(self, hostports, pool_size=10): # We want a private copy of this list: either to detach the argument # default, or to detach from the caller's list. - self.hostports = hostports[:] - - # Open a connection to each specified host/port. On single-threaded - # systems, this will create a round-robin across all specified servers. - # When multi-threaded, this will give us an initial set for all the - # threads to work with (and more will be created, according to demand). - self.conns = [self.connection_class(host, port) - for host, port in hostports] - - def add_hostport(self, host, port): - self.hostports.append((host, port)) - - # Open an initial connection. For single-threaded, this adds to the - # round-robin pool. On multi-threaded, it simply gives us an extra - # connectiong for the load-balancing across the servers. - self.conns.append(self.connection_class(host, port)) - - def remove_host(self, host, port=None): - if port is None: - self.hostports = [(h, p) for h, p in self.hostports - if h != host] + if type(hostports) is list: + if len(hostports) > pool_size: + raise Exception("pool_size cannot be larger than hostports") + + self.pool_size = pool_size + self.hostports = hostports + + queue_limit = pool_size * len(hostports) if type(hostports) is list else pool_size + self.queue = Queue(queue_limit) + + # Patch httplib if we are using that, also patch the sockets + monkey.patch_all(select=False, os=False, ssl=False, subprocess=False) + + if type(self.hostports) is list: + for host, port in self.hostports: + for i in range(0, pool_size): + self.queue.put(self.connection_class(host, port)) else: - self.hostports.remove((host, port)) - - # Now that the host/port pair has been removed from self.hostports, - # no connections on this pair will be added in .giveback(). Thus, the - # existing connections are all that may exist at this time. We'll - # snapshot the list, and look for offending connections, then try and - # remove them, being wary that race conditions may remove them before - # we can remove it. - for conn in self.conns[:]: - if conn.host == host and (port is None or conn.port == port): - try: - self.conns.remove(conn) - except ValueError: - # Another thread removed the connection. It won't be coming back, - # so we have nothing to do here. - pass - else: - # If the connection was still present (no ValueError), then we - # should go ahead and close it down. - conn.close() - - # Just in case somebody uses a host/port combo and typos... - remove_hostport = remove_host - - def take(self): - if len(self.conns) == 0: - # RACE: in a multi-threaded environment, a conn might arrive in - # self.conns, but... no biggy. If we're bouncing up against - # needing a new connection, then we'll just create one. - return self._new_connection() - - # RACE: self.conns might empty out right now, so we need to protect - # our access to it. - try: - # round-robin: take from the front, we'll append when it comes back - return self.conns.pop(0) - except IndexError: - return self._new_connection() - - def giveback(self, conn): - # Connections using a host/port pair that is NOT in self.hostports - # should be ignored. Likely, remove_host() was called while this - # connection was borrowed for some work. - if (conn.host, conn.port) in self.hostports: - self.conns.append(conn) + for i in range(0, pool_size): + self.queue.put(self.connection_class(hostports[0], hostports[1])) + + def checkout(self): + """Checkout a connection from the queue - block if we have + none in the queue till one arrives.""" + + return self.queue.get() + + def checkin(self, conn): + """Checkin a connection back into the queue. + + If it doesn't exist in the hostports list then + close it. + """ + + # If we are tring to checkin a connection and the pool has + # already been re-filled, throw away the connection silently + if self.queue.qsize() >= self.pool_size: + conn.close() + + if ((type(self.hostports) is list) and (conn.host, conn.port) in self.hostports) or (conn.host, conn.port) == self.hostports: + self.queue.put(conn) else: # Proactively close the connection. The caller won't know whether # we put it into our list, or left the connection for the caller @@ -106,45 +83,22 @@ def giveback(self, conn): @contextlib.contextmanager def withconn(self): - conn = self.take() + """Context managaer method.""" + conn = self.checkout() try: yield conn finally: - self.giveback(conn) - - def _new_connection(self): - if len(self.hostports) == 0: - raise NoHostsDefined() - - # Grab the first host/port combo. We'll put this at the end, so that - # we do a round-robin on the host/port pairs. - host, port = self.hostports[0] - conn = self.connection_class(host, port) - - if len(self.hostports) == 1: - # No rotation needed. - return conn - - # Be careful about rotating. We want to append before removing, so that - # we never hit a len==0 race condition (which could prevent the creation - # of needed connections). - self.hostports.append((host, port)) - - # RACE: another thread may have appended the same host/port pair. We - # will add another pair. Each thread will remove one (either [0], or - # one that had been appened), resulting in a correct state of a single - # pair in the list. - # RACE: another thread may get the host/port pair from hostports[0] - # before we have a chance to remove it. We don't need precision - # round-robin behavior; just something close. - # RACE: another thread may have removed hostports[0] (which we are - # also trying to remove), but it will have placed another copy at - # the end before doing so. We have also added a host/port pair, and - # will remove one, leaving the list in a correct state. - self.hostports.remove((host, port)) - - return conn - + self.checkin(conn) + + def new(self): + """Create a new connection if we have room for it.""" + + if self.queue.qsize() >= self.pool_size: + raise Exception("the queue is already full") + + host, port = self.hostports[0] if type(self.hostports) is list else self.hostports + + self.queue.put(self.connection_class(host, port)) class Socket(object): @@ -172,9 +126,9 @@ def close(self): class FactoryConnectionManager(ConnectionManager): - def __init__(self, connection_class, hostports=[]): + def __init__(self, connection_class, hostports, pool_size): self.connection_class = connection_class - ConnectionManager.__init__(self, hostports) + ConnectionManager.__init__(self, hostports, pool_size) def cm_using(connection_class): diff --git a/setup.py b/setup.py index d6a1d020..3524b4af 100755 --- a/setup.py +++ b/setup.py @@ -13,7 +13,8 @@ def make_docs(): os.rename(name, 'docs/%s' % name) install_requires = ["riak_pb >=1.2.0, < 1.3.0"] -requires = ["riak_pb(>=1.2.0,<1.3.0)"] +requires = ["riak_pb(>=1.2.0,<1.3.0)", + "gevent"] tests_require = [] if platform.python_version() < '2.7': tests_require.append("unittest2")