connection pool is now a real pool. we no longer rely on threading.local for anything in redis-py. yaa!

1 files changed, 30 insertions, 24 deletions

diff --git a/redis/connection.py b/redis/connection.py
index dc4d40c..b5e7a49 100644
--- a/redis/connection.py
+++ b/redis/connection.py
@@ -1,8 +1,7 @@
 import errno
 import os
 import socket
-import threading
-from itertools import imap
+from itertools import chain, imap
 from redis.exceptions import ConnectionError, ResponseError, InvalidResponse
 
 class PythonParser(object):
@@ -235,36 +234,43 @@ class UnixDomainSocketConnection(Connection):
         sock.connect(self.path)
         return sock
 
+# TODO: add ability to block waiting on a connection to be released
 class ConnectionPool(object):
-    """
-    A connection pool that maintains only one connection. Great for
-    single-threaded apps with no sharding
-    """
-    def __init__(self, connection_class=Connection, **kwargs):
+    "Generic connection pool"
+    def __init__(self, connection_class=Connection, max_connections=None,
+                 **connection_kwargs):
         self.connection_class = connection_class
-        self.kwargs = kwargs
-        self._connection = None
-        self._in_use = False
-
-    def copy(self):
-        "Return a new instance of this class with the same parameters"
-        return self.__class__(self.connection_class, **self.kwargs)
+        self.connection_kwargs = connection_kwargs
+        self.max_connections = max_connections or 2**31
+        self._created_connections = 0
+        self._available_connections = []
+        self._in_use_connections = set()
 
     def get_connection(self, command_name, *keys):
         "Get a connection from the pool"
-        if self._in_use:
-            raise ConnectionError("Connection already in-use")
-        if not self._connection:
-            self._connection = self.connection_class(**self.kwargs)
-        self._in_use = True
-        return self._connection
+        try:
+            connection = self._available_connections.pop()
+        except IndexError:
+            connection = self.make_connection()
+        self._in_use_connections.add(connection)
+        return connection
+
+    def make_connection(self):
+        "Create a new connection"
+        if self._created_connections >= self.max_connections:
+            raise Exception("Too many connections")
+        self._created_connections += 1
+        return self.connection_class(**self.connection_kwargs)
 
     def release(self, connection):
         "Releases the connection back to the pool"
-        assert self._connection == connection
-        self._in_use = False
+        # assert self._connection == connection
+        # self._in_use = False
+        self._in_use_connections.remove(connection)
+        self._available_connections.append(connection)
 
     def disconnect(self):
         "Disconnects all connections in the pool"
-        if self._connection:
-            self._connection.disconnect()
+        all_conns = chain(self._available_connections, self._in_use_connections)
+        for connection in all_conns:
+            connection.disconnect()