Merge ffae551 into c13e336

aiomcache/pool.py

@@ -14,9 +14,11 @@ def __init__(self, host, port, *, minsize, maxsize, loop=None):
         self._host = host
         self._port = port
         self._minsize = minsize
+        self._maxsize = maxsize
         self._loop = loop
         self._pool = asyncio.Queue(maxsize, loop=loop)
         self._in_use = set()
+        self._size = 0
 
     @asyncio.coroutine
     def clear(self):
@@ -26,6 +28,7 @@ def clear(self):
             self._do_close(conn)
 
     def _do_close(self, conn):
+        self._size -= 1
         conn.reader.feed_eof()
         conn.writer.close()
 
@@ -36,8 +39,12 @@ def acquire(self):
 
         :return: ``tuple`` (reader, writer)
         """
-        while self.size() < self._minsize:
+        while self._size < self._minsize:
             _conn = yield from self._create_new_conn()
+
+            # Could not create new connection
+            if _conn is None:
+                break
             yield from self._pool.put(_conn)
 
         conn = None
@@ -53,6 +60,9 @@ def acquire(self):
             if conn is None:
                 conn = yield from self._create_new_conn()
 
+            # Give up control
+            yield from asyncio.sleep(0)
+
         self._in_use.add(conn)
         return conn
 
@@ -73,9 +83,17 @@ def release(self, conn):
 
     @asyncio.coroutine
     def _create_new_conn(self):
-        reader, writer = yield from asyncio.open_connection(
-            self._host, self._port, loop=self._loop)
-        return _connection(reader, writer)
+        if self._size < self._maxsize:
+            self._size += 1
+            try:
+                reader, writer = yield from asyncio.open_connection(
+                    self._host, self._port, loop=self._loop)
+            except:
+                self._size -= 1
+                raise
+            return _connection(reader, writer)
+        else:
+            return None
 
     def size(self):
-        return len(self._in_use) + self._pool.qsize()
+        return self._size

tests/pool_test.py

@@ -43,22 +43,6 @@ def test_pool_clear(mcache_params, loop):
     assert pool._pool.qsize() == 0
 
 
-@pytest.mark.run_loop
-def test_pool_is_full(mcache_params, loop):
-    pool = MemcachePool(minsize=1, maxsize=2, loop=loop, **mcache_params)
-    conn = yield from pool.acquire()
-
-    # put garbage to the pool make it look like full
-    mocked_conns = [_connection(0, 0), _connection(1, 1)]
-    yield from pool._pool.put(mocked_conns[0])
-    yield from pool._pool.put(mocked_conns[1])
-
-    # try to return connection back
-    assert pool.size() == 3
-    pool.release(conn)
-    assert pool.size() == 2
-
-
 @pytest.mark.run_loop
 def test_acquire_dont_create_new_connection_if_have_conn_in_pool(mcache_params,
                                                                  loop):
@@ -73,3 +57,29 @@ def test_acquire_dont_create_new_connection_if_have_conn_in_pool(mcache_params,
     conn = yield from pool.acquire()
     assert conn is _conn
     assert pool.size() == 1
+
+
+@pytest.mark.run_loop
+def test_acquire_limit_maxsize(mcache_params,
+                               loop):
+    pool = MemcachePool(minsize=1, maxsize=1, loop=loop, **mcache_params)
+    assert pool.size() == 0
+
+    # Create up to max connections
+    _conn = yield from pool.acquire()
+    assert pool.size() == 1
+    pool.release(_conn)
+
+    @asyncio.coroutine
+    def acquire_wait_release():
+        conn = yield from pool.acquire()
+        assert conn is _conn
+        yield from asyncio.sleep(0.01, loop=loop)
+        assert len(pool._in_use) == 1
+        assert pool.size() == 1
+        assert pool._pool.qsize() == 0
+        pool.release(conn)
+
+    yield from asyncio.gather(*([acquire_wait_release()] * 3), loop=loop)
+    assert pool.size() == 1
+    assert len(pool._in_use) == 0