From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Fri, 16 Dec 2022 03:21:38 +0000 (+0000)
Subject: fix(pool): make sure that check() fills an empty pool
X-Git-Tag: pool-3.1.5~6^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=f196fb4184ca02f4f1b9ce1cefdc4f0f20e363e5;p=thirdparty%2Fpsycopg.git

fix(pool): make sure that check() fills an empty pool

Close #438.
---

diff --git a/docs/news_pool.rst b/docs/news_pool.rst
index 866b04548..1477a6199 100644
--- a/docs/news_pool.rst
+++ b/docs/news_pool.rst
@@ -13,7 +13,9 @@ Future releases
 psycopg_pool 3.1.5 (unreleased)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-- Avoid error in pyright caused by aliasing TypeAlias (:ticket:`#439`).
+- Make sure that `!ConnectionPool.check()` refills an empty pool
+  (:ticket:`#438`).
+- Avoid error in Pyright caused by aliasing `!TypeAlias` (:ticket:`#439`).
 
 
 Current release
diff --git a/psycopg_pool/psycopg_pool/pool.py b/psycopg_pool/psycopg_pool/pool.py
index 766c0c64a..609d95dfe 100644
--- a/psycopg_pool/psycopg_pool/pool.py
+++ b/psycopg_pool/psycopg_pool/pool.py
@@ -402,6 +402,11 @@ class ConnectionPool(BasePool[Connection[Any]]):
             conns = list(self._pool)
             self._pool.clear()
 
+            # Give a chance to the pool to grow if it has no connection.
+            # In case there are enough connection, or the pool is already
+            # growing, this is a no-op.
+            self._maybe_grow_pool()
+
         while conns:
             conn = conns.pop()
             try:
diff --git a/psycopg_pool/psycopg_pool/pool_async.py b/psycopg_pool/psycopg_pool/pool_async.py
index d411d8bc0..0ea6e9a40 100644
--- a/psycopg_pool/psycopg_pool/pool_async.py
+++ b/psycopg_pool/psycopg_pool/pool_async.py
@@ -337,6 +337,11 @@ class AsyncConnectionPool(BasePool[AsyncConnection[Any]]):
             conns = list(self._pool)
             self._pool.clear()
 
+            # Give a chance to the pool to grow if it has no connection.
+            # In case there are enough connection, or the pool is already
+            # growing, this is a no-op.
+            self._maybe_grow_pool()
+
         while conns:
             conn = conns.pop()
             try:
diff --git a/tests/pool/test_pool.py b/tests/pool/test_pool.py
index b83ca08e9..30c790b10 100644
--- a/tests/pool/test_pool.py
+++ b/tests/pool/test_pool.py
@@ -953,6 +953,37 @@ def test_reconnect_after_grow_failed(proxy):
         assert len(p._pool) == p.min_size == 4
 
 
+@pytest.mark.slow
+def test_refill_on_check(proxy):
+    proxy.start()
+    ev = Event()
+
+    def failed(pool):
+        ev.set()
+
+    with pool.ConnectionPool(
+        proxy.client_dsn, min_size=4, reconnect_timeout=1.0, reconnect_failed=failed
+    ) as p:
+        # The pool is full
+        p.wait(timeout=2)
+
+        # Break all the connection
+        proxy.stop()
+
+        # Checking the pool will empty it
+        p.check()
+        assert ev.wait(timeout=2)
+        assert len(p._pool) == 0
+
+        # Allow to connect again
+        proxy.start()
+
+        # Make sure that check has refilled the pool
+        p.check()
+        p.wait(timeout=2)
+        assert len(p._pool) == 4
+
+
 @pytest.mark.slow
 def test_uniform_use(dsn):
     with pool.ConnectionPool(dsn, min_size=4) as p:
diff --git a/tests/pool/test_pool_async.py b/tests/pool/test_pool_async.py
index 7d6ca6d79..286a77524 100644
--- a/tests/pool/test_pool_async.py
+++ b/tests/pool/test_pool_async.py
@@ -905,6 +905,37 @@ async def test_reconnect_after_grow_failed(proxy):
         assert len(p._pool) == p.min_size == 4
 
 
+@pytest.mark.slow
+async def test_refill_on_check(proxy):
+    proxy.start()
+    ev = asyncio.Event()
+
+    def failed(pool):
+        ev.set()
+
+    async with pool.AsyncConnectionPool(
+        proxy.client_dsn, min_size=4, reconnect_timeout=1.0, reconnect_failed=failed
+    ) as p:
+        # The pool is full
+        await p.wait(timeout=2)
+
+        # Break all the connection
+        proxy.stop()
+
+        # Checking the pool will empty it
+        await p.check()
+        await asyncio.wait_for(ev.wait(), 2.0)
+        assert len(p._pool) == 0
+
+        # Allow to connect again
+        proxy.start()
+
+        # Make sure that check has refilled the pool
+        await p.check()
+        await p.wait(timeout=2)
+        assert len(p._pool) == 4
+
+
 @pytest.mark.slow
 async def test_uniform_use(dsn):
     async with pool.AsyncConnectionPool(dsn, min_size=4) as p: