basho · ixmatus · Aug 9, 2012 · Aug 9, 2012 · Aug 9, 2012 · Aug 9, 2012
diff --git 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
@@ -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
@@ -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")